From 5548e914e54538f4bc3f5d3b77591a0083b689cc Mon Sep 17 00:00:00 2001
From: Ivan Krutov <vania-pooh@vania-pooh.com>
Date: Tue, 19 Dec 2017 18:20:18 +0300
Subject: [PATCH] Added ability to proxy video (fixes #142)

---
 proxy.go      | 33 +++++++++++++++++++++++++++++++--
 proxy_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/proxy.go b/proxy.go
index e05cbd8..a0a4fcd 100644
--- a/proxy.go
+++ b/proxy.go
@@ -35,6 +35,7 @@ const (
 	routePath      string = "/wd/hub/session"
 	proxyPath      string = routePath + "/"
 	vncPath        string = "/vnc/"
+	videoPath      string = "/video/"
 	head           int    = len(proxyPath)
 	md5SumLength   int    = 32
 	tail           int    = head + md5SumLength
@@ -477,12 +478,11 @@ 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)
+		log.Printf("[INVALID_VNC_REQUEST_URL] [%s]\n", path)
 		return
 	}
 	sum := path[head:tail]
@@ -499,6 +499,7 @@ func vnc(wsconn *websocket.Conn) {
 			port = vncInfo.Port
 			path = vncInfo.Path
 		}
+		sessionId := strings.Split(wsconn.Request().URL.Path, "/")[2]
 		switch scheme {
 		case vncScheme:
 			proxyVNC(wsconn, sessionId, host, port)
@@ -548,6 +549,33 @@ func proxyConn(wsconn *websocket.Conn, conn net.Conn, err error, sessionId strin
 	log.Printf("[VNC_CLIENT_DISCONNECTED] [%s] [%s]\n", sessionId, address)
 }
 
+func video(w http.ResponseWriter, r *http.Request) {
+	confLock.RLock()
+	defer confLock.RUnlock()
+
+	head := len(videoPath)
+	tail := head + md5SumLength
+	path := r.URL.Path
+	if len(path) < tail {
+		log.Printf("[INVALID_VIDEO_REQUEST_URL] [%s]\n", path)
+		reply(w, errMsg("invalid video request URL"), http.StatusNotFound)
+		return
+	}
+	sum := path[head:tail]
+	sessionId := path[tail:]
+	h, ok := routes[sum]
+	if ok {
+		(&httputil.ReverseProxy{Director: func(r *http.Request) {
+			r.URL.Scheme = "http"
+			r.URL.Host = h.net()
+			r.URL.Path = fmt.Sprintf("/video/%s.mp4", sessionId)
+		}}).ServeHTTP(w, r)
+	} else {
+		log.Printf("[UNKNOWN_VIDEO_HOST] [%s]\n", sum)
+		reply(w, errMsg("unknown video host"), http.StatusNotFound)
+	}
+}
+
 func mux() http.Handler {
 	mux := http.NewServeMux()
 	authenticator := auth.NewBasicAuthenticator(
@@ -560,5 +588,6 @@ func mux() http.Handler {
 	mux.HandleFunc(routePath, withCloseNotifier(WithSuitableAuthentication(authenticator, postOnly(route))))
 	mux.Handle(proxyPath, &httputil.ReverseProxy{Director: proxy})
 	mux.Handle(vncPath, websocket.Handler(vnc))
+	mux.HandleFunc(videoPath, WithSuitableAuthentication(authenticator, video))
 	return mux
 }
diff --git a/proxy_test.go b/proxy_test.go
index cd994a4..cc5a015 100644
--- a/proxy_test.go
+++ b/proxy_test.go
@@ -239,6 +239,54 @@ func TestProxyScreenWebSocketsProtocol(t *testing.T) {
 
 }
 
+func TestProxyVideoFileWithoutAuth(t *testing.T) {
+	rsp, err := http.Get(gridrouter("/video/123.mp4"))
+
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusUnauthorized})
+}
+
+func TestProxyVideoFile(t *testing.T) {
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/video/123.mp4", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	fileServer := httptest.NewServer(mux)
+	defer fileServer.Close()
+
+	host, port := hostportnum(fileServer.URL)
+	node := Host{Name: host, Port: port, Count: 1}
+
+	test.Lock()
+	defer test.Unlock()
+
+	browsers := Browsers{Browsers: []Browser{
+		{Name: "browser", DefaultVersion: "1.0", Versions: []Version{
+			{Number: "1.0", Regions: []Region{
+				{Hosts: Hosts{
+					node,
+				}},
+			}},
+		}}}}
+	updateQuota(user, browsers)
+
+	sessionId := node.sum() + "123"
+
+	rsp, err := doBasicHTTPRequest(http.MethodGet, gridrouter(fmt.Sprintf("/video/%s", sessionId)), nil)
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusOK})
+
+	rsp, err = doBasicHTTPRequest(http.MethodGet, gridrouter("/video/missing-file"), nil)
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusNotFound})
+
+	rsp, err = doBasicHTTPRequest(http.MethodGet, gridrouter("/video/f7fd94f75c79c36e547c091632da440f_missing-file"), nil)
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusNotFound})
+}
+
 func TestCreateSessionGet(t *testing.T) {
 	req, _ := http.NewRequest(http.MethodGet, gridrouter("/wd/hub/session"), nil)
 	req.SetBasicAuth("test", "test")