From 1831141deb737e2e9ab813f859e9b838783bd885 Mon Sep 17 00:00:00 2001
From: James Rouzier <jrouzier@inverse.ca>
Date: Mon, 8 Jan 2024 20:47:06 +0000
Subject: [PATCH] Add the ability to hook into server

---
 server/server.go         | 3 +++
 server/server_handler.go | 7 ++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/server/server.go b/server/server.go
index 8a702fce..d9007817 100644
--- a/server/server.go
+++ b/server/server.go
@@ -32,6 +32,7 @@ type Config struct {
 	Reverse   bool
 	KeepAlive time.Duration
 	TLS       TLSConfig
+	HTTPHook  http.Handler
 }
 
 // Server respresent a chisel service
@@ -45,6 +46,7 @@ type Server struct {
 	sessions     *settings.Users
 	sshConfig    *ssh.ServerConfig
 	users        *settings.UserIndex
+	httpHook     http.Handler
 }
 
 var upgrader = websocket.Upgrader{
@@ -60,6 +62,7 @@ func NewServer(c *Config) (*Server, error) {
 		httpServer: cnet.NewHTTPServer(),
 		Logger:     cio.NewLogger("server"),
 		sessions:   settings.NewUsers(),
+		httpHook:   c.HTTPHook,
 	}
 	server.Info = true
 	server.users = settings.NewUserIndex(server.Logger)
diff --git a/server/server_handler.go b/server/server_handler.go
index 952aa4d8..461d7afe 100644
--- a/server/server_handler.go
+++ b/server/server_handler.go
@@ -19,7 +19,7 @@ func (s *Server) handleClientHandler(w http.ResponseWriter, r *http.Request) {
 	//websockets upgrade AND has chisel prefix
 	upgrade := strings.ToLower(r.Header.Get("Upgrade"))
 	protocol := r.Header.Get("Sec-WebSocket-Protocol")
-	if upgrade == "websocket"  {
+	if upgrade == "websocket" {
 		if protocol == chshare.ProtocolVersion {
 			s.handleWebsocket(w, r)
 			return
@@ -42,6 +42,11 @@ func (s *Server) handleClientHandler(w http.ResponseWriter, r *http.Request) {
 		w.Write([]byte(chshare.BuildVersion))
 		return
 	}
+
+	if s.httpHook != nil {
+		s.httpHook.ServeHTTP(w, r)
+		return
+	}
 	//missing :O
 	w.WriteHeader(404)
 	w.Write([]byte("Not found"))