From 3f3c2d9387d06fa27666067bc7f9984b98a42268 Mon Sep 17 00:00:00 2001
From: Andrei Mihu <andrei@heroiclabs.com>
Date: Fri, 11 Nov 2022 15:58:06 +0000
Subject: [PATCH] Add support for custom response headers set in server
 configuration.

---
 CHANGELOG.md     |  1 +
 server/api.go    | 19 ++++++++++++++++++-
 server/config.go | 13 +++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ef7096033..e56fbe972 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@ The format is based on [keep a changelog](http://keepachangelog.com) and this pr
 ## [Unreleased]
 ### Added
 - Allow the socket acceptor to read session tokens from request headers.
+- Add support for custom response headers set in server configuration.
 
 ### Changed
 - Stricter validation of limit in runtime storage list operations.
diff --git a/server/api.go b/server/api.go
index 977d2c659..af482b82e 100644
--- a/server/api.go
+++ b/server/api.go
@@ -253,13 +253,30 @@ func StartApiServer(logger *zap.Logger, startupLogger *zap.Logger, db *sql.DB, p
 	CORSMethods := handlers.AllowedMethods([]string{"GET", "HEAD", "POST", "PUT", "DELETE"})
 	handlerWithCORS := handlers.CORS(CORSHeaders, CORSOrigins, CORSMethods)(grpcGatewayRouter)
 
+	// Enable configured response headers, if any are set. Do not override values that may have been set by server processing.
+	optionalResponseHeaderHandler := handlerWithCORS
+	if headers := config.GetSocket().Headers; len(headers) > 0 {
+		optionalResponseHeaderHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Preemptively set custom response headers. Further processing will override them if needed for proper functionality.
+			wHeaders := w.Header()
+			for key, value := range headers {
+				if wHeaders.Get(key) == "" {
+					wHeaders.Set(key, value)
+				}
+			}
+
+			// Allow core server processing to handle the request.
+			handlerWithCORS.ServeHTTP(w, r)
+		})
+	}
+
 	// Set up and start GRPC Gateway server.
 	s.grpcGatewayServer = &http.Server{
 		ReadTimeout:    time.Millisecond * time.Duration(int64(config.GetSocket().ReadTimeoutMs)),
 		WriteTimeout:   time.Millisecond * time.Duration(int64(config.GetSocket().WriteTimeoutMs)),
 		IdleTimeout:    time.Millisecond * time.Duration(int64(config.GetSocket().IdleTimeoutMs)),
 		MaxHeaderBytes: 5120,
-		Handler:        handlerWithCORS,
+		Handler:        optionalResponseHeaderHandler,
 	}
 	if config.GetSocket().TLSCert != nil {
 		s.grpcGatewayServer.TLSConfig = &tls.Config{Certificates: config.GetSocket().TLSCert}
diff --git a/server/config.go b/server/config.go
index 88f882750..c35888799 100644
--- a/server/config.go
+++ b/server/config.go
@@ -352,6 +352,17 @@ func CheckConfig(logger *zap.Logger, config Config) map[string]string {
 		configWarnings["runtime.read_only_globals"] = "Deprecated configuration parameter"
 	}
 
+	if l := len(config.GetSocket().ResponseHeaders); l > 0 {
+		config.GetSocket().Headers = make(map[string]string, l)
+		for _, header := range config.GetSocket().ResponseHeaders {
+			parts := strings.SplitN(header, "=", 2)
+			if len(parts) != 2 {
+				logger.Fatal("Response headers configuration invalid, format must be 'key=value'", zap.String("param", "socket.response_headers"))
+			}
+			config.GetSocket().Headers[parts[0]] = parts[1]
+		}
+	}
+
 	// Log warnings for SSL usage.
 	if config.GetSocket().SSLCertificate != "" && config.GetSocket().SSLPrivateKey == "" {
 		logger.Fatal("SSL configuration invalid, specify both socket.ssl_certificate and socket.ssl_private_key", zap.String("param", "socket.ssl_certificate"))
@@ -660,6 +671,8 @@ type SocketConfig struct {
 	OutgoingQueueSize    int               `yaml:"outgoing_queue_size" json:"outgoing_queue_size" usage:"The maximum number of messages waiting to be sent to the client. If this is exceeded the client is considered too slow and will disconnect. Used when processing real-time connections."`
 	SSLCertificate       string            `yaml:"ssl_certificate" json:"ssl_certificate" usage:"Path to certificate file if you want the server to use SSL directly. Must also supply ssl_private_key. NOT recommended for production use."`
 	SSLPrivateKey        string            `yaml:"ssl_private_key" json:"ssl_private_key" usage:"Path to private key file if you want the server to use SSL directly. Must also supply ssl_certificate. NOT recommended for production use."`
+	ResponseHeaders      []string          `yaml:"response_headers" json:"response_headers" usage:"Additional headers to send to clients with every response. Values here are only used if the response would not otherwise contain a value for the specified headers."`
+	Headers              map[string]string `yaml:"-" json:"-"` // Created by parsing ResponseHeaders above, not set from input args directly.
 	CertPEMBlock         []byte            `yaml:"-" json:"-"` // Created by fully reading the file contents of SSLCertificate, not set from input args directly.
 	KeyPEMBlock          []byte            `yaml:"-" json:"-"` // Created by fully reading the file contents of SSLPrivateKey, not set from input args directly.
 	TLSCert              []tls.Certificate `yaml:"-" json:"-"` // Created by processing CertPEMBlock and KeyPEMBlock, not set from input args directly.
-- 
GitLab