package cookies import ( "fmt" "net/http" "net/http/httptest" "strings" "testing" ) func TestSetCookie(t *testing.T) { tests := []struct { name string cookie string path string value string maxAge int expected string }{ { name: "basic cookie", cookie: "test", path: "/", value: "value", maxAge: 3600, expected: "test=value; Path=/; Max-Age=3600; HttpOnly", }, { name: "zero max age", cookie: "session", path: "/api", value: "abc123", maxAge: 0, expected: "session=abc123; Path=/api; HttpOnly", }, { name: "negative max age", cookie: "temp", path: "/", value: "temp", maxAge: -1, expected: "temp=temp; Path=/; Max-Age=0; HttpOnly", }, { name: "empty value", cookie: "empty", path: "/", value: "", maxAge: 3600, expected: "empty=; Path=/; Max-Age=3600; HttpOnly", }, { name: "special characters in value", cookie: "data", path: "/", value: "test@123!#$%", maxAge: 7200, expected: "data=test@123!#$%; Path=/; Max-Age=7200; HttpOnly", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { w := httptest.NewRecorder() SetCookie(w, tt.cookie, tt.path, tt.value, tt.maxAge) headers := w.Header()["Set-Cookie"] if len(headers) != 1 { t.Errorf("Expected 1 Set-Cookie header, got %d", len(headers)) return } // Parse the cookie header to check individual components cookieHeader := headers[0] // Check that all expected components are present if !strings.Contains(cookieHeader, tt.cookie+"="+tt.value) { t.Errorf("Expected cookie name/value not found in: %s", cookieHeader) } if !strings.Contains(cookieHeader, "Path="+tt.path) { t.Errorf("Expected path not found in: %s", cookieHeader) } if !strings.Contains(cookieHeader, "HttpOnly") { t.Errorf("Expected HttpOnly not found in: %s", cookieHeader) } if tt.maxAge != 0 { expectedMaxAge := fmt.Sprintf("Max-Age=%d", tt.maxAge) if tt.maxAge < 0 { expectedMaxAge = "Max-Age=0" // Go normalizes negative Max-Age to 0 } if !strings.Contains(cookieHeader, expectedMaxAge) { t.Errorf("Expected Max-Age not found in: %s", cookieHeader) } } }) } } func TestDeleteCookie(t *testing.T) { tests := []struct { name string cookie string path string expected string }{ { name: "basic deletion", cookie: "test", path: "/", expected: "test=; Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Max-Age=0; HttpOnly", }, { name: "delete with specific path", cookie: "session", path: "/api", expected: "session=; Path=/api; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Max-Age=0; HttpOnly", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { w := httptest.NewRecorder() DeleteCookie(w, tt.cookie, tt.path) headers := w.Header()["Set-Cookie"] if len(headers) != 1 { t.Errorf("Expected 1 Set-Cookie header, got %d", len(headers)) return } cookieHeader := headers[0] // Check deletion-specific components if !strings.Contains(cookieHeader, tt.cookie+"=") { t.Errorf("Expected cookie name not found in: %s", cookieHeader) } if !strings.Contains(cookieHeader, "Path="+tt.path) { t.Errorf("Expected path not found in: %s", cookieHeader) } if !strings.Contains(cookieHeader, "Max-Age=0") { t.Errorf("Expected Max-Age=0 not found in: %s", cookieHeader) } if !strings.Contains(cookieHeader, "Expires=") { t.Errorf("Expected Expires not found in: %s", cookieHeader) } if !strings.Contains(cookieHeader, "HttpOnly") { t.Errorf("Expected HttpOnly not found in: %s", cookieHeader) } }) } } func TestCheckPageFrom(t *testing.T) { tests := []struct { name string cookieValue string cookiePath string expectedResult string shouldSet bool }{ { name: "valid pagefrom cookie", cookieValue: "/dashboard", cookiePath: "/", expectedResult: "/dashboard", shouldSet: true, }, { name: "no pagefrom cookie", cookieValue: "", cookiePath: "", expectedResult: "/", shouldSet: false, }, { name: "empty pagefrom cookie", cookieValue: "", cookiePath: "/", expectedResult: "", shouldSet: true, }, { name: "pagefrom with query params", cookieValue: "/search?q=test", cookiePath: "/", expectedResult: "/search?q=test", shouldSet: true, }, { name: "pagefrom with special path", cookieValue: "/api/v1/users", cookiePath: "/api", expectedResult: "/api/v1/users", shouldSet: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { w := httptest.NewRecorder() r := &http.Request{ Header: make(http.Header), } if tt.shouldSet { cookie := &http.Cookie{ Name: "pagefrom", Value: tt.cookieValue, Path: tt.cookiePath, } r.AddCookie(cookie) } result := CheckPageFrom(w, r) if result != tt.expectedResult { t.Errorf("CheckPageFrom() = %v, want %v", result, tt.expectedResult) } // Verify that the cookie was deleted if tt.shouldSet { headers := w.Header()["Set-Cookie"] if len(headers) != 1 { t.Errorf("Expected 1 Set-Cookie header for deletion, got %d", len(headers)) return } cookieHeader := headers[0] if !strings.Contains(cookieHeader, "pagefrom=") { t.Errorf("Expected pagefrom cookie deletion not found in: %s", cookieHeader) } if !strings.Contains(cookieHeader, "Max-Age=0") { t.Errorf("Expected Max-Age=0 for deletion not found in: %s", cookieHeader) } } }) } } func TestSetPageFrom(t *testing.T) { tests := []struct { name string referer string trustedHost string expectedSet bool expectedValue string }{ { name: "valid trusted host referer", referer: "http://example.com/dashboard", trustedHost: "example.com", expectedSet: true, expectedValue: "/dashboard", }, { name: "valid trusted host with https", referer: "https://example.com/profile", trustedHost: "example.com", expectedSet: true, expectedValue: "/profile", }, { name: "untrusted host", referer: "http://evil.com/dashboard", trustedHost: "example.com", expectedSet: true, expectedValue: "/", }, { name: "empty path", referer: "http://example.com", trustedHost: "example.com", expectedSet: true, expectedValue: "/", }, { name: "login path - should not set", referer: "http://example.com/login", trustedHost: "example.com", expectedSet: false, expectedValue: "", }, { name: "register path - should not set", referer: "http://example.com/register", trustedHost: "example.com", expectedSet: false, expectedValue: "", }, { name: "invalid referer URL", referer: "not-a-url", trustedHost: "example.com", expectedSet: true, expectedValue: "/", }, { name: "empty referer", referer: "", trustedHost: "example.com", expectedSet: true, expectedValue: "/", }, { name: "root path", referer: "http://example.com/", trustedHost: "example.com", expectedSet: true, expectedValue: "/", }, { name: "path with query string", referer: "http://example.com/search?q=test", trustedHost: "example.com", expectedSet: true, expectedValue: "/search", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { w := httptest.NewRecorder() r := &http.Request{ Header: make(http.Header), } if tt.referer != "" { r.Header.Set("Referer", tt.referer) } SetPageFrom(w, r, tt.trustedHost) headers := w.Header()["Set-Cookie"] if tt.expectedSet { if len(headers) != 1 { t.Errorf("Expected 1 Set-Cookie header, got %d", len(headers)) return } cookieHeader := headers[0] if !strings.Contains(cookieHeader, "pagefrom="+tt.expectedValue) { t.Errorf("Expected pagefrom=%s not found in: %s", tt.expectedValue, cookieHeader) } } else { if len(headers) != 0 { t.Errorf("Expected no Set-Cookie header, got %d", len(headers)) } } }) } } func TestIntegration(t *testing.T) { // Test the complete flow: SetPageFrom -> CheckPageFrom t.Run("complete flow", func(t *testing.T) { // Step 1: Set pagefrom cookie w1 := httptest.NewRecorder() r1 := &http.Request{ Header: make(http.Header), } r1.Header.Set("Referer", "http://example.com/dashboard") SetPageFrom(w1, r1, "example.com") // Extract the cookie from the response headers1 := w1.Header()["Set-Cookie"] if len(headers1) != 1 { t.Errorf("Expected 1 Set-Cookie header, got %d", len(headers1)) return } // Verify the cookie was set correctly cookieHeader := headers1[0] if !strings.Contains(cookieHeader, "pagefrom=/dashboard") { t.Errorf("Expected pagefrom=/dashboard not found in: %s", cookieHeader) } // Step 2: Check pagefrom cookie (should delete it) w2 := httptest.NewRecorder() r2 := &http.Request{ Header: make(http.Header), } r2.AddCookie(&http.Cookie{ Name: "pagefrom", Value: "/dashboard", Path: "/", }) result := CheckPageFrom(w2, r2) if result != "/dashboard" { t.Errorf("Expected result /dashboard, got %s", result) } // Verify the cookie was deleted headers2 := w2.Header()["Set-Cookie"] if len(headers2) != 1 { t.Errorf("Expected 1 Set-Cookie header for deletion, got %d", len(headers2)) return } cookieHeader2 := headers2[0] // Check for deletion indicators (Max-Age=0 with Expires in the past) if !(strings.Contains(cookieHeader2, "Max-Age=0") && strings.Contains(cookieHeader2, "Expires=Thu, 01 Jan 1970")) { t.Errorf("Expected cookie deletion, got: %s", cookieHeader2) } }) }