From 407bfbbbbee4207ad52ee5b24369dfa405cbfa37 Mon Sep 17 00:00:00 2001
From: Michael J Feher <github@phearzero.com>
Date: Wed, 26 Apr 2023 01:03:41 -0500
Subject: [PATCH] Add CrossOrigin to the service (#68)

Adds Cross Origin as an optional flag to allow users to turn off the CORs check.
---
 main.go          |  6 ++++--
 server/server.go | 16 +++++++++++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/main.go b/main.go
index a2584fd..c8715a4 100644
--- a/main.go
+++ b/main.go
@@ -16,13 +16,14 @@ import (
 
 var version string = "0.0.0"
 
-func createServer(frontListenAddress string, frontendPath string, pty server.PTYHandler, sessionID string, allowTunneling bool) *server.TTYServer {
+func createServer(frontListenAddress string, frontendPath string, pty server.PTYHandler, sessionID string, allowTunneling bool, crossOrigin bool) *server.TTYServer {
 	config := ttyServer.TTYServerConfig{
 		FrontListenAddress: frontListenAddress,
 		FrontendPath:       frontendPath,
 		PTY:                pty,
 		SessionID:          sessionID,
 		AllowTunneling:     allowTunneling,
+		CrossOrigin:        crossOrigin,
 	}
 
 	server := ttyServer.NewTTYServer(config)
@@ -84,6 +85,7 @@ Flags:
 	detachKeys := flag.String("detach-keys", "ctrl-o,ctrl-c", "[c] Sequence of keys to press for closing the connection. Supported: https://godoc.org/github.com/moby/term#pkg-variables.")
 	allowTunneling := flag.Bool("A", false, "[s] Allow clients to create a TCP tunnel")
 	tunnelConfig := flag.String("L", "", "[c] TCP tunneling addresses: local_port:remote_host:remote_port. The client will listen on local_port for TCP connections, and will forward those to the from the server side to remote_host:remote_port")
+	crossOrgin := flag.Bool("cross-origin", false, "[s] Allow cross origin requests to the server")
 
 	verbose := flag.Bool("verbose", false, "Verbose logging")
 	flag.Usage = func() {
@@ -195,7 +197,7 @@ Flags:
 		pty = &nilPTY{}
 	}
 
-	server := createServer(*listenAddress, *frontendPath, pty, sessionID, *allowTunneling)
+	server := createServer(*listenAddress, *frontendPath, pty, sessionID, *allowTunneling, *crossOrgin)
 	if cols, rows, e := ptyMaster.GetWinSize(); e == nil {
 		server.WindowSize(cols, rows)
 	}
diff --git a/server/server.go b/server/server.go
index ffbb129..f3d8e83 100644
--- a/server/server.go
+++ b/server/server.go
@@ -41,6 +41,7 @@ type TTYServerConfig struct {
 	PTY                PTYHandler
 	SessionID          string
 	AllowTunneling     bool
+	CrossOrigin        bool
 }
 
 // TTYServer represents the instance of a tty server
@@ -126,7 +127,7 @@ func NewTTYServer(config TTYServerConfig) (server *TTYServer) {
 			server.handleWithTemplateHtml(w, r, "tty-share.in.html", templateModel)
 		})
 		routesHandler.HandleFunc(ttyWsPath, func(w http.ResponseWriter, r *http.Request) {
-			server.handleTTYWebsocket(w, r)
+			server.handleTTYWebsocket(w, r, config.CrossOrigin)
 		})
 		if server.config.AllowTunneling {
 			// tunnel websockets connection
@@ -151,16 +152,25 @@ func NewTTYServer(config TTYServerConfig) (server *TTYServer) {
 	return server
 }
 
-func (server *TTYServer) handleTTYWebsocket(w http.ResponseWriter, r *http.Request) {
+func (server *TTYServer) handleTTYWebsocket(w http.ResponseWriter, r *http.Request, crossOrigin bool) {
 	if r.Method != "GET" {
 		w.WriteHeader(http.StatusForbidden)
 		return
 	}
-
 	upgrader := websocket.Upgrader{
 		ReadBufferSize:  1024,
 		WriteBufferSize: 1024,
 	}
+	if crossOrigin {
+		upgrader = websocket.Upgrader{
+			ReadBufferSize:  1024,
+			WriteBufferSize: 1024,
+			CheckOrigin: func(r *http.Request) bool {
+				return true
+			},
+		}
+	}
+
 	conn, err := upgrader.Upgrade(w, r, nil)
 
 	if err != nil {