1package middleware
2
3import (
4 "crypto/tls"
5 "net/http"
6 "net/http/httptest"
7 "testing"
8
9 "github.com/stretchr/testify/assert"
10)
11
12func TestSecurityHeaders(t *testing.T) {
13 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
14 w.WriteHeader(http.StatusOK)
15 })
16
17 wrapped := SecurityHeaders(handler)
18
19 req := httptest.NewRequest(http.MethodGet, "/", nil)
20 rec := httptest.NewRecorder()
21 wrapped.ServeHTTP(rec, req)
22
23 assert.Equal(t, http.StatusOK, rec.Code)
24
25 // Verify all security headers are set
26 assert.Equal(t, "nosniff", rec.Header().Get("X-Content-Type-Options"))
27 assert.Equal(t, "DENY", rec.Header().Get("X-Frame-Options"))
28 assert.Equal(t, "1; mode=block", rec.Header().Get("X-XSS-Protection"))
29 assert.Equal(t, "strict-origin-when-cross-origin", rec.Header().Get("Referrer-Policy"))
30 assert.Equal(t, "geolocation=(), microphone=(), camera=()", rec.Header().Get("Permissions-Policy"))
31
32 // CSP should NOT be set globally (would break proxied sites)
33 assert.Empty(t, rec.Header().Get("Content-Security-Policy"))
34
35 // HSTS should NOT be set on plain HTTP
36 assert.Empty(t, rec.Header().Get("Strict-Transport-Security"))
37}
38
39func TestSecurityHeaders_HSTS_WithTLS(t *testing.T) {
40 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
41 w.WriteHeader(http.StatusOK)
42 })
43
44 wrapped := SecurityHeaders(handler)
45
46 req := httptest.NewRequest(http.MethodGet, "/", nil)
47 req.TLS = &tls.ConnectionState{} // Simulate TLS connection
48 rec := httptest.NewRecorder()
49 wrapped.ServeHTTP(rec, req)
50
51 assert.Equal(t, "max-age=31536000; includeSubDomains", rec.Header().Get("Strict-Transport-Security"))
52}
53
54func TestSecurityHeaders_HSTS_ForceWithoutTLS(t *testing.T) {
55 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
56 w.WriteHeader(http.StatusOK)
57 })
58
59 wrapped := SecurityHeadersWithOptions(true)(handler)
60
61 req := httptest.NewRequest(http.MethodGet, "/", nil)
62 // r.TLS is nil — no TLS on the connection, but forceHSTS is true
63 rec := httptest.NewRecorder()
64 wrapped.ServeHTTP(rec, req)
65
66 assert.Equal(t, "max-age=31536000; includeSubDomains", rec.Header().Get("Strict-Transport-Security"),
67 "HSTS should be set when forceHSTS is true even without TLS")
68}
69
70func TestSecurityHeaders_HSTS_WithForwardedProto_NotTrusted(t *testing.T) {
71 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
72 w.WriteHeader(http.StatusOK)
73 })
74
75 wrapped := SecurityHeaders(handler)
76
77 // X-Forwarded-Proto is client-spoofable and should NOT trigger HSTS
78 req := httptest.NewRequest(http.MethodGet, "/", nil)
79 req.Header.Set("X-Forwarded-Proto", "https")
80 rec := httptest.NewRecorder()
81 wrapped.ServeHTTP(rec, req)
82
83 assert.Empty(t, rec.Header().Get("Strict-Transport-Security"), "HSTS should not be set based on spoofable X-Forwarded-Proto")
84}
85
86func TestSecurityHeaders_PreservesExistingHeaders(t *testing.T) {
87 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
88 w.Header().Set("Content-Type", "application/json")
89 w.Header().Set("X-Custom-Header", "custom-value")
90 w.WriteHeader(http.StatusOK)
91 })
92
93 wrapped := SecurityHeaders(handler)
94
95 req := httptest.NewRequest(http.MethodGet, "/", nil)
96 rec := httptest.NewRecorder()
97 wrapped.ServeHTTP(rec, req)
98
99 // Security headers should be set
100 assert.Equal(t, "nosniff", rec.Header().Get("X-Content-Type-Options"))
101
102 // Handler's headers should also be present
103 assert.Equal(t, "application/json", rec.Header().Get("Content-Type"))
104 assert.Equal(t, "custom-value", rec.Header().Get("X-Custom-Header"))
105}
106
107func TestSecurityHeaders_CallsNextHandler(t *testing.T) {
108 called := false
109 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
110 called = true
111 w.WriteHeader(http.StatusOK)
112 })
113
114 wrapped := SecurityHeaders(handler)
115
116 req := httptest.NewRequest(http.MethodGet, "/", nil)
117 rec := httptest.NewRecorder()
118 wrapped.ServeHTTP(rec, req)
119
120 assert.True(t, called, "next handler should have been called")
121}