auto update script executed

mod/sshprox/sshprox.go

@@ -10,10 +10,12 @@ import (
 	"path/filepath"
 	"runtime"
 	"strconv"
+	"strings"
 
 	"github.com/google/uuid"
 	"imuslab.com/zoraxy/mod/reverseproxy"
 	"imuslab.com/zoraxy/mod/utils"
+	"imuslab.com/zoraxy/mod/websocketproxy"
 )
 
 /*
@@ -31,10 +33,13 @@ type Manager struct {
 }
 
 type Instance struct {
-	UUID     string
-	ExecPath string
-	conn     *reverseproxy.ReverseProxy //HTTP proxy
-	tty      *exec.Cmd                  //SSH connection ported to web interface
+	UUID         string
+	ExecPath     string
+	RemoteAddr   string
+	RemotePort   int
+	AssignedPort int
+	conn         *reverseproxy.ReverseProxy //HTTP proxy
+	tty          *exec.Cmd                  //SSH connection ported to web interface
 }
 
 func NewSSHProxyManager() *Manager {
@@ -69,6 +74,21 @@ func (m *Manager) HandleHttpByInstanceId(instanceId string, w http.ResponseWrite
 		return
 	}
 
+	r.Header.Set("X-Forwarded-Host", r.Host)
+	requestURL := r.URL.String()
+	if r.Header["Upgrade"] != nil && strings.ToLower(r.Header["Upgrade"][0]) == "websocket" {
+		//Handle WebSocket request. Forward the custom Upgrade header and rewrite origin
+		r.Header.Set("A-Upgrade", "websocket")
+		if len(requestURL) > 0 && requestURL[:1] == "/" {
+			//Remove starting / from request URL if exists
+			requestURL = requestURL[1:]
+		}
+		u, _ := url.Parse("ws://127.0.0.1:" + strconv.Itoa(targetInstance.AssignedPort) + requestURL)
+		wspHandler := websocketproxy.NewProxy(u)
+		wspHandler.ServeHTTP(w, r)
+		return
+	}
+
 	targetInstance.conn.ProxyHTTP(w, r)
 }
 
@@ -101,8 +121,9 @@ func (m *Manager) NewSSHProxy(binaryRoot string) (*Instance, error) {
 	}
 
 	thisInstance := Instance{
-		UUID:     uuid.New().String(),
-		ExecPath: realpath,
+		UUID:         uuid.New().String(),
+		ExecPath:     realpath,
+		AssignedPort: -1,
 	}
 
 	m.Instances = append(m.Instances, &thisInstance)
@@ -120,6 +141,9 @@ func (i *Instance) CreateNewConnection(listenPort int, remoteIpAddr string, remo
 		cmd.Run()
 	}()
 	i.tty = cmd
+	i.AssignedPort = listenPort
+	i.RemoteAddr = remoteIpAddr
+	i.RemotePort = remotePort
 
 	//Create a new proxy agent for this root
 	path, err := url.Parse("http://127.0.0.1:" + strconv.Itoa(listenPort))