query-orchestration/internal/cognitoauth/middleware_test.go

package cognitoauth

import (
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/labstack/echo/v4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestJWTAuthMiddleware tests the JWTAuthMiddleware function to ensure it correctly processes
// authentication tokens from both headers and cookies. It verifies that:
// - Public paths (/login, /callback, /home) are skipped without authentication
// - The /logout path properly clears authentication cookies
// - Requests with Authorization headers are passed through
// - Requests with auth_token cookies have their tokens moved to Authorization headers
// - Requests without authentication are redirected to the login page
func TestJWTAuthMiddleware(t *testing.T) {
	// Setup
	e := echo.New()
	config := &mockConfigProvider{
		region:     "us-east-2",
		userPoolID: "us-east-2_testpool",
	}

	// Setup logger
	logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
		Level: slog.LevelDebug,
	})
	logger := slog.New(logHandler)
	config.SetAuthLogger(logger)

	// Create middleware
	middleware := JWTAuthMiddleware(config)

	// Test handler that simply returns success
	testHandler := func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	}

	// Test cases
	tests := []struct {
		name           string
		path           string
		setupRequest   func(*http.Request)
		expectStatus   int
		expectLocation string
		expectHeader   string
	}{
		{
			name: "skip login path",
			path: "/login",
			setupRequest: func(req *http.Request) {
				// No setup needed
			},
			expectStatus: http.StatusOK,
		},
		{
			name: "skip callback path",
			path: "/callback",
			setupRequest: func(req *http.Request) {
				// No setup needed
			},
			expectStatus: http.StatusOK,
		},
		{
			name: "skip home path",
			path: "/home",
			setupRequest: func(req *http.Request) {
				// No setup needed
			},
			expectStatus: http.StatusOK,
		},
		{
			name: "logout path sets empty cookie",
			path: "/logout",
			setupRequest: func(req *http.Request) {
				// No setup needed
			},
			expectStatus: http.StatusOK,
		},
		{
			name: "use authorization header",
			path: "/api/test",
			setupRequest: func(req *http.Request) {
				req.Header.Set("Authorization", "Bearer test-token")
			},
			expectStatus: http.StatusOK,
		},
		{
			name: "use cookie if no authorization header",
			path: "/api/test",
			setupRequest: func(req *http.Request) {
				cookie := &http.Cookie{
					Name:  "auth_token",
					Value: "cookie-token",
					Path:  "/",
				}
				req.AddCookie(cookie)
			},
			expectStatus: http.StatusOK,
			expectHeader: "Bearer cookie-token",
		},
		{
			name: "continue to next middleware if no token",
			path: "/api/test",
			setupRequest: func(req *http.Request) {
				// No auth header or cookie
			},
			expectStatus: http.StatusOK, // Should continue to test handler
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, tt.path, nil)
			rec := httptest.NewRecorder()

			// Setup request (add headers, cookies, etc.)
			tt.setupRequest(req)

			// Create Echo context
			c := e.NewContext(req, rec)
			c.SetPath(tt.path)

			// Run middleware
			err := middleware(testHandler)(c)
			if err != nil {
				t.Errorf("Middleware returned error: %v", err)
			}

			// Check status code
			if rec.Code != tt.expectStatus {
				t.Errorf("Expected status code %d, got %d", tt.expectStatus, rec.Code)
			}

			// Check location header (for redirects)
			if tt.expectLocation != "" {
				location := rec.Header().Get("Location")
				if location != tt.expectLocation {
					t.Errorf("Expected Location header %q, got %q", tt.expectLocation, location)
				}
			}

			// Check Authorization header (if token taken from cookie)
			if tt.expectHeader != "" {
				authHeader := req.Header.Get("Authorization")
				if authHeader != tt.expectHeader {
					t.Errorf("Expected Authorization header %q, got %q", tt.expectHeader, authHeader)
				}
			}

			// Special case for logout - check cookie
			if tt.path == "/logout" {
				cookies := rec.Result().Cookies()
				found := false
				for _, cookie := range cookies {
					if cookie.Name == "auth_token" {
						found = true
						if cookie.Value != "" {
							t.Errorf("Expected empty auth_token cookie, got %q", cookie.Value)
						}
						if !cookie.Expires.Before(time.Now()) {
							t.Errorf("Expected auth_token cookie to be expired, got %v", cookie.Expires)
						}
					}
				}
				if !found {
					t.Error("Expected auth_token cookie to be set for logout")
				}
			}
		})
	}
}

// TestTokenValidationMiddleware tests the TokenValidationMiddleware function to verify that it:
// - Allows access to public paths (/home, /logout) without validation
// - Properly handles OAuth callbacks by processing code and state parameters
// - Initiates login by redirecting to Cognito when accessing the login path
// - Returns 401 Unauthorized for protected paths without Authorization header
// - Attempts to validate tokens for protected paths with Authorization headers
func TestTokenValidationMiddleware(t *testing.T) {
	// Setup
	e := echo.New()
	config := &mockConfigProvider{
		region:     "us-east-2",
		userPoolID: "us-east-2_testpool",
	}

	// Setup logger
	logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
		Level: slog.LevelDebug,
	})
	logger := slog.New(logHandler)
	config.SetAuthLogger(logger)

	// Create middleware
	middleware := TokenValidationMiddleware(config, nil)

	// Test handler that simply returns success
	testHandler := func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	}

	// Test cases
	tests := []struct {
		name         string
		path         string
		setupRequest func(*http.Request)
		expectStatus int
		checkBody    func(*testing.T, []byte)
	}{
		{
			name: "skip public paths - home",
			path: "/home",
			setupRequest: func(req *http.Request) {
				// No setup needed
			},
			expectStatus: http.StatusOK,
		},
		{
			name: "skip public paths - logout",
			path: "/logout",
			setupRequest: func(req *http.Request) {
				// No setup needed
			},
			expectStatus: http.StatusOK,
		},
		{
			name: "handle oauth callback",
			path: "/callback",
			setupRequest: func(req *http.Request) {
				q := req.URL.Query()
				q.Add("code", "test-code")
				q.Add("state", "test-state")
				req.URL.RawQuery = q.Encode()
			},
			// This will attempt to call handleOAuthCallback which will fail
			// but we're just testing the routing logic
			expectStatus: http.StatusBadRequest,
		},
		{
			name: "handle login initiation",
			path: "/login",
			setupRequest: func(req *http.Request) {
				// No setup needed
			},
			// The login initiation actually redirects to the Cognito login page
			expectStatus: http.StatusFound, // 302 Found (redirect)
			// We can't check the exact URL since it depends on generateRandomString output
			// But we can check that a Location header exists
			checkBody: func(t *testing.T, body []byte) {
				// No body check needed, we'll verify the status code
			},
		},
		{
			name: "missing authorization header",
			path: "/api/test",
			setupRequest: func(req *http.Request) {
				// No auth header
			},
			expectStatus: http.StatusUnauthorized,
			checkBody: func(t *testing.T, body []byte) {
				var response map[string]string
				err := json.Unmarshal(body, &response)
				if err != nil {
					t.Errorf("Failed to parse response: %v", err)
				}
				if response["error"] != "Authorization required" {
					t.Errorf("Expected error 'Authorization required', got %q", response["error"])
				}
			},
		},
		{
			name: "invalid token",
			path: "/api/test",
			setupRequest: func(req *http.Request) {
				req.Header.Set("Authorization", "Bearer invalid-token")
			},
			expectStatus: http.StatusInternalServerError, // Since we can't mock JWKS fetching
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodGet, tt.path, nil)
			rec := httptest.NewRecorder()

			// Setup request
			tt.setupRequest(req)

			// Create Echo context
			c := e.NewContext(req, rec)
			c.SetPath(tt.path)

			// Run middleware
			_ = middleware(testHandler)(c)

			// Check status code
			if rec.Code != tt.expectStatus {
				t.Errorf("Expected status code %d, got %d", tt.expectStatus, rec.Code)
			}

			// Check response body if needed
			if tt.checkBody != nil {
				tt.checkBody(t, rec.Body.Bytes())
			}
		})
	}
}

// TestTokenValidationMiddleware_NilClientWithEnvConfigured documents the fail-open behavior
// when COGNITO_ENVIRONMENT_ID is configured but the Cognito client is nil (e.g., AWS config
// failed at startup). In this scenario, the environment check at middleware.go:403 is silently
// skipped because the guard `if cognitoClient != nil` evaluates to false.
//
// SECURITY NOTE: This is a known fail-open behavior. If COGNITO_ENVIRONMENT_ID is set but
// the Cognito client cannot be initialized (bad AWS credentials, network issue), the
// environment isolation check is bypassed while other security layers (JWT validation,
// Permit.io authorization) still operate. This means a user with a valid JWT from a
// different environment could access this server's resources.
func TestTokenValidationMiddleware_NilClientWithEnvConfigured(t *testing.T) {
	e := echo.New()

	// Config has environment ID set, simulating a deployment with env isolation enabled
	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-2",
			userPoolID: "us-east-2_testpool",
		},
		envID: "acmehealth",
	}
	logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug})
	config.SetAuthLogger(slog.New(logHandler))

	// cognitoClient is nil -- simulating AWS config failure at startup
	middleware := TokenValidationMiddleware(config, nil)

	testHandler := func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	}

	// A request to a protected path without auth gets 401 (token validation),
	// NOT a 403 from the env check. The env check is never reached because
	// cognitoClient is nil.
	req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
	rec := httptest.NewRecorder()
	c := e.NewContext(req, rec)
	c.SetPath("/api/test")

	_ = middleware(testHandler)(c)

	// The request fails at token validation (401), not env check (403).
	// This confirms the env check is skipped when cognitoClient is nil.
	if rec.Code != http.StatusUnauthorized {
		t.Errorf("Expected 401 (auth required), got %d -- env check should not trigger with nil client", rec.Code)
	}
}

// TestIsAuthOnlyPath verifies that certain paths bypass authorization while still requiring authentication.
// The /identity endpoint should be accessible to all authenticated users regardless of their Permit.io roles,
// allowing them to discover their own roles/permissions.
func TestIsAuthOnlyPath(t *testing.T) {
	tests := []struct {
		name     string
		path     string
		expected bool
	}{
		{
			name:     "identity endpoint should bypass authorization",
			path:     "/identity",
			expected: true,
		},
		{
			name:     "api endpoint should require authorization",
			path:     "/api/clients",
			expected: false,
		},
		{
			name:     "documents endpoint should require authorization",
			path:     "/documents",
			expected: false,
		},
		{
			name:     "health endpoint is not auth-only (it's public)",
			path:     "/health",
			expected: false,
		},
		{
			name:     "similar but different path should require authorization",
			path:     "/identity/something",
			expected: false,
		},
		{
			name:     "prefix match should not work",
			path:     "/identityinfo",
			expected: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := isAuthOnlyPath(tt.path)
			if result != tt.expected {
				t.Errorf("isAuthOnlyPath(%q) = %v, expected %v", tt.path, result, tt.expected)
			}
		})
	}
}

// --- Test infrastructure for performPermitIOAuthorization ---

// stubPermitChecker controls permit behavior in tests without a live PDP.
type stubPermitChecker struct {
	allowed bool
	err     error
	called  bool
}

func (s *stubPermitChecker) CheckPermission(userID, action, resourceType string) (bool, error) {
	s.called = true
	return s.allowed, s.err
}

// newAuthContext creates an Echo context with optional JWT claims pre-set.
func newAuthContext(e *echo.Echo, method, path, userSubject string) (echo.Context, *httptest.ResponseRecorder) {
	req := httptest.NewRequest(method, path, nil)
	rec := httptest.NewRecorder()
	c := e.NewContext(req, rec)
	if userSubject != "" {
		c.Set("user_claims", map[string]interface{}{"sub": userSubject})
	}
	return c, rec
}

// assertSingleJSONResponse verifies the response body contains exactly one JSON object.
// This is the critical regression guard for the double-write bug.
func assertSingleJSONResponse(t *testing.T, rec *httptest.ResponseRecorder) {
	t.Helper()
	body := strings.TrimSpace(rec.Body.String())
	require.NotEmpty(t, body, "response body should not be empty")

	decoder := json.NewDecoder(strings.NewReader(body))
	var first json.RawMessage
	require.NoError(t, decoder.Decode(&first), "response body should contain valid JSON")
	require.False(t, decoder.More(), "response body contains multiple JSON objects (double-write bug)")
}

// TestPerformPermitIOAuthorization tests the authorization function directly,
// verifying each error path returns a single *HTTPError (not c.JSON) so
// the caller writes exactly one response.
func TestPerformPermitIOAuthorization(t *testing.T) {
	e := echo.New()
	cfg := &mockConfigProvider{
		region:     "us-east-2",
		userPoolID: "us-east-2_testpool",
	}

	t.Run("NoUserSubject_Returns401", func(t *testing.T) {
		// No user_claims in context -- empty userSubject
		stub := &stubPermitChecker{allowed: false}
		c, _ := newAuthContext(e, http.MethodGet, "/document", "")

		err := performPermitIOAuthorization(c, "/document", cfg, stub)

		require.Error(t, err)
		httpErr, ok := err.(*HTTPError)
		require.True(t, ok, "expected *HTTPError, got %T", err)
		assert.Equal(t, http.StatusUnauthorized, httpErr.StatusCode)

		resp, ok := httpErr.Response.(map[string]string)
		require.True(t, ok)
		assert.Equal(t, "Authentication required", resp["error"])
		assert.False(t, stub.called, "permit checker should not be called without a user subject")
	})

	t.Run("PermitServiceError_Returns500", func(t *testing.T) {
		stub := &stubPermitChecker{allowed: false, err: errors.New("connection refused")}
		c, _ := newAuthContext(e, http.MethodGet, "/document", "user-123")

		err := performPermitIOAuthorization(c, "/document", cfg, stub)

		require.Error(t, err)
		httpErr, ok := err.(*HTTPError)
		require.True(t, ok, "expected *HTTPError, got %T", err)
		assert.Equal(t, http.StatusInternalServerError, httpErr.StatusCode)

		resp, ok := httpErr.Response.(map[string]string)
		require.True(t, ok)
		assert.Equal(t, "Authorization service unavailable", resp["error"])
		assert.True(t, stub.called)
	})

	t.Run("PermissionDenied_Returns403", func(t *testing.T) {
		stub := &stubPermitChecker{allowed: false, err: nil}
		c, _ := newAuthContext(e, http.MethodPost, "/document", "user-123")

		err := performPermitIOAuthorization(c, "/document", cfg, stub)

		require.Error(t, err)
		httpErr, ok := err.(*HTTPError)
		require.True(t, ok, "expected *HTTPError, got %T", err)
		assert.Equal(t, http.StatusForbidden, httpErr.StatusCode)

		resp, ok := httpErr.Response.(map[string]interface{})
		require.True(t, ok)
		assert.Equal(t, "Insufficient permissions", resp["error"])
		assert.Equal(t, "user-123", resp["user"])
		assert.Equal(t, "document", resp["resource"])
		assert.Equal(t, "post", resp["action"])
		assert.True(t, stub.called)
	})

	t.Run("PermissionGranted_ReturnsNil", func(t *testing.T) {
		stub := &stubPermitChecker{allowed: true, err: nil}
		c, _ := newAuthContext(e, http.MethodGet, "/document", "user-123")

		err := performPermitIOAuthorization(c, "/document", cfg, stub)

		assert.NoError(t, err)
		assert.True(t, stub.called)
	})

	t.Run("AuthOnlyPath_SkipsPermitCheck", func(t *testing.T) {
		stub := &stubPermitChecker{allowed: false, err: errors.New("should not be called")}
		c, _ := newAuthContext(e, http.MethodGet, "/identity", "user-123")

		err := performPermitIOAuthorization(c, "/identity", cfg, stub)

		assert.NoError(t, err)
		assert.False(t, stub.called, "permit checker should not be called for auth-only paths")
	})

	t.Run("SwaggerPath_SkipsPermitCheck", func(t *testing.T) {
		stub := &stubPermitChecker{allowed: false, err: errors.New("should not be called")}
		c, _ := newAuthContext(e, http.MethodGet, "/swagger/index.html", "user-123")

		err := performPermitIOAuthorization(c, "/swagger/index.html", cfg, stub)

		assert.NoError(t, err)
		assert.False(t, stub.called, "permit checker should not be called for swagger paths")
	})
}

// TestPermitIOAuthorization_SingleResponseIntegration tests the full middleware call site
// to confirm that *HTTPError from performPermitIOAuthorization flows through the caller
// and produces exactly one JSON response (not the pre-fix double-write).
func TestPermitIOAuthorization_SingleResponseIntegration(t *testing.T) {
	e := echo.New()
	cfg := &mockConfigProvider{
		region:          "us-east-2",
		userPoolID:      "us-east-2_testpool",
		noJWTValidation: true,
	}

	// handlerCalled tracks whether the downstream handler runs
	handlerCalled := false
	testHandler := func(c echo.Context) error {
		handlerCalled = true
		return c.JSON(http.StatusOK, map[string]string{"status": "ok"})
	}

	// simulateMiddlewareCallSite reproduces the caller logic from TokenValidationMiddleware
	// (lines 409-418) so we can test *HTTPError unwrapping without needing full JWT flow.
	simulateMiddlewareCallSite := func(c echo.Context, requestPath string, checker PermitChecker) error {
		if err := performPermitIOAuthorization(c, requestPath, cfg, checker); err != nil {
			if httpErr, ok := err.(*HTTPError); ok {
				return c.JSON(httpErr.StatusCode, httpErr.Response)
			}
			return err
		}
		return testHandler(c)
	}

	t.Run("PermitDenied_SingleJSONResponse", func(t *testing.T) {
		handlerCalled = false
		stub := &stubPermitChecker{allowed: false, err: nil}
		c, rec := newAuthContext(e, http.MethodGet, "/document", "user-123")

		err := simulateMiddlewareCallSite(c, "/document", stub)

		assert.NoError(t, err)
		assert.Equal(t, http.StatusForbidden, rec.Code)
		assert.False(t, handlerCalled, "handler should not run when authorization is denied")
		assertSingleJSONResponse(t, rec)
	})

	t.Run("PermitServiceError_SingleJSONResponse", func(t *testing.T) {
		handlerCalled = false
		stub := &stubPermitChecker{allowed: false, err: errors.New("connection refused")}
		c, rec := newAuthContext(e, http.MethodGet, "/document", "user-123")

		err := simulateMiddlewareCallSite(c, "/document", stub)

		assert.NoError(t, err)
		assert.Equal(t, http.StatusInternalServerError, rec.Code)
		assert.False(t, handlerCalled, "handler should not run when permit service fails")
		assertSingleJSONResponse(t, rec)
	})

	t.Run("NoUserSubject_SingleJSONResponse", func(t *testing.T) {
		handlerCalled = false
		stub := &stubPermitChecker{allowed: false}
		c, rec := newAuthContext(e, http.MethodGet, "/document", "") // no subject

		err := simulateMiddlewareCallSite(c, "/document", stub)

		assert.NoError(t, err)
		assert.Equal(t, http.StatusUnauthorized, rec.Code)
		assert.False(t, handlerCalled, "handler should not run without user subject")
		assertSingleJSONResponse(t, rec)
	})

	t.Run("PermitGranted_SingleJSONResponse", func(t *testing.T) {
		handlerCalled = false
		stub := &stubPermitChecker{allowed: true, err: nil}
		c, rec := newAuthContext(e, http.MethodGet, "/document", "user-123")

		err := simulateMiddlewareCallSite(c, "/document", stub)

		assert.NoError(t, err)
		assert.Equal(t, http.StatusOK, rec.Code)
		assert.True(t, handlerCalled, "handler should run when authorization is granted")
		assertSingleJSONResponse(t, rec)
	})

	t.Run("AuthOnlyPath_SingleJSONResponse", func(t *testing.T) {
		handlerCalled = false
		stub := &stubPermitChecker{allowed: false, err: errors.New("should not be called")}
		c, rec := newAuthContext(e, http.MethodGet, "/identity", "user-123")

		err := simulateMiddlewareCallSite(c, "/identity", stub)

		assert.NoError(t, err)
		assert.Equal(t, http.StatusOK, rec.Code)
		assert.True(t, handlerCalled, "handler should run for auth-only paths")
		assert.False(t, stub.called, "permit checker should not be called for auth-only paths")
		assertSingleJSONResponse(t, rec)
	})
}