query-orchestration/internal/cognitoauth/envcheck_test.go

package cognitoauth

import (
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
	"github.com/labstack/echo/v4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"queryorchestration/internal/usermanagement"
)

// roundTripFunc is a function adapter for http.RoundTripper.
type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return f(req)
}

// newTestCognitoClientForEnvCheck creates a Cognito client with a custom HTTP transport.
func newTestCognitoClientForEnvCheck(t *testing.T, rt roundTripFunc) *cognitoidentityprovider.Client {
	t.Helper()

	cfg, err := awsconfig.LoadDefaultConfig(t.Context(),
		awsconfig.WithRegion("us-east-1"),
		awsconfig.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("AKID", "SECRET", "SESSION")),
		awsconfig.WithHTTPClient(&http.Client{Transport: rt}),
	)
	require.NoError(t, err)

	return cognitoidentityprovider.NewFromConfig(cfg)
}

// jsonResp builds an HTTP response with a JSON body.
func jsonResp(statusCode int, body interface{}) *http.Response {
	b, _ := json.Marshal(body)
	return &http.Response{
		StatusCode: statusCode,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       io.NopCloser(strings.NewReader(string(b))),
	}
}

// envCheckMockConfig extends mockConfigProvider to support a configurable environment ID.
type envCheckMockConfig struct {
	mockConfigProvider
	envID string
}

// GetCognitoEnvironmentID returns the configured environment ID.
func (m *envCheckMockConfig) GetCognitoEnvironmentID() string {
	return m.envID
}

// SetCognitoEnvironmentID sets the environment ID.
func (m *envCheckMockConfig) SetCognitoEnvironmentID(id string) {
	m.envID = id
}

// createEnvCheckEchoContext builds an echo.Context with user_claims set for testing.
func createEnvCheckEchoContext(sub string) echo.Context {
	e := echo.New()
	req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
	rec := httptest.NewRecorder()
	c := e.NewContext(req, rec)
	c.Set("user_claims", map[string]interface{}{
		"sub": sub,
	})
	return c
}

// TestCheckUserEnvironmentAccess_NoServerEnv verifies that when the server has no
// COGNITO_ENVIRONMENT_ID set, all users pass the environment check.
func TestCheckUserEnvironmentAccess_NoServerEnv(t *testing.T) {
	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "", // no environment set
	}

	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		t.Fatal("Cognito should NOT be called when server has no environment ID")
		return nil, nil
	})

	c := createEnvCheckEchoContext("sub-123")

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	assert.NoError(t, err, "all users should pass when server has no environment ID")
}

// TestCheckUserEnvironmentAccess_UserEnvMatchesServer verifies that when the user's
// environment_id matches the server's, the check passes.
func TestCheckUserEnvironmentAccess_UserEnvMatchesServer(t *testing.T) {
	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "acmehealth",
	}

	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		resp := map[string]interface{}{
			"Username":   "user@acme.com",
			"Enabled":    true,
			"UserStatus": "CONFIRMED",
			"UserAttributes": []map[string]string{
				{"Name": "sub", "Value": "sub-match"},
				{"Name": "email", "Value": "user@acme.com"},
				{"Name": usermanagement.CognitoEnvIDAttrName, "Value": "acmehealth"},
			},
		}
		return jsonResp(200, resp), nil
	})

	c := createEnvCheckEchoContext("sub-match")

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	assert.NoError(t, err, "user with matching environment should pass")
}

// TestCheckUserEnvironmentAccess_UntaggedUser verifies that an untagged user
// (no environment_id attribute) passes the check (backward compatibility).
func TestCheckUserEnvironmentAccess_UntaggedUser(t *testing.T) {
	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "acmehealth",
	}

	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		resp := map[string]interface{}{
			"Username":   "legacy@acme.com",
			"Enabled":    true,
			"UserStatus": "CONFIRMED",
			"UserAttributes": []map[string]string{
				{"Name": "sub", "Value": "sub-legacy"},
				{"Name": "email", "Value": "legacy@acme.com"},
				// No environment_id attribute
			},
		}
		return jsonResp(200, resp), nil
	})

	c := createEnvCheckEchoContext("sub-legacy")

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	assert.NoError(t, err, "untagged user should pass (backward compat)")
}

// TestCheckUserEnvironmentAccess_MismatchNotSuperAdmin verifies that a user with a
// mismatched environment_id who is NOT a super_admin gets a 403 error.
func TestCheckUserEnvironmentAccess_MismatchNotSuperAdmin(t *testing.T) {
	// Set up Permit.io mock server for role check
	permitMux := http.NewServeMux()
	permitMux.HandleFunc("/v2/api-key/scope", func(w http.ResponseWriter, r *http.Request) {
		resp := map[string]string{
			"organization_id": "test-org",
			"project_id":      "test-project",
			"environment_id":  "test-env",
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(resp)
	})
	permitMux.HandleFunc("/v2/facts/test-project/test-env/role_assignments", func(w http.ResponseWriter, r *http.Request) {
		// User has "editor" role, NOT super_admin
		resp := []map[string]interface{}{
			{"role": "editor", "user": "sub-mismatch", "tenant": "default"},
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(resp)
	})
	permitServer := httptest.NewServer(permitMux)
	defer permitServer.Close()

	t.Setenv("PERMIT_IO_BASE_URL", permitServer.URL)

	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "acmehealth",
	}

	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		resp := map[string]interface{}{
			"Username":   "user@ford.com",
			"Enabled":    true,
			"UserStatus": "CONFIRMED",
			"UserAttributes": []map[string]string{
				{"Name": "sub", "Value": "sub-mismatch"},
				{"Name": "email", "Value": "user@ford.com"},
				{"Name": usermanagement.CognitoEnvIDAttrName, "Value": "fordmotorco"},
			},
		}
		return jsonResp(200, resp), nil
	})

	c := createEnvCheckEchoContext("sub-mismatch")

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	assert.Error(t, err, "user with mismatched environment should be denied")
}

// TestCheckUserEnvironmentAccess_MismatchButSuperAdmin verifies that a super_admin user
// with a mismatched environment_id is allowed through (super_admins bypass env checks).
func TestCheckUserEnvironmentAccess_MismatchButSuperAdmin(t *testing.T) {
	// Set up Permit.io mock server for role check
	permitMux := http.NewServeMux()
	permitMux.HandleFunc("/v2/api-key/scope", func(w http.ResponseWriter, r *http.Request) {
		resp := map[string]string{
			"organization_id": "test-org",
			"project_id":      "test-project",
			"environment_id":  "test-env",
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(resp)
	})
	permitMux.HandleFunc("/v2/facts/test-project/test-env/role_assignments", func(w http.ResponseWriter, r *http.Request) {
		// User has super_admin role
		resp := []map[string]interface{}{
			{"role": "super_admin", "user": "sub-superadmin", "tenant": "default"},
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(resp)
	})
	permitServer := httptest.NewServer(permitMux)
	defer permitServer.Close()

	t.Setenv("PERMIT_IO_BASE_URL", permitServer.URL)

	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "acmehealth",
	}

	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		resp := map[string]interface{}{
			"Username":   "admin@ford.com",
			"Enabled":    true,
			"UserStatus": "CONFIRMED",
			"UserAttributes": []map[string]string{
				{"Name": "sub", "Value": "sub-superadmin"},
				{"Name": "email", "Value": "admin@ford.com"},
				{"Name": usermanagement.CognitoEnvIDAttrName, "Value": "fordmotorco"},
			},
		}
		return jsonResp(200, resp), nil
	})

	c := createEnvCheckEchoContext("sub-superadmin")

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	assert.NoError(t, err, "super_admin should bypass environment check")
}

// TestCheckUserEnvironmentAccess_CacheHit verifies that repeated calls for the same
// user sub use the cache and do NOT make additional Cognito/Permit.io API calls.
func TestCheckUserEnvironmentAccess_CacheHit(t *testing.T) {
	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "acmehealth",
	}

	callCount := 0
	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		callCount++
		resp := map[string]interface{}{
			"Username":   "user@acme.com",
			"Enabled":    true,
			"UserStatus": "CONFIRMED",
			"UserAttributes": []map[string]string{
				{"Name": "sub", "Value": "sub-cached"},
				{"Name": "email", "Value": "user@acme.com"},
				{"Name": usermanagement.CognitoEnvIDAttrName, "Value": "acmehealth"},
			},
		}
		return jsonResp(200, resp), nil
	})

	// First call -- should hit Cognito
	c1 := createEnvCheckEchoContext("sub-cached")
	err := CheckUserEnvironmentAccess(c1, config, cognitoClient)
	assert.NoError(t, err)
	firstCallCount := callCount

	// Second call with same sub -- should use cache, NOT call Cognito again
	c2 := createEnvCheckEchoContext("sub-cached")
	err = CheckUserEnvironmentAccess(c2, config, cognitoClient)
	assert.NoError(t, err)

	assert.Equal(t, firstCallCount, callCount,
		"second call should use cache and not make additional Cognito API calls")
}

// TestCheckUserEnvironmentAccess_MissingSub verifies that when user_claims has no
// "sub" field, CheckUserEnvironmentAccess returns 403 without calling Cognito.
func TestCheckUserEnvironmentAccess_MissingSub(t *testing.T) {
	ResetEnvCheckCache()

	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "acmehealth",
	}

	// Cognito client that fatals if called -- proves no Cognito call is made
	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		t.Fatal("Cognito should NOT be called when sub is missing")
		return nil, nil
	})

	// Create context with empty user_claims (no sub)
	e := echo.New()
	req := httptest.NewRequest(http.MethodGet, "/api/test", nil)
	rec := httptest.NewRecorder()
	c := e.NewContext(req, rec)
	c.Set("user_claims", map[string]interface{}{})

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	require.Error(t, err, "Missing sub should return an error")
	assert.Contains(t, err.Error(), "missing user identity")
}

// TestCheckUserEnvironmentAccess_NegativeCacheHit verifies that a denied user is
// cached and subsequent calls do NOT make additional Cognito API calls.
func TestCheckUserEnvironmentAccess_NegativeCacheHit(t *testing.T) {
	ResetEnvCheckCache()

	// Set up Permit.io mock -- user has "editor" role, NOT super_admin
	permitMux := http.NewServeMux()
	permitMux.HandleFunc("/v2/api-key/scope", func(w http.ResponseWriter, r *http.Request) {
		resp := map[string]string{
			"organization_id": "test-org",
			"project_id":      "test-project",
			"environment_id":  "test-env",
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(resp)
	})
	permitMux.HandleFunc("/v2/facts/test-project/test-env/role_assignments", func(w http.ResponseWriter, r *http.Request) {
		resp := []map[string]interface{}{
			{"role": "editor", "user": "sub-negcache", "tenant": "default"},
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(resp)
	})
	permitServer := httptest.NewServer(permitMux)
	defer permitServer.Close()
	t.Setenv("PERMIT_IO_BASE_URL", permitServer.URL)

	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_TestPool",
		},
		envID: "acmehealth",
	}

	cognitoCallCount := 0
	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		cognitoCallCount++
		resp := map[string]interface{}{
			"Username":   "user@ford.com",
			"Enabled":    true,
			"UserStatus": "CONFIRMED",
			"UserAttributes": []map[string]string{
				{"Name": "sub", "Value": "sub-negcache"},
				{"Name": "email", "Value": "user@ford.com"},
				{"Name": usermanagement.CognitoEnvIDAttrName, "Value": "fordmotorco"},
			},
		}
		return jsonResp(200, resp), nil
	})

	// First call -- should call Cognito, get denied, and cache the denial
	c1 := createEnvCheckEchoContext("sub-negcache")
	err := CheckUserEnvironmentAccess(c1, config, cognitoClient)
	assert.Error(t, err, "First call should be denied (env mismatch, not super_admin)")
	assert.Equal(t, 1, cognitoCallCount, "First call should hit Cognito once")

	// Second call with same sub -- should use negative cache, NOT call Cognito
	c2 := createEnvCheckEchoContext("sub-negcache")
	err = CheckUserEnvironmentAccess(c2, config, cognitoClient)
	assert.Error(t, err, "Second call should also be denied (from cache)")
	assert.Contains(t, err.Error(), "environment access denied")
	assert.Equal(t, 1, cognitoCallCount, "Second call should use cache, not call Cognito again")
}

// TestEnvCheckCacheTTL verifies the cache TTL constant is set to 30 minutes.
func TestEnvCheckCacheTTL(t *testing.T) {
	assert.Equal(t, 30*60, int(EnvCheckCacheTTL.Seconds()),
		"cache TTL should be 30 minutes")
}

// TestCheckUserEnvironmentAccess_ServiceAuthSource_UsesServicePool verifies
// that when the request was tagged with auth_source="service", the
// env check looks the sub up in pool B (the service-accounts pool)
// rather than pool A. We assert this by inspecting the AdminGetUser
// request body, which carries the pool id.
func TestCheckUserEnvironmentAccess_ServiceAuthSource_UsesServicePool(t *testing.T) {
	ResetEnvCheckCache()

	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:        "us-east-1",
			userPoolID:    "us-east-1_PoolA",
			svcUserPoolID: "us-east-1_PoolB",
		},
		envID: "acmehealth",
	}

	var seenPoolID string
	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		body, _ := io.ReadAll(req.Body)
		var parsed map[string]string
		_ = json.Unmarshal(body, &parsed)
		seenPoolID = parsed["UserPoolId"]
		resp := map[string]interface{}{
			"Username":   "svc-1",
			"Enabled":    true,
			"UserStatus": "CONFIRMED",
			"UserAttributes": []map[string]string{
				{"Name": "sub", "Value": "sub-svc"},
				{"Name": "email", "Value": "svc-1@svcacct.local"},
				{"Name": usermanagement.CognitoEnvIDAttrName, "Value": "acmehealth"},
			},
		}
		return jsonResp(200, resp), nil
	})

	c := createEnvCheckEchoContext("sub-svc")
	c.Set("auth_source", AuthSourceService)

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	assert.NoError(t, err, "matching-env service token should pass")
	assert.Equal(t, "us-east-1_PoolB", seenPoolID,
		"service auth_source must look up against pool B, not pool A")
}

// TestCheckUserEnvironmentAccess_ServiceAuthSource_PoolBUnconfigured
// verifies that a service-tagged request fails closed (403) when
// COGNITO_SVC_USER_POOL_ID is empty — we never silently fall back to
// looking up the sub in pool A, because the sub does not exist there.
func TestCheckUserEnvironmentAccess_ServiceAuthSource_PoolBUnconfigured(t *testing.T) {
	ResetEnvCheckCache()

	config := &envCheckMockConfig{
		mockConfigProvider: mockConfigProvider{
			region:     "us-east-1",
			userPoolID: "us-east-1_PoolA",
			// svcUserPoolID intentionally empty
		},
		envID: "acmehealth",
	}

	cognitoCalled := false
	cognitoClient := newTestCognitoClientForEnvCheck(t, func(req *http.Request) (*http.Response, error) {
		cognitoCalled = true
		return jsonResp(200, map[string]interface{}{}), nil
	})

	c := createEnvCheckEchoContext("sub-svc")
	c.Set("auth_source", AuthSourceService)

	err := CheckUserEnvironmentAccess(c, config, cognitoClient)
	assert.Error(t, err, "service token without pool B configured must be rejected")
	assert.False(t, cognitoCalled, "must not call Cognito when pool B is unset")
}