From 2df8a291433e61fa837657c720c47fe57962f772 Mon Sep 17 00:00:00 2001
From: Ro Santalla <roberto@santalla.io>
Date: Tue, 23 Jul 2024 09:25:24 +0200
Subject: [PATCH] http: implement `/proxy/{id}` to proxy requests to chromium

---
 go.mod       |  5 +++++
 go.sum       |  4 ++++
 http/http.go | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/go.mod b/go.mod
index b1af83a..19a6562 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,8 @@
 module github.com/grafana/crocochrome
 
 go 1.22
+
+require (
+	github.com/gorilla/websocket v1.5.3 // indirect
+	github.com/koding/websocketproxy v0.0.0-20181220232114-7ed82d81a28c // indirect
+)
diff --git a/go.sum b/go.sum
index e69de29..5e41d3f 100644
--- a/go.sum
+++ b/go.sum
@@ -0,0 +1,4 @@
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/koding/websocketproxy v0.0.0-20181220232114-7ed82d81a28c h1:N7A4JCA2G+j5fuFxCsJqjFU/sZe0mj8H0sSoSwbaikw=
+github.com/koding/websocketproxy v0.0.0-20181220232114-7ed82d81a28c/go.mod h1:Nn5wlyECw3iJrzi0AhIWg+AJUb4PlRQVW4/3XHH1LZA=
diff --git a/http/http.go b/http/http.go
index 2a26c51..a554890 100644
--- a/http/http.go
+++ b/http/http.go
@@ -8,6 +8,7 @@ import (
 	"net/url"
 
 	"github.com/grafana/crocochrome"
+	"github.com/koding/websocketproxy"
 )
 
 type Server struct {
@@ -28,6 +29,7 @@ func New(logger *slog.Logger, supervisor *crocochrome.Supervisor) *Server {
 	mux.HandleFunc("GET /sessions", api.List)
 	mux.HandleFunc("POST /sessions", api.Create)
 	mux.HandleFunc("DELETE /sessions/{id}", api.Delete)
+	mux.HandleFunc("/proxy/{id}", api.Proxy)
 
 	return api
 }
@@ -82,6 +84,42 @@ func (s *Server) Delete(rw http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// Proxy checks an open session for the given session ID (from path) and proxies the request to the URL present in that
+// session.
+// This is needed as recent versions of chromium do not support listening in addresses other than localhost, so to make
+// chromium reachable from the outside we need to proxy it.
+func (s *Server) Proxy(rw http.ResponseWriter, r *http.Request) {
+	sessionID := r.PathValue("id")
+	if sessionID == "" {
+		rw.WriteHeader(http.StatusBadRequest)
+		return
+	}
+
+	sessionInfo := s.supervisor.Session(sessionID)
+	if sessionInfo == nil {
+		s.logger.Warn("sessionID not found", "sessionID", sessionID)
+		rw.WriteHeader(http.StatusNotFound)
+		return
+	}
+
+	rawUrl := sessionInfo.ChromiumVersion.WebSocketDebuggerURL
+	chromiumURL, err := url.Parse(rawUrl)
+	if err != nil {
+		s.logger.Warn("could not parse ws URL form chromium response", "sessionID", sessionID, "url", rawUrl, "err", err)
+		rw.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+
+	s.logger.Debug("Proxying WS connection", "sessionID", sessionID, "chromiumURL", rawUrl)
+
+	wsp := websocketproxy.WebsocketProxy{
+		Backend: func(r *http.Request) *url.URL {
+			return chromiumURL
+		},
+	}
+	wsp.ServeHTTP(rw, r)
+}
+
 // replaceHost returns a new url with its hostname replaced with host. The port is kept as it is.
 func replaceHost(urlStr, host string) (string, error) {
 	parsedURL, err := url.Parse(urlStr)