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_WithForwardedProto_NotTrusted(t *testing.T) {
55 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
56 w.WriteHeader(http.StatusOK)
57 })
58
59 wrapped := SecurityHeaders(handler)
60
61 // X-Forwarded-Proto is client-spoofable and should NOT trigger HSTS
62 req := httptest.NewRequest(http.MethodGet, "/", nil)
63 req.Header.Set("X-Forwarded-Proto", "https")
64 rec := httptest.NewRecorder()
65 wrapped.ServeHTTP(rec, req)
66
67 assert.Empty(t, rec.Header().Get("Strict-Transport-Security"), "HSTS should not be set based on spoofable X-Forwarded-Proto")
68}
69
70func TestSecurityHeaders_PreservesExistingHeaders(t *testing.T) {
71 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
72 w.Header().Set("Content-Type", "application/json")
73 w.Header().Set("X-Custom-Header", "custom-value")
74 w.WriteHeader(http.StatusOK)
75 })
76
77 wrapped := SecurityHeaders(handler)
78
79 req := httptest.NewRequest(http.MethodGet, "/", nil)
80 rec := httptest.NewRecorder()
81 wrapped.ServeHTTP(rec, req)
82
83 // Security headers should be set
84 assert.Equal(t, "nosniff", rec.Header().Get("X-Content-Type-Options"))
85
86 // Handler's headers should also be present
87 assert.Equal(t, "application/json", rec.Header().Get("Content-Type"))
88 assert.Equal(t, "custom-value", rec.Header().Get("X-Custom-Header"))
89}
90
91func TestSecurityHeaders_CallsNextHandler(t *testing.T) {
92 called := false
93 handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
94 called = true
95 w.WriteHeader(http.StatusOK)
96 })
97
98 wrapped := SecurityHeaders(handler)
99
100 req := httptest.NewRequest(http.MethodGet, "/", nil)
101 rec := httptest.NewRecorder()
102 wrapped.ServeHTTP(rec, req)
103
104 assert.True(t, called, "next handler should have been called")
105}