golib/cookies/cookies_test.go

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)
		}
	})
}