From 83fa2dee728ec2e7aeb69c8b394b454fea22a839 Mon Sep 17 00:00:00 2001
From: Ivan Krutov <vania-pooh@vania-pooh.com>
Date: Wed, 18 Jul 2018 13:42:19 +0300
Subject: [PATCH] Ability to proxy logs to upstream host

---
 docs/log-files.adoc |  3 +++
 proxy.go            |  8 ++++++++
 proxy_test.go       | 28 ++++++++++++++++++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/docs/log-files.adoc b/docs/log-files.adoc
index 6ee33de..9c48b04 100644
--- a/docs/log-files.adoc
+++ b/docs/log-files.adoc
@@ -40,11 +40,13 @@ The following statuses are available:
 | INIT | Server initialization messages 
 | INVALID_HOST_VNC_URL | Failed to parse VNC host URL specified in quota configuration 
 | INVALID_DOWNLOAD_REQUEST_URL | Download request URL do not contain enough information to determine upstream host 
+| INVALID_LOG_REQUEST_URL | Log request URL do not contain enough information to determine upstream host 
 | INVALID_VNC_REQUEST_URL | VNC request URL do not contain enough information to determine upstream host 
 | INVALID_VIDEO_REQUEST_URL | Video request URL do not contain enough information to determine upstream host 
 | INVALID_URL | Session ID does not contain information about host where it was created 
 | PROXYING | Proxying Selenium request (shown in verbose mode only) 
 | PROXYING_DOWNLOAD | Starting to proxy downloaded file from upstream host 
+| PROXYING_LOG | Starting to proxy log from upstream host 
 | PROXYING_TO_VNC | Starting to proxy VNC traffic 
 | PROXYING_VIDEO | Starting to proxy video from upstream host 
 | QUOTA_INFO_REQUESTED | Quota information request arrived 
@@ -56,6 +58,7 @@ The following statuses are available:
 | SESSION_NOT_CREATED | Attempts to create a new session on all hosts failed. An error was returned to user. 
 | SHUTTING_DOWN | Server is shutting down and waiting graceful shutdown timeout for currently proxied requests to finish  
 | UNKNOWN_DOWNLOAD_HOST | Requested to proxy downloaded file to host not present in quota 
+| UNKNOWN_LOG_HOST | Requested to proxy log to host not present in quota 
 | UNKNOWN_VNC_HOST | Requested to proxy VNC to host not present in quota 
 | UNKNOWN_VIDEO_HOST | Requested to proxy video to host not present in quota 
 | UNSUPPORTED_BROWSER | Requested browser name and version is not present in quota 
diff --git a/proxy.go b/proxy.go
index 7da1d63..31360a6 100644
--- a/proxy.go
+++ b/proxy.go
@@ -38,6 +38,7 @@ const (
 	proxyPath      = routePath + "/"
 	vncPath        = "/vnc/"
 	videoPath      = "/video/"
+	logsPath       = "/logs/"
 	downloadPath   = "/download/"
 	head           = len(proxyPath)
 	md5SumLength   = 32
@@ -679,6 +680,12 @@ func video(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
+func logs(w http.ResponseWriter, r *http.Request) {
+	proxyStatic(w, r, logsPath, "INVALID_LOG_REQUEST_URL", "PROXYING_LOG", "UNKNOWN_LOG_HOST", func(sessionId string) string {
+		return fmt.Sprintf("/logs/%s.log", sessionId)
+	})
+}
+
 func download(w http.ResponseWriter, r *http.Request) {
 	proxyStatic(w, r, downloadPath, "INVALID_DOWNLOAD_REQUEST_URL", "PROXYING_DOWNLOAD", "UNKNOWN_DOWNLOAD_HOST", func(remainder string) string {
 		return fmt.Sprintf("/download/%s", remainder)
@@ -730,6 +737,7 @@ func mux() http.Handler {
 	mux.Handle(proxyPath, &httputil.ReverseProxy{Director: proxy})
 	mux.Handle(vncPath, websocket.Handler(vnc))
 	mux.HandleFunc(videoPath, WithSuitableAuthentication(authenticator, video))
+	mux.HandleFunc(logsPath, WithSuitableAuthentication(authenticator, logs))
 	mux.HandleFunc(downloadPath, WithSuitableAuthentication(authenticator, download))
 	return mux
 }
diff --git a/proxy_test.go b/proxy_test.go
index 51e6db1..e644daf 100644
--- a/proxy_test.go
+++ b/proxy_test.go
@@ -364,6 +364,34 @@ func prepareMockFileServer(path string) (*httptest.Server, string) {
 	return fileServer, sessionID
 }
 
+func TestProxyLogsWithoutAuth(t *testing.T) {
+	rsp, err := http.Get(gridrouter("/logs/123"))
+
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusUnauthorized})
+}
+
+func TestProxyLogs(t *testing.T) {
+
+	test.Lock()
+	defer test.Unlock()
+
+	fileServer, sessionID := prepareMockFileServer("/logs/123.log")
+	defer fileServer.Close()
+
+	rsp, err := doBasicHTTPRequest(http.MethodGet, gridrouter(fmt.Sprintf("/logs/%s", sessionID)), nil)
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusOK})
+
+	rsp, err = doBasicHTTPRequest(http.MethodGet, gridrouter("/logs/missing-session-id"), nil)
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusNotFound})
+
+	rsp, err = doBasicHTTPRequest(http.MethodGet, gridrouter("/logs/f7fd94f75c79c36e547c091632da440f_missing-file"), nil)
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, rsp, Code{http.StatusNotFound})
+}
+
 func TestProxyDownloadWithoutAuth(t *testing.T) {
 	rsp, err := http.Get(gridrouter("/download/123"))