Repository: nats-io/nats-server
Pull request #5230 (2024-03-19 20:05:58)

[ADD] WebSocket custom response headers
<!-- Please make sure to read CONTRIBUTING.md, then delete this notice and replace it with your PR description. The below sign-off certifies that the contribution is your original work and that you license the work to the project under the Apache-2.0 license. We cannot accept contributions without it. -->

Allow to configure extra headers for the WebSocket Upgrade Response:

For Example:
```
websocket {
    port: 8080
    no_tls: true
    headers {
        Strict-Transport-Security: "max-age=15768000 ; includeSubDomains"
        X-Frame-Options: "DENY"
    }
}
```


Related Issue: https://github.com/nats-io/nats-server/issues/5226

Signed-off-by: Ramon Berrutti <ramonberrutti@gmail.com>


ramonberrutti: @kozlovic I reimplement to support equaly between the config and the server options. As we don't support reloads, I precalculate the headers before starting the WebSocket server. Let me know what you think. In a bit I also include the missing tests.

Patch:
diff --git a/server/monitor.go b/server/monitor.go
index a5721956..49ac5ab8 100644
--- a/server/monitor.go
+++ b/server/monitor.go
@@ -2934,7 +2934,7 @@ func (s *Server) Jsz(opts *JSzOptions) (*JSInfo, error) {
 	if opts.Consumer {
 		opts.Streams = true
 	}
-	if opts.Streams && opts.Account == _EMPTY_ {
+	if opts.Streams {
 		opts.Accounts = true
 	}
 
diff --git a/server/opts.go b/server/opts.go
index a39c6187..2ca6f846 100644
--- a/server/opts.go
+++ b/server/opts.go
@@ -498,6 +498,10 @@ type WebsocketOpts struct {
 	// time needed for the TLS Handshake.
 	HandshakeTimeout time.Duration
 
+	// Headers to be added to the upgrade response.
+	// Useful for adding custom headers like Strict-Transport-Security.
+	Headers map[string]string
+
 	// Snapshot of configured TLS options.
 	tlsConfigOpts *TLSConfigOpts
 }
@@ -4810,6 +4814,23 @@ func parseWebsocket(v interface{}, o *Options, errors *[]error, warnings *[]erro
 			o.Websocket.TokenCookie = mv.(string)
 		case "no_auth_user":
 			o.Websocket.NoAuthUser = mv.(string)
+		case "headers":
+			m, ok := mv.(map[string]interface{})
+			if !ok {
+				err := &configErr{tk, fmt.Sprintf("error parsing headers: unsupported type %T", mv)}
+				*errors = append(*errors, err)
+				continue
+			}
+			o.Websocket.Headers = make(map[string]string)
+			for key, val := range m {
+				tk, val = unwrapValue(val, &lt)
+				if headerValue, ok := val.(string); !ok {
+					*errors = append(*errors, &configErr{tk, fmt.Sprintf("error parsing header key %s: unsupported type %T", key, val)})
+					continue
+				} else {
+					o.Websocket.Headers[key] = headerValue
+				}
+			}
 		default:
 			if !tk.IsUsedVariable() {
 				err := &unknownConfigFieldErr{
diff --git a/server/websocket.go b/server/websocket.go
index 7c9bb718..7efd976b 100644
--- a/server/websocket.go
+++ b/server/websocket.go
@@ -133,7 +133,8 @@ type srvWebsocket struct {
 	sameOrigin     bool
 	connectURLs    []string
 	connectURLsMap refCountedUrlSet
-	authOverride   bool // indicate if there is auth override in websocket config
+	authOverride   bool   // indicate if there is auth override in websocket config
+	rawHeaders     string // raw headers to be used in the upgrade response.
 }
 
 type allowedOrigin struct {
@@ -777,6 +778,9 @@ func (s *Server) wsUpgrade(w http.ResponseWriter, r *http.Request) (*wsUpgradeRe
 	if kind == MQTT {
 		p = append(p, wsMQTTSecProto...)
 	}
+	if s.websocket.rawHeaders != _EMPTY_ {
+		p = append(p, s.websocket.rawHeaders...)
+	}
 	p = append(p, _CRLF_...)
 
 	if _, err = conn.Write(p); err != nil {
@@ -1019,6 +1023,24 @@ func validateWebsocketOptions(o *Options) error {
 	if err := validatePinnedCerts(wo.TLSPinnedCerts); err != nil {
 		return fmt.Errorf("websocket: %v", err)
 	}
+
+	// Check for invalid headers here.
+	for key := range wo.Headers {
+		k := strings.ToLower(key)
+		switch k {
+		case "host",
+			"content-length",
+			"connection",
+			"upgrade",
+			"nats-no-masking":
+			return fmt.Errorf("websocket: invalid header %q not allowed", key)
+		}
+
+		if strings.HasPrefix(k, "sec-websocket-") {
+			return fmt.Errorf("websocket: invalid header %q, \"Sec-WebSocket-\" prefix not allowed", key)
+		}
+	}
+
 	return nil
 }
 
@@ -1050,6 +1072,21 @@ func (s *Server) wsSetOriginOptions(o *WebsocketOpts) {
 	}
 }
 
+// Calculate the raw headers for websocket upgrade response.
+func (s *Server) wsSetHeadersOptions(o *WebsocketOpts) {
+	var sb strings.Builder
+	for k, v := range o.Headers {
+		sb.WriteString(k)
+		sb.WriteString(": ")
+		sb.WriteString(v)
+		sb.WriteString(_CRLF_)
+	}
+	ws := &s.websocket
+	ws.mu.Lock()
+	defer ws.mu.Unlock()
+	ws.rawHeaders = sb.String()
+}
+
 // Given the websocket options, we check if any auth configuration
 // has been provided. If so, possibly create users/nkey users and
 // store them in s.websocket.users/nkeys.
@@ -1071,6 +1108,7 @@ func (s *Server) startWebsocketServer() {
 	o := &sopts.Websocket
 
 	s.wsSetOriginOptions(o)
+	s.wsSetHeadersOptions(o)
 
 	var hl net.Listener
 	var proto string


Test patch:
diff --git a/server/monitor_test.go b/server/monitor_test.go
index 307923a0..b2ff195a 100644
--- a/server/monitor_test.go
+++ b/server/monitor_test.go
@@ -4572,15 +4572,7 @@ func TestMonitorJsz(t *testing.T) {
 	})
 	t.Run("account-non-existing", func(t *testing.T) {
 		for _, url := range []string{monUrl1, monUrl2} {
-			info := readJsInfo(url + "?acc=DOES_NOT_EXIST")
-			if len(info.AccountDetails) != 0 {
-				t.Fatalf("expected no account to be returned by %s but got %v", url, info)
-			}
-		}
-	})
-	t.Run("account-non-existing-with-stream-details", func(t *testing.T) {
-		for _, url := range []string{monUrl1, monUrl2} {
-			info := readJsInfo(url + "?acc=DOES_NOT_EXIST&streams=true")
+			info := readJsInfo(url + "?acc=DOES_NOT_EXIT")
 			if len(info.AccountDetails) != 0 {
 				t.Fatalf("expected no account to be returned by %s but got %v", url, info)
 			}
diff --git a/server/websocket_test.go b/server/websocket_test.go
index 1eacf4b7..a989abcc 100644
--- a/server/websocket_test.go
+++ b/server/websocket_test.go
@@ -1442,6 +1442,61 @@ func TestWSCompressNegotiation(t *testing.T) {
 	}
 }
 
+func TestWSSetHeader(t *testing.T) {
+	opts := testWSOptions()
+	opts.Websocket.Headers = map[string]string{
+		"X-Header":         "some-value",
+		"X-Another-Header": "another-value",
+	}
+	s := &Server{opts: opts}
+	s.wsSetHeadersOptions(&opts.Websocket)
+	rw := &testResponseWriter{}
+	req := testWSCreateValidReq()
+	res, err := s.wsUpgrade(rw, req)
+	if res == nil || err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	buf := bufio.NewReader(&rw.conn.wbuf)
+	resp, err := http.ReadResponse(buf, req)
+	if err != nil {
+		t.Fatalf("Error reading request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	// Check that the response is a 101
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		t.Fatalf("Expected 101, got %v", resp.StatusCode)
+	}
+
+	headers := resp.Header.Clone()
+
+	// Compare all the headers
+	for k, v := range opts.Websocket.Headers {
+		if got := headers.Get(k); got != v {
+			t.Fatalf("Expected %q for header %q, got %q", v, k, got)
+		}
+		headers.Del(k)
+	}
+
+	// Check remain headers
+	for k, v := range map[string]string{
+		"Upgrade":              "websocket",
+		"Connection":           "Upgrade",
+		"Sec-Websocket-Accept": wsAcceptKey(req.Header.Get("Sec-Websocket-Key")),
+	} {
+		if got := headers.Get(k); got != v {
+			t.Fatalf("Expected %q for header %q, got %q", v, k, got)
+		}
+		headers.Del(k)
+	}
+
+	// Check that we have no more headers
+	if len(headers) > 0 {
+		t.Fatalf("Unexpected headers: %v", headers)
+	}
+}
+
 func TestWSParseOptions(t *testing.T) {
 	for _, test := range []struct {
 		name     string
@@ -1461,6 +1516,9 @@ func TestWSParseOptions(t *testing.T) {
 		{"bad allowed origins values", `websocket: { allowed_origins: [ {} ] }`, nil, "unsupported type in array"},
 		{"bad handshake timeout type", `websocket: { handshake_timeout: [] }`, nil, "unsupported type"},
 		{"bad handshake timeout duration", `websocket: { handshake_timeout: "abc" }`, nil, "invalid duration"},
+		{"bad header type", `websocket: { headers: 123 }`, nil, "unsupported type"},
+		{"bad header type", `websocket: { headers: [] }`, nil, "unsupported type"},
+		{"bad header value", `websocket: { headers: { "key": 123 } }`, nil, "unsupported type"},
 		{"unknown field", `websocket: { this_does_not_exist: 123 }`, nil, "unknown"},
 		// Positive tests
 		{"listen port only", `websocket { listen: 1234 }`, func(wo *WebsocketOpts) error {
@@ -1609,6 +1667,29 @@ func TestWSParseOptions(t *testing.T) {
 				}
 				return nil
 			}, ""},
+		{"headers block",
+			`
+			websocket {
+				headers {
+					"X-Header": "some-value"
+					"X-Another-Header": "another-value"
+				}
+			}
+			`, func(wo *WebsocketOpts) error {
+				if len(wo.Headers) != 2 {
+					return fmt.Errorf("Expected 2 headers, got %v", len(wo.Headers))
+				}
+
+				for k, v := range map[string]string{
+					"X-Header":         "some-value",
+					"X-Another-Header": "another-value",
+				} {
+					if got, ok := wo.Headers[k]; !ok || got != v {
+						return fmt.Errorf("Invalid value for %q: %q", k, got)
+					}
+				}
+				return nil
+			}, ""},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			conf := createConfFile(t, []byte(test.content))
@@ -1661,6 +1742,36 @@ func TestWSValidateOptions(t *testing.T) {
 			o.Websocket.Token = "mytoken"
 			return o
 		}, "websocket authentication token not compatible with presence of users/nkeys"},
+		{"headers with sec-websocket- prefix not allowed", func() *Options {
+			o := wso.Clone()
+			o.Websocket.Headers = map[string]string{"Sec-WebSocket-Key": "123"}
+			return o
+		}, `invalid header "Sec-WebSocket-Key", "Sec-WebSocket-" prefix not allowed`},
+		{"header with host", func() *Options {
+			o := wso.Clone()
+			o.Websocket.Headers = map[string]string{"Host": "http://localhost:8080"}
+			return o
+		}, `websocket: invalid header "Host" not allowed`},
+		{"header with content-length", func() *Options {
+			o := wso.Clone()
+			o.Websocket.Headers = map[string]string{"Content-Length": "0"}
+			return o
+		}, `websocket: invalid header "Content-Length" not allowed`},
+		{"header with connection", func() *Options {
+			o := wso.Clone()
+			o.Websocket.Headers = map[string]string{"Connection": "Upgrade"}
+			return o
+		}, `websocket: invalid header "Connection" not allowed`},
+		{"header with upgrade", func() *Options {
+			o := wso.Clone()
+			o.Websocket.Headers = map[string]string{"Upgrade": "websocket"}
+			return o
+		}, `websocket: invalid header "Upgrade" not allowed`},
+		{"header with Nats-No-Masking", func() *Options {
+			o := wso.Clone()
+			o.Websocket.Headers = map[string]string{"Nats-No-Masking": "false"}
+			return o
+		}, `websocket: invalid header "Nats-No-Masking" not allowed`},
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			err := validateWebsocketOptions(test.getOpts())
@@ -1792,12 +1903,13 @@ func TestWSAbnormalFailureOfWebServer(t *testing.T) {
 }
 
 type testWSClientOptions struct {
-	compress, web bool
-	host          string
-	port          int
-	extraHeaders  map[string][]string
-	noTLS         bool
-	path          string
+	compress, web        bool
+	host                 string
+	port                 int
+	extraHeaders         map[string][]string
+	noTLS                bool
+	path                 string
+	extraResponseHeaders map[string]string
 }
 
 func testNewWSClient(t testing.TB, o testWSClientOptions) (net.Conn, *bufio.Reader, []byte) {
@@ -1855,6 +1967,11 @@ func testNewWSClientWithError(t testing.TB, o testWSClientOptions) (net.Conn, *b
 	if resp.StatusCode != http.StatusSwitchingProtocols {
 		return nil, nil, nil, fmt.Errorf("Expected response status %v, got %v", http.StatusSwitchingProtocols, resp.StatusCode)
 	}
+	for k, v := range o.extraResponseHeaders {
+		if value := resp.Header.Get(k); value != v {
+			return nil, nil, nil, fmt.Errorf("Expected header %q to be %q, got %q", k, v, value)
+		}
+	}
 	var info []byte
 	if o.path == mqttWSPath {
 		if v := resp.Header[wsSecProto]; len(v) != 1 || v[0] != wsMQTTSecProtoVal {
@@ -3778,6 +3895,25 @@ func TestWSNkeyAuth(t *testing.T) {
 	}
 }
 
+func TestWSSetHeaderServer(t *testing.T) {
+	o := testWSOptions()
+	o.Websocket.Headers = map[string]string{
+		"X-Custom-Header": "custom-value",
+	}
+
+	s := RunServer(o)
+	defer s.Shutdown()
+
+	opts := testWSClientOptions{
+		host:                 o.Websocket.Host,
+		port:                 o.Websocket.Port,
+		extraResponseHeaders: o.Websocket.Headers,
+	}
+
+	c, _, _ := testNewWSClient(t, opts)
+	defer c.Close()
+}
+
 func TestWSJWTWithAllowedConnectionTypes(t *testing.T) {
 	o := testWSOptions()
 	setupAddTrusted(o)