From 5752ade642170015d1c03a7baaaf135547c5e875 Mon Sep 17 00:00:00 2001
From: yuli <stremovsky@gmail.com>
Date: Mon, 4 Nov 2024 09:14:18 +0200
Subject: [PATCH] add initContext() function

---
 src/bunker.go     | 4 ++--
 src/conf.go       | 5 +++++
 src/service.go    | 2 +-
 src/utils_test.go | 2 +-
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/bunker.go b/src/bunker.go
index 4af6d8f..9112e10 100644
--- a/src/bunker.go
+++ b/src/bunker.go
@@ -422,10 +422,10 @@ func (w *CustomResponseWriter) WriteHeader(statusCode int) {
 var statusCounter = 0
 var statusErrorCounter = 0
 
-func reqMiddleware(handler http.Handler) http.Handler {
+func (e mainEnv) reqMiddleware(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		//log.Printf("Set host %s\n", r.Host)
-		autocontext.Set(r, "host", r.Host)
+		e.initContext(r)
 		w.Header().Set("Access-Control-Allow-Origin", "*")
 		w.Header().Set("X-Frame-Options", "SAMEORIGIN")
 		w.Header().Set("Strict-Transport-Security", "max-age=31536000; includeSubDomains; preload")
diff --git a/src/conf.go b/src/conf.go
index 6a2cf91..77c19d4 100644
--- a/src/conf.go
+++ b/src/conf.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 
 	"github.com/julienschmidt/httprouter"
+	"github.com/securitybunker/databunker/src/autocontext"
 )
 
 func (e mainEnv) setupConfRouter(router *httprouter.Router) *httprouter.Router {
@@ -15,6 +16,10 @@ func (e mainEnv) setupConfRouter(router *httprouter.Router) *httprouter.Router {
 	return router
 }
 
+func (e mainEnv) initContext(r *http.Request) {
+	autocontext.Set(r, "host", r.Host)
+}
+
 func (e mainEnv) cookieSettings(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
 	resultJSON, scriptsJSON, _, err := e.db.getLegalBasisCookieConf()
 	if err != nil {
diff --git a/src/service.go b/src/service.go
index 4f9d8fe..f9da007 100644
--- a/src/service.go
+++ b/src/service.go
@@ -113,7 +113,7 @@ func loadService() {
 		},
 	}
 	listener := cfg.Server.Host + ":" + cfg.Server.Port
-	srv := &http.Server{Addr: listener, Handler: reqMiddleware(router), TLSConfig: tlsConfig}
+	srv := &http.Server{Addr: listener, Handler: e.reqMiddleware(router), TLSConfig: tlsConfig}
 
 	stop := make(chan os.Signal, 2)
 	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
diff --git a/src/utils_test.go b/src/utils_test.go
index 930aa69..9538a2a 100644
--- a/src/utils_test.go
+++ b/src/utils_test.go
@@ -94,7 +94,7 @@ func TestUtilGetJSONPost(t *testing.T) {
 }
 
 func TestUtilSMS(t *testing.T) {
-	server := httptest.NewServer(reqMiddleware(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+	server := httptest.NewServer(e.reqMiddleware(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		rw.Header().Set("Content-Type", "application/json")
 		rw.WriteHeader(200)
 		defer req.Body.Close()