1package middleware
2
3import (
4 "net/http"
5 "net/http/httptest"
6 "strings"
7 "testing"
8
9 "github.com/stretchr/testify/assert"
10)
11
12func TestNewResponseWriter(t *testing.T) {
13 originalWriter := httptest.NewRecorder()
14 rw := NewResponseWriter(originalWriter)
15
16 assert.NotNil(t, rw)
17 assert.Equal(t, originalWriter, rw.ResponseWriter)
18 assert.Equal(t, http.StatusOK, rw.StatusCode()) // Default status
19 assert.Equal(t, 0, rw.BytesWritten()) // No bytes written yet
20}
21
22func TestResponseWriter_WriteHeader(t *testing.T) {
23 originalWriter := httptest.NewRecorder()
24 rw := NewResponseWriter(originalWriter)
25
26 // Test setting status code
27 rw.WriteHeader(http.StatusNotFound)
28 assert.Equal(t, http.StatusNotFound, rw.StatusCode())
29
30 // Note: In HTTP, calling WriteHeader multiple times is not typical behavior.
31 // Our implementation captures the first call, which is the standard behavior.
32}
33
34func TestResponseWriter_Write(t *testing.T) {
35 originalWriter := httptest.NewRecorder()
36 rw := NewResponseWriter(originalWriter)
37
38 // Test writing data
39 data := []byte("Hello, World!")
40 n, err := rw.Write(data)
41
42 assert.NoError(t, err)
43 assert.Equal(t, len(data), n)
44 assert.Equal(t, len(data), rw.BytesWritten())
45
46 // Test multiple writes
47 moreData := []byte(" More data")
48 n2, err2 := rw.Write(moreData)
49
50 assert.NoError(t, err2)
51 assert.Equal(t, len(moreData), n2)
52 assert.Equal(t, len(data)+len(moreData), rw.BytesWritten())
53
54 // Verify data was written to underlying writer
55 result := originalWriter.Body.String()
56 assert.Equal(t, "Hello, World! More data", result)
57}
58
59func TestResponseWriter_StatusCode(t *testing.T) {
60 tests := []struct {
61 name string
62 statusCode int
63 }{
64 {"OK", http.StatusOK},
65 {"Not Found", http.StatusNotFound},
66 {"Internal Server Error", http.StatusInternalServerError},
67 {"Created", http.StatusCreated},
68 {"Bad Request", http.StatusBadRequest},
69 }
70
71 for _, tt := range tests {
72 t.Run(tt.name, func(t *testing.T) {
73 originalWriter := httptest.NewRecorder()
74 rw := NewResponseWriter(originalWriter)
75
76 rw.WriteHeader(tt.statusCode)
77 assert.Equal(t, tt.statusCode, rw.StatusCode())
78 })
79 }
80}
81
82func TestResponseWriter_BytesWritten(t *testing.T) {
83 originalWriter := httptest.NewRecorder()
84 rw := NewResponseWriter(originalWriter)
85
86 // Initially zero
87 assert.Equal(t, 0, rw.BytesWritten())
88
89 // Write some data
90 rw.Write([]byte("test"))
91 assert.Equal(t, 4, rw.BytesWritten())
92
93 // Write more data
94 rw.Write([]byte("ing"))
95 assert.Equal(t, 7, rw.BytesWritten())
96
97 // Write empty data
98 rw.Write([]byte(""))
99 assert.Equal(t, 7, rw.BytesWritten()) // Should remain the same
100}
101
102func TestGetClientIP(t *testing.T) {
103 tests := []struct {
104 name string
105 setupReq func(*http.Request)
106 expectedIP string
107 }{
108 {
109 name: "X-Forwarded-For single IP",
110 setupReq: func(r *http.Request) {
111 r.Header.Set("X-Forwarded-For", "192.168.1.100")
112 r.RemoteAddr = "10.0.0.1:12345"
113 },
114 expectedIP: "192.168.1.100",
115 },
116 {
117 name: "X-Forwarded-For multiple IPs",
118 setupReq: func(r *http.Request) {
119 r.Header.Set("X-Forwarded-For", "192.168.1.100, 10.0.0.1, 172.16.0.1")
120 r.RemoteAddr = "10.0.0.1:12345"
121 },
122 expectedIP: "192.168.1.100",
123 },
124 {
125 name: "X-Real-IP",
126 setupReq: func(r *http.Request) {
127 r.Header.Set("X-Real-IP", "203.0.113.195")
128 r.RemoteAddr = "10.0.0.1:12345"
129 },
130 expectedIP: "203.0.113.195",
131 },
132 {
133 name: "X-Forwarded",
134 setupReq: func(r *http.Request) {
135 r.Header.Set("X-Forwarded", "198.51.100.178")
136 r.RemoteAddr = "10.0.0.1:12345"
137 },
138 expectedIP: "198.51.100.178",
139 },
140 {
141 name: "RemoteAddr fallback",
142 setupReq: func(r *http.Request) {
143 r.RemoteAddr = "203.0.113.45:54321"
144 },
145 expectedIP: "203.0.113.45:54321",
146 },
147 {
148 name: "X-Forwarded-For priority over X-Real-IP",
149 setupReq: func(r *http.Request) {
150 r.Header.Set("X-Forwarded-For", "1.2.3.4")
151 r.Header.Set("X-Real-IP", "5.6.7.8")
152 r.RemoteAddr = "10.0.0.1:12345"
153 },
154 expectedIP: "1.2.3.4",
155 },
156 {
157 name: "X-Real-IP priority over X-Forwarded",
158 setupReq: func(r *http.Request) {
159 r.Header.Set("X-Real-IP", "1.2.3.4")
160 r.Header.Set("X-Forwarded", "5.6.7.8")
161 r.RemoteAddr = "10.0.0.1:12345"
162 },
163 expectedIP: "1.2.3.4",
164 },
165 {
166 name: "Empty headers fallback to RemoteAddr",
167 setupReq: func(r *http.Request) {
168 r.Header.Set("X-Forwarded-For", "")
169 r.Header.Set("X-Real-IP", "")
170 r.RemoteAddr = "192.168.1.1:8080"
171 },
172 expectedIP: "192.168.1.1:8080",
173 },
174 }
175
176 for _, tt := range tests {
177 t.Run(tt.name, func(t *testing.T) {
178 req := httptest.NewRequest("GET", "/test", nil)
179 tt.setupReq(req)
180
181 ip := getClientIP(req)
182 assert.Equal(t, tt.expectedIP, ip)
183 })
184 }
185}
186
187func TestRequestLogger(t *testing.T) {
188 // Create a test handler
189 testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
190 w.WriteHeader(http.StatusOK)
191 w.Write([]byte("test response"))
192 })
193
194 // Apply the middleware
195 handler := RequestLogger(testHandler)
196
197 // Create test request
198 req := httptest.NewRequest("GET", "/test/path?param=value", nil)
199 req.Header.Set("User-Agent", "test-agent")
200 req.Header.Set("Referer", "https://example.com")
201 req.Header.Set("X-Forwarded-For", "192.168.1.100")
202 req.Host = "test.example.com"
203
204 rr := httptest.NewRecorder()
205 handler.ServeHTTP(rr, req)
206
207 // Verify response
208 assert.Equal(t, http.StatusOK, rr.Code)
209 assert.Equal(t, "test response", rr.Body.String())
210
211 // Note: In a real scenario, you'd want to capture the log output
212 // For now, we just verify the middleware doesn't break the request flow
213}
214
215func TestContainerRequestLogger(t *testing.T) {
216 containerID := "container123"
217 domain := "app.example.com"
218
219 // Create a test handler
220 testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
221 w.WriteHeader(http.StatusCreated)
222 w.Write([]byte("container response"))
223 })
224
225 // Apply the middleware
226 middleware := ContainerRequestLogger(containerID, domain)
227 handler := middleware(testHandler)
228
229 // Create test request
230 req := httptest.NewRequest("POST", "/api/data", strings.NewReader("test data"))
231 req.Header.Set("Content-Type", "application/json")
232 req.Header.Set("X-Real-IP", "10.0.0.5")
233
234 rr := httptest.NewRecorder()
235 handler.ServeHTTP(rr, req)
236
237 // Verify response
238 assert.Equal(t, http.StatusCreated, rr.Code)
239 assert.Equal(t, "container response", rr.Body.String())
240}
241
242func TestPanicRecovery(t *testing.T) {
243 tests := []struct {
244 name string
245 handler http.HandlerFunc
246 expectPanic bool
247 expectedCode int
248 expectedBody string
249 }{
250 {
251 name: "normal operation",
252 handler: func(w http.ResponseWriter, r *http.Request) {
253 w.WriteHeader(http.StatusOK)
254 w.Write([]byte("success"))
255 },
256 expectPanic: false,
257 expectedCode: http.StatusOK,
258 expectedBody: "success",
259 },
260 {
261 name: "panic recovery",
262 handler: func(w http.ResponseWriter, r *http.Request) {
263 panic("test panic")
264 },
265 expectPanic: true,
266 expectedCode: http.StatusInternalServerError,
267 expectedBody: "Internal Server Error\n",
268 },
269 {
270 name: "nil panic",
271 handler: func(w http.ResponseWriter, r *http.Request) {
272 panic(nil)
273 },
274 expectPanic: true,
275 expectedCode: http.StatusInternalServerError,
276 expectedBody: "Internal Server Error\n",
277 },
278 {
279 name: "string panic",
280 handler: func(w http.ResponseWriter, r *http.Request) {
281 panic("string error")
282 },
283 expectPanic: true,
284 expectedCode: http.StatusInternalServerError,
285 expectedBody: "Internal Server Error\n",
286 },
287 {
288 name: "error panic",
289 handler: func(w http.ResponseWriter, r *http.Request) {
290 panic(assert.AnError)
291 },
292 expectPanic: true,
293 expectedCode: http.StatusInternalServerError,
294 expectedBody: "Internal Server Error\n",
295 },
296 }
297
298 for _, tt := range tests {
299 t.Run(tt.name, func(t *testing.T) {
300 // Apply panic recovery middleware
301 handler := PanicRecovery(tt.handler)
302
303 req := httptest.NewRequest("GET", "/test", nil)
304 rr := httptest.NewRecorder()
305
306 // This should not panic due to the middleware
307 assert.NotPanics(t, func() {
308 handler.ServeHTTP(rr, req)
309 })
310
311 assert.Equal(t, tt.expectedCode, rr.Code)
312 assert.Equal(t, tt.expectedBody, rr.Body.String())
313 })
314 }
315}
316
317func TestCORS(t *testing.T) {
318 tests := []struct {
319 name string
320 method string
321 expectedStatus int
322 checkHeaders bool
323 }{
324 {
325 name: "GET request",
326 method: "GET",
327 expectedStatus: http.StatusOK,
328 checkHeaders: true,
329 },
330 {
331 name: "POST request",
332 method: "POST",
333 expectedStatus: http.StatusOK,
334 checkHeaders: true,
335 },
336 {
337 name: "OPTIONS preflight",
338 method: "OPTIONS",
339 expectedStatus: http.StatusOK,
340 checkHeaders: true,
341 },
342 {
343 name: "PUT request",
344 method: "PUT",
345 expectedStatus: http.StatusOK,
346 checkHeaders: true,
347 },
348 {
349 name: "DELETE request",
350 method: "DELETE",
351 expectedStatus: http.StatusOK,
352 checkHeaders: true,
353 },
354 }
355
356 for _, tt := range tests {
357 t.Run(tt.name, func(t *testing.T) {
358 // Create test handler that should only be called for non-OPTIONS requests
359 called := false
360 testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
361 called = true
362 w.WriteHeader(http.StatusOK)
363 w.Write([]byte("test response"))
364 })
365
366 // Apply CORS middleware
367 handler := CORS(testHandler)
368
369 req := httptest.NewRequest(tt.method, "/test", nil)
370 rr := httptest.NewRecorder()
371 handler.ServeHTTP(rr, req)
372
373 assert.Equal(t, tt.expectedStatus, rr.Code)
374
375 if tt.checkHeaders {
376 assert.Equal(t, "*", rr.Header().Get("Access-Control-Allow-Origin"))
377 assert.Equal(t, "GET, POST, PUT, DELETE, OPTIONS", rr.Header().Get("Access-Control-Allow-Methods"))
378 assert.Equal(t, "Content-Type, Authorization", rr.Header().Get("Access-Control-Allow-Headers"))
379 }
380
381 // For OPTIONS request, the handler should not be called
382 if tt.method == "OPTIONS" {
383 assert.False(t, called, "Handler should not be called for OPTIONS request")
384 assert.Empty(t, rr.Body.String())
385 } else {
386 assert.True(t, called, "Handler should be called for non-OPTIONS request")
387 assert.Equal(t, "test response", rr.Body.String())
388 }
389 })
390 }
391}
392
393func TestChain(t *testing.T) {
394 // Create test middlewares that add headers
395 middleware1 := func(next http.Handler) http.Handler {
396 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
397 w.Header().Set("X-Middleware-1", "applied")
398 next.ServeHTTP(w, r)
399 })
400 }
401
402 middleware2 := func(next http.Handler) http.Handler {
403 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
404 w.Header().Set("X-Middleware-2", "applied")
405 next.ServeHTTP(w, r)
406 })
407 }
408
409 middleware3 := func(next http.Handler) http.Handler {
410 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
411 w.Header().Set("X-Middleware-3", "applied")
412 next.ServeHTTP(w, r)
413 })
414 }
415
416 // Create final handler
417 finalHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
418 w.WriteHeader(http.StatusOK)
419 w.Write([]byte("final handler"))
420 })
421
422 tests := []struct {
423 name string
424 middlewares []func(http.Handler) http.Handler
425 expectHeaders map[string]string
426 }{
427 {
428 name: "no middlewares",
429 middlewares: []func(http.Handler) http.Handler{},
430 expectHeaders: map[string]string{},
431 },
432 {
433 name: "single middleware",
434 middlewares: []func(http.Handler) http.Handler{middleware1},
435 expectHeaders: map[string]string{
436 "X-Middleware-1": "applied",
437 },
438 },
439 {
440 name: "multiple middlewares",
441 middlewares: []func(http.Handler) http.Handler{middleware1, middleware2, middleware3},
442 expectHeaders: map[string]string{
443 "X-Middleware-1": "applied",
444 "X-Middleware-2": "applied",
445 "X-Middleware-3": "applied",
446 },
447 },
448 }
449
450 for _, tt := range tests {
451 t.Run(tt.name, func(t *testing.T) {
452 // Apply chain
453 chain := Chain(tt.middlewares...)
454 handler := chain(finalHandler)
455
456 req := httptest.NewRequest("GET", "/test", nil)
457 rr := httptest.NewRecorder()
458 handler.ServeHTTP(rr, req)
459
460 assert.Equal(t, http.StatusOK, rr.Code)
461 assert.Equal(t, "final handler", rr.Body.String())
462
463 // Check expected headers
464 for header, value := range tt.expectHeaders {
465 assert.Equal(t, value, rr.Header().Get(header))
466 }
467 })
468 }
469}
470
471func TestChain_ExecutionOrder(t *testing.T) {
472 // Test that middlewares are executed in the correct order
473 var executionOrder []string
474
475 middleware1 := func(next http.Handler) http.Handler {
476 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
477 executionOrder = append(executionOrder, "middleware1-before")
478 next.ServeHTTP(w, r)
479 executionOrder = append(executionOrder, "middleware1-after")
480 })
481 }
482
483 middleware2 := func(next http.Handler) http.Handler {
484 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
485 executionOrder = append(executionOrder, "middleware2-before")
486 next.ServeHTTP(w, r)
487 executionOrder = append(executionOrder, "middleware2-after")
488 })
489 }
490
491 finalHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
492 executionOrder = append(executionOrder, "final-handler")
493 w.WriteHeader(http.StatusOK)
494 })
495
496 // Apply chain
497 chain := Chain(middleware1, middleware2)
498 handler := chain(finalHandler)
499
500 req := httptest.NewRequest("GET", "/test", nil)
501 rr := httptest.NewRecorder()
502 handler.ServeHTTP(rr, req)
503
504 expectedOrder := []string{
505 "middleware1-before",
506 "middleware2-before",
507 "final-handler",
508 "middleware2-after",
509 "middleware1-after",
510 }
511
512 assert.Equal(t, expectedOrder, executionOrder)
513}
514
515func TestMiddleware_Integration(t *testing.T) {
516 // Test multiple middlewares working together
517 finalHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
518 w.WriteHeader(http.StatusOK)
519 w.Write([]byte("integration test"))
520 })
521
522 // Chain multiple middlewares
523 handler := Chain(
524 PanicRecovery,
525 CORS,
526 RequestLogger,
527 )(finalHandler)
528
529 req := httptest.NewRequest("GET", "/integration", nil)
530 req.Header.Set("X-Forwarded-For", "192.168.1.1")
531
532 rr := httptest.NewRecorder()
533 handler.ServeHTTP(rr, req)
534
535 assert.Equal(t, http.StatusOK, rr.Code)
536 assert.Equal(t, "integration test", rr.Body.String())
537
538 // Check CORS headers are present
539 assert.Equal(t, "*", rr.Header().Get("Access-Control-Allow-Origin"))
540 assert.Equal(t, "GET, POST, PUT, DELETE, OPTIONS", rr.Header().Get("Access-Control-Allow-Methods"))
541}