From 1ce030af5d8af01c1c7d5edbeed8d58ba51a75a0 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@getprobo.com>
Date: Mon, 23 Jun 2025 09:07:01 -0700
Subject: [PATCH] Add support for extra header

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
---
 pkg/probod/api_config.go |  5 +++--
 pkg/probod/probod.go     |  1 +
 pkg/server/server.go     | 19 +++++++++++++------
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/pkg/probod/api_config.go b/pkg/probod/api_config.go
index 56bef0787..d2e42a4f1 100644
--- a/pkg/probod/api_config.go
+++ b/pkg/probod/api_config.go
@@ -20,7 +20,8 @@ type (
 	}
 
 	apiConfig struct {
-		Addr string     `json:"addr"`
-		Cors corsConfig `json:"cors"`
+		Addr              string            `json:"addr"`
+		Cors              corsConfig        `json:"cors"`
+		ExtraHeaderFields map[string]string `json:"extra-header-fields"`
 	}
 )
diff --git a/pkg/probod/probod.go b/pkg/probod/probod.go
index 592590f70..82a628094 100644
--- a/pkg/probod/probod.go
+++ b/pkg/probod/probod.go
@@ -224,6 +224,7 @@ func (impl *Implm) Run(
 	serverHandler, err := server.NewServer(
 		server.Config{
 			AllowedOrigins:    impl.cfg.Api.Cors.AllowedOrigins,
+			ExtraHeaderFields: impl.cfg.Api.ExtraHeaderFields,
 			Probo:             proboService,
 			Usrmgr:            usrmgrService,
 			ConnectorRegistry: defaultConnectorRegistry,
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 954365cbb..50e9f7071 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -34,6 +34,7 @@ import (
 // Config holds the configuration for the server
 type Config struct {
 	AllowedOrigins    []string
+	ExtraHeaderFields map[string]string
 	Probo             *probo.Service
 	Usrmgr            *usrmgr.Service
 	Auth              console_v1.AuthConfig
@@ -45,9 +46,10 @@ type Config struct {
 
 // Server represents the main server that handles both API and frontend requests
 type Server struct {
-	apiServer *api.Server
-	webServer *web.Server
-	router    *chi.Mux
+	apiServer         *api.Server
+	webServer         *web.Server
+	router            *chi.Mux
+	extraHeaderFields map[string]string
 }
 
 // NewServer creates a new server instance
@@ -77,9 +79,10 @@ func NewServer(cfg Config) (*Server, error) {
 	router := chi.NewRouter()
 
 	server := &Server{
-		apiServer: apiServer,
-		webServer: webServer,
-		router:    router,
+		apiServer:         apiServer,
+		webServer:         webServer,
+		router:            router,
+		extraHeaderFields: cfg.ExtraHeaderFields,
 	}
 
 	// Set up routes
@@ -106,5 +109,9 @@ func (s *Server) setupRoutes() {
 
 // ServeHTTP implements the http.Handler interface
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	for key, value := range s.extraHeaderFields {
+		w.Header().Set(key, value)
+	}
+
 	s.router.ServeHTTP(w, r)
 }