Merge pull request #120 from vania-pooh/master

Ability to proxy VNC (fixes #117)

config.go

@@ -43,8 +43,18 @@ type Host struct {
 	Count    int    `xml:"count,attr"`
 	Username string `xml:"username,attr,omitempty"`
 	Password string `xml:"password,attr,omitempty"`
+	VNC      string `xml:"vnc,attr,omitempty"`
 	region   string
+	vncInfo  *vncInfo
 }
+
+type vncInfo struct {
+	Scheme string
+	Host   string
+	Port   string
+	Path   string
+}
+
 type set interface {
 	contains(el string) bool
 	add(el string)

proxy.go

@@ -18,6 +18,8 @@ import (
 	"context"
 
 	"github.com/abbot/go-http-auth"
+	"golang.org/x/net/websocket"
+	"io"
 )
 
 const (
@@ -27,15 +29,19 @@ const (
 )
 
 const (
-	pingPath     string = "/ping"
-	errPath      string = "/err"
-	hostPath     string = "/host/"
-	routePath    string = "/wd/hub/session"
-	proxyPath    string = routePath + "/"
-	head         int    = len(proxyPath)
-	md5SumLength int    = 32
-	tail         int    = head + md5SumLength
-	sessPart     int    = 4 // /wd/hub/session/{various length session}
+	pingPath       string = "/ping"
+	errPath        string = "/err"
+	hostPath       string = "/host/"
+	routePath      string = "/wd/hub/session"
+	proxyPath      string = routePath + "/"
+	vncPath        string = "/vnc/"
+	head           int    = len(proxyPath)
+	md5SumLength   int    = 32
+	tail           int    = head + md5SumLength
+	sessPart       int    = 4 // /wd/hub/session/{various length session}
+	defaultVNCPort string = "5900"
+	vncScheme      string = "vnc"
+	wsScheme       string = "ws"
 )
 
 var (
@@ -398,15 +404,40 @@ func appendRoutes(routes Routes, config *Browsers) Routes {
 		for _, v := range b.Versions {
 			for _, r := range v.Regions {
 				for i, h := range r.Hosts {
-					r.Hosts[i].region = r.Name
-					routes[h.sum()] = &r.Hosts[i]
+					host := r.Hosts[i]
+					host.region = r.Name
+					host.vncInfo = createVNCInfo(host)
+					routes[h.sum()] = &host
 				}
 			}
 		}
 	}
 	return routes
 }
 
+func createVNCInfo(h Host) *vncInfo {
+	vncUrl := h.VNC
+	if vncUrl != "" {
+		u, err := url.Parse(vncUrl)
+		if err != nil {
+			log.Printf("[INVALID_HOST_VNC_URL] [%s] [%s]\n", fmt.Sprintf("%s:%s", h.Name, h.Port), vncUrl)
+			return nil
+		}
+		if u.Scheme != vncScheme && u.Scheme != wsScheme {
+			log.Printf("[UNSUPPORTED_HOST_VNC_SCHEME] [%s] [%s]\n", fmt.Sprintf("%s:%s", h.Name, h.Port), vncUrl)
+			return nil
+		}
+		vncInfo := vncInfo{
+			Scheme: u.Scheme,
+			Path:   u.Path,
+		}
+		vncInfo.Scheme = u.Scheme
+		vncInfo.Host, vncInfo.Port, _ = net.SplitHostPort(u.Host)
+		return &vncInfo
+	}
+	return nil
+}
+
 func requireBasicAuth(authenticator *auth.BasicAuth, handler func(http.ResponseWriter, *http.Request)) func(http.ResponseWriter, *http.Request) {
 	return authenticator.Wrap(func(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
 		handler(w, &r.Request)
@@ -435,6 +466,78 @@ func WithSuitableAuthentication(authenticator *auth.BasicAuth, handler func(http
 	}
 }
 
+func vnc(wsconn *websocket.Conn) {
+	defer wsconn.Close()
+	confLock.RLock()
+	defer confLock.RUnlock()
+	sessionId := strings.Split(wsconn.Request().URL.Path, "/")[2]
+	head := len(vncPath)
+	tail := head + md5SumLength
+	path := wsconn.Request().URL.Path
+	if len(path) < tail {
+		log.Printf("[INVALID_VNC_REQUEST_URL] [%s]\n", wsconn.Request().URL.Path)
+		return
+	}
+	sum := path[head:tail]
+	h, ok := routes[sum]
+	if ok {
+		vncInfo := h.vncInfo
+		scheme := vncScheme
+		host := h.Name
+		port := defaultVNCPort
+		path := ""
+		if vncInfo != nil {
+			scheme = vncInfo.Scheme
+			host = vncInfo.Host
+			port = vncInfo.Port
+			path = vncInfo.Path
+		}
+		switch scheme {
+		case vncScheme: proxyVNC(wsconn, sessionId, host, port)
+		case wsScheme: proxyWebSockets(wsconn, sessionId, host, port, path)
+		default: {
+			log.Printf("[UNSUPPORTED_HOST_VNC_SCHEME] [%s]\n", scheme)
+			return
+		}
+		}
+	} else {
+		log.Printf("[UNKNOWN_VNC_HOST] [%s]\n", sum)
+	}
+
+}
+
+func proxyVNC(wsconn *websocket.Conn, sessionId string, host string, port string) {
+	var d net.Dialer
+	address := fmt.Sprintf("%s:%s", host, port)
+	conn, err := d.DialContext(wsconn.Request().Context(), "tcp", address)
+	proxyConn(wsconn, conn, err, sessionId, address)
+}
+
+func proxyWebSockets(wsconn *websocket.Conn, sessionId string, host string, port string, path string) {
+	origin := "http://localhost/"
+	u := fmt.Sprintf("ws://%s:%s%s/%s", host, port, path, sessionId)
+	//TODO: consider context from wsconn
+	conn, err := websocket.Dial(u, "", origin)
+	proxyConn(wsconn, conn, err, sessionId, u)
+}
+
+func proxyConn(wsconn *websocket.Conn, conn net.Conn, err error, sessionId string, address string) {
+	log.Printf("[PROXYING_TO_VNC] [%s] [%s]\n", sessionId, address)
+	if err != nil {
+		log.Printf("[VNC_ERROR] [%s] [%s] [%v]\n", sessionId, address, err)
+		return
+	}
+	defer conn.Close()
+	wsconn.PayloadType = websocket.BinaryFrame
+	go func() {
+		io.Copy(wsconn, conn)
+		wsconn.Close()
+		log.Printf("[VNC_SESSION_CLOSED] [%s] [%s]\n", sessionId, address)
+	}()
+	io.Copy(conn, wsconn)
+	log.Printf("[VNC_CLIENT_DISCONNECTED] [%s] [%s]\n", sessionId, address)
+}
+
 func mux() http.Handler {
 	mux := http.NewServeMux()
 	authenticator := auth.NewBasicAuthenticator(
@@ -446,5 +549,6 @@ func mux() http.Handler {
 	mux.HandleFunc(hostPath, WithSuitableAuthentication(authenticator, host))
 	mux.HandleFunc(routePath, withCloseNotifier(WithSuitableAuthentication(authenticator, postOnly(route))))
 	mux.Handle(proxyPath, &httputil.ReverseProxy{Director: proxy})
+	mux.Handle(vncPath, websocket.Handler(vnc))
 	return mux
 }

proxy_test.go

@@ -19,6 +19,8 @@ import (
 
 	. "github.com/aandryashin/matchers"
 	. "github.com/aandryashin/matchers/httpresp"
+	"golang.org/x/net/websocket"
+	"strings"
 )
 
 var (
@@ -151,6 +153,90 @@ func testGetHost(t *testing.T, sessionId string, statusCode int) *Host {
 	return &host
 }
 
+func TestProxyScreenVNCProtocol(t *testing.T) {
+
+	test.Lock()
+	defer test.Unlock()
+
+	const testData = "vnc-data"
+	server := testTcpServer(testData)
+	defer server.Close()
+
+	vncHost := Host{Name: "example.com", Port: 4444, Count: 1, VNC: fmt.Sprintf("vnc://%s", server.Addr().String())}
+
+	browsers := Browsers{Browsers: []Browser{
+		{Name: "browser", DefaultVersion: "1.0", Versions: []Version{
+			{Number: "1.0", Regions: []Region{
+				{Hosts: Hosts{
+					vncHost,
+				}},
+			}},
+		}}}}
+	updateQuota(user, browsers)
+
+	testDataReceived(vncHost, testData, t)
+}
+
+func testDataReceived(host Host, correctData string, t *testing.T) {
+	sessionId := host.sum() + "123"
+
+	origin := "http://localhost/"
+	u := fmt.Sprintf("ws://%s/vnc/%s", srv.Listener.Addr(), sessionId)
+	ws, err := websocket.Dial(u, "", origin)
+	AssertThat(t, err, Is{nil})
+
+	var data = make([]byte, len(correctData))
+	_, err = ws.Read(data)
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, strings.TrimSpace(string(data)), EqualTo{correctData})
+}
+
+func testTcpServer(data string) net.Listener {
+	l, _ := net.Listen("tcp", "127.0.0.1:0")
+	go func() {
+		for {
+			conn, err := l.Accept()
+			if err != nil {
+				continue
+			}
+			defer conn.Close()
+			io.WriteString(conn, data)
+			return
+		}
+	}()
+	return l
+}
+
+func TestProxyScreenWebSocketsProtocol(t *testing.T) {
+
+	test.Lock()
+	defer test.Unlock()
+
+	const testData = "ws-data"
+	mux := http.NewServeMux()
+	mux.Handle("/vnc/", websocket.Handler(func(wsconn *websocket.Conn) {
+		wsconn.Write([]byte(testData))
+	}))
+
+	wsServer := httptest.NewServer(mux)
+	defer wsServer.Close()
+
+	wsHost := Host{Name: "example.com", Port: 4444, Count: 1, VNC: fmt.Sprintf("ws://%s/vnc", wsServer.Listener.Addr().String())}
+
+	browsers := Browsers{Browsers: []Browser{
+		{Name: "browser", DefaultVersion: "1.0", Versions: []Version{
+			{Number: "1.0", Regions: []Region{
+				{Hosts: Hosts{
+					wsHost,
+				}},
+			}},
+		}}}}
+	updateQuota(user, browsers)
+
+	testDataReceived(wsHost, testData, t)
+
+}
+
 func TestCreateSessionGet(t *testing.T) {
 	req, _ := http.NewRequest(http.MethodGet, gridrouter("/wd/hub/session"), nil)
 	req.SetBasicAuth("test", "test")

vendor/vendor.json

@@ -37,6 +37,12 @@
 			"path": "golang.org/x/net/context",
 			"revision": "f2499483f923065a842d38eb4c7f1927e6fc6e6d",
 			"revisionTime": "2017-01-14T04:22:49Z"
+		},
+		{
+			"checksumSHA1": "7EZyXN0EmZLgGxZxK01IJua4c8o=",
+			"path": "golang.org/x/net/websocket",
+			"revision": "f2499483f923065a842d38eb4c7f1927e6fc6e6d",
+			"revisionTime": "2017-01-14T04:22:49Z"
 		}
 	],
 	"rootPath": "github.com/aerokube/ggr"