seaweedfs/weed/s3api/s3api_sts_get_federation_token_test.go

package s3api

import (
	"context"
	"encoding/xml"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sort"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/seaweedfs/seaweedfs/weed/iam/integration"
	"github.com/seaweedfs/seaweedfs/weed/iam/policy"
	"github.com/seaweedfs/seaweedfs/weed/iam/sts"
	"github.com/seaweedfs/seaweedfs/weed/pb/iam_pb"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// mockUserStore implements integration.UserStore for testing GetPoliciesForUser
type mockUserStore struct {
	users map[string]*iam_pb.Identity
}

func (m *mockUserStore) GetUser(_ context.Context, username string) (*iam_pb.Identity, error) {
	u, ok := m.users[username]
	if !ok {
		return nil, nil
	}
	return u, nil
}

// TestGetFederationToken_BasicFlow tests basic credential generation for GetFederationToken
func TestGetFederationToken_BasicFlow(t *testing.T) {
	stsService, _ := setupTestSTSService(t)

	iam := &IdentityAccessManagement{
		iamIntegration: &MockIAMIntegration{},
	}
	stsHandlers := NewSTSHandlers(stsService, iam)

	// Simulate the core logic of handleGetFederationToken
	name := "BobApp"
	callerIdentity := &Identity{
		Name:         "alice",
		PrincipalArn: fmt.Sprintf("arn:aws:iam::%s:user/alice", defaultAccountID),
		PolicyNames:  []string{"S3ReadPolicy"},
	}

	accountID := stsHandlers.getAccountID()

	// Generate session ID and credentials
	sessionId, err := sts.GenerateSessionId()
	require.NoError(t, err)

	expiration := time.Now().Add(12 * time.Hour)
	federatedUserArn := fmt.Sprintf("arn:aws:sts::%s:federated-user/%s", accountID, name)
	federatedUserId := fmt.Sprintf("%s:%s", accountID, name)

	claims := sts.NewSTSSessionClaims(sessionId, stsService.Config.Issuer, expiration).
		WithSessionName(name).
		WithRoleInfo(callerIdentity.PrincipalArn, federatedUserId, federatedUserArn).
		WithPolicies(callerIdentity.PolicyNames)

	sessionToken, err := stsService.GetTokenGenerator().GenerateJWTWithClaims(claims)
	require.NoError(t, err)

	// Validate the session token
	sessionInfo, err := stsService.ValidateSessionToken(context.Background(), sessionToken)
	require.NoError(t, err)
	require.NotNil(t, sessionInfo)

	// Verify the session info contains caller's policies
	assert.Equal(t, []string{"S3ReadPolicy"}, sessionInfo.Policies)

	// Verify principal is the federated user ARN
	assert.Equal(t, federatedUserArn, sessionInfo.Principal)

	// Verify the RoleArn points to the caller's identity (for policy resolution)
	assert.Equal(t, callerIdentity.PrincipalArn, sessionInfo.RoleArn)

	// Verify session name
	assert.Equal(t, name, sessionInfo.SessionName)
}

// TestGetFederationToken_WithSessionPolicy tests session policy scoping
func TestGetFederationToken_WithSessionPolicy(t *testing.T) {
	stsService, _ := setupTestSTSService(t)

	stsHandlers := NewSTSHandlers(stsService, &IdentityAccessManagement{
		iamIntegration: &MockIAMIntegration{},
	})

	accountID := stsHandlers.getAccountID()
	name := "ScopedApp"

	sessionPolicyJSON := `{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Action":["s3:GetObject"],"Resource":["arn:aws:s3:::my-bucket/*"]}]}`
	normalizedPolicy, err := sts.NormalizeSessionPolicy(sessionPolicyJSON)
	require.NoError(t, err)

	sessionId, err := sts.GenerateSessionId()
	require.NoError(t, err)

	expiration := time.Now().Add(12 * time.Hour)
	federatedUserArn := fmt.Sprintf("arn:aws:sts::%s:federated-user/%s", accountID, name)
	federatedUserId := fmt.Sprintf("%s:%s", accountID, name)

	claims := sts.NewSTSSessionClaims(sessionId, stsService.Config.Issuer, expiration).
		WithSessionName(name).
		WithRoleInfo("arn:aws:iam::000000000000:user/caller", federatedUserId, federatedUserArn).
		WithPolicies([]string{"S3FullAccess"}).
		WithSessionPolicy(normalizedPolicy)

	sessionToken, err := stsService.GetTokenGenerator().GenerateJWTWithClaims(claims)
	require.NoError(t, err)

	sessionInfo, err := stsService.ValidateSessionToken(context.Background(), sessionToken)
	require.NoError(t, err)
	require.NotNil(t, sessionInfo)

	// Verify session policy is embedded
	assert.NotEmpty(t, sessionInfo.SessionPolicy)
	assert.Contains(t, sessionInfo.SessionPolicy, "s3:GetObject")

	// Verify caller's policies are still present
	assert.Equal(t, []string{"S3FullAccess"}, sessionInfo.Policies)
}

// TestGetFederationToken_RejectTemporaryCredentials tests that requests
// authenticated with temporary credentials are rejected.
func TestGetFederationToken_RejectTemporaryCredentials(t *testing.T) {
	stsService, _ := setupTestSTSService(t)
	iam := &IdentityAccessManagement{
		iamIntegration: &MockIAMIntegration{
			validateSessionFunc: func(context.Context, string) (*sts.SessionInfo, error) {
				return &sts.SessionInfo{
					AssumedRoleUser: "role/user",
					Principal:       "arn:aws:sts:::assumed-role/role/user",
					Credentials: &sts.Credentials{
						AccessKeyId:     "STSACCESSKEY",
						SecretAccessKey: "sts-secret-key",
					},
					ExpiresAt: time.Now().Add(time.Hour),
				}, nil
			},
		},
		hashes:       make(map[string]*sync.Pool),
		hashCounters: make(map[string]*int32),
	}
	stsHandlers := NewSTSHandlers(stsService, iam)

	tests := []struct {
		name     string
		setToken func(r *http.Request)
	}{
		{
			name: "SessionTokenInHeader",
			setToken: func(r *http.Request) {
				r.Header.Set("X-Amz-Security-Token", "session-token")
			},
		},
		{
			name: "SessionTokenInQuery",
			setToken: func(r *http.Request) {
				q := r.URL.Query()
				q.Set("X-Amz-Security-Token", "session-token")
				r.URL.RawQuery = q.Encode()
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			form := url.Values{}
			form.Set("Action", "GetFederationToken")
			form.Set("Name", "TestUser")
			form.Set("Version", "2011-06-15")

			req := mustNewRequest(http.MethodPost, "http://sts.amazonaws.com/",
				int64(len(form.Encode())), strings.NewReader(form.Encode()), t)
			req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
			tt.setToken(req)
			require.NoError(t, signRequestV4(req, "STSACCESSKEY", "sts-secret-key"))

			rr := httptest.NewRecorder()
			stsHandlers.HandleSTSRequest(rr, req)

			assert.Equal(t, http.StatusForbidden, rr.Code)
			assert.Contains(t, rr.Body.String(), "AccessDenied")
			assert.Contains(t, rr.Body.String(), "cannot be called with temporary credentials")
		})
	}
}

// A session token attached to a statically configured credential (as credential
// vendors emit) must not disqualify the caller as temporary.
func TestGetFederationToken_VendedStaticCredential(t *testing.T) {
	stsService, _ := setupTestSTSService(t)
	iam := newTestIAMWithCreds()
	stsHandlers := NewSTSHandlers(stsService, iam)

	form := url.Values{}
	form.Set("Action", "GetFederationToken")
	form.Set("Name", "VendedApp")
	form.Set("Version", "2011-06-15")

	req := mustNewRequest(http.MethodPost, "http://sts.amazonaws.com/",
		int64(len(form.Encode())), strings.NewReader(form.Encode()), t)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("X-Amz-Security-Token", "vendor-token")
	require.NoError(t, signRequestV4(req, "access_key_1", "secret_key_1"))

	rr := httptest.NewRecorder()
	stsHandlers.HandleSTSRequest(rr, req)

	assert.Equal(t, http.StatusOK, rr.Code, rr.Body.String())
	assert.Contains(t, rr.Body.String(), "GetFederationTokenResponse")
}

// TestGetFederationToken_MissingName tests that a missing Name parameter returns an error
func TestGetFederationToken_MissingName(t *testing.T) {
	stsService, _ := setupTestSTSService(t)
	stsHandlers := NewSTSHandlers(stsService, &IdentityAccessManagement{
		iamIntegration: &MockIAMIntegration{},
	})

	req := httptest.NewRequest("POST", "/", nil)
	req.Form = url.Values{}
	req.Form.Set("Action", "GetFederationToken")
	req.Form.Set("Version", "2011-06-15")
	// Name is intentionally omitted

	rr := httptest.NewRecorder()
	stsHandlers.HandleSTSRequest(rr, req)

	assert.Equal(t, http.StatusBadRequest, rr.Code)
	assert.Contains(t, rr.Body.String(), "Name is required")
}

// TestGetFederationToken_NameValidation tests Name parameter validation
func TestGetFederationToken_NameValidation(t *testing.T) {
	stsService, _ := setupTestSTSService(t)
	stsHandlers := NewSTSHandlers(stsService, &IdentityAccessManagement{
		iamIntegration: &MockIAMIntegration{},
	})

	tests := []struct {
		name        string
		federName   string
		expectError bool
		errContains string
	}{
		{
			name:        "TooShort",
			federName:   "A",
			expectError: true,
			errContains: "between 2 and 32",
		},
		{
			name:        "TooLong",
			federName:   strings.Repeat("A", 33),
			expectError: true,
			errContains: "between 2 and 32",
		},
		{
			name:        "MinLength",
			federName:   "AB",
			expectError: false,
		},
		{
			name:        "MaxLength",
			federName:   strings.Repeat("A", 32),
			expectError: false,
		},
		{
			name:        "ValidSpecialChars",
			federName:   "user+=,.@-test",
			expectError: false,
		},
		{
			name:        "InvalidChars_Space",
			federName:   "bad name",
			expectError: true,
			errContains: "invalid characters",
		},
		{
			name:        "InvalidChars_Slash",
			federName:   "bad/name",
			expectError: true,
			errContains: "invalid characters",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest("POST", "/", nil)
			req.Form = url.Values{}
			req.Form.Set("Action", "GetFederationToken")
			req.Form.Set("Name", tt.federName)
			req.Form.Set("Version", "2011-06-15")

			rr := httptest.NewRecorder()
			stsHandlers.HandleSTSRequest(rr, req)

			if tt.expectError {
				assert.Equal(t, http.StatusBadRequest, rr.Code)
				assert.Contains(t, rr.Body.String(), tt.errContains)
			} else {
				// Valid name should proceed past validation — will fail at SigV4
				// (returns 403 because we have no real signature)
				assert.NotEqual(t, http.StatusBadRequest, rr.Code,
					"Valid name should not produce a 400 for name validation")
			}
		})
	}
}

// TestGetFederationToken_DurationValidation tests DurationSeconds validation
func TestGetFederationToken_DurationValidation(t *testing.T) {
	stsService, _ := setupTestSTSService(t)
	stsHandlers := NewSTSHandlers(stsService, &IdentityAccessManagement{
		iamIntegration: &MockIAMIntegration{},
	})

	tests := []struct {
		name        string
		duration    string
		expectError bool
		errContains string
	}{
		{
			name:        "BelowMinimum",
			duration:    "899",
			expectError: true,
			errContains: "between",
		},
		{
			name:        "AboveMaximum",
			duration:    "129601",
			expectError: true,
			errContains: "between",
		},
		{
			name:        "InvalidFormat",
			duration:    "not-a-number",
			expectError: true,
			errContains: "invalid DurationSeconds",
		},
		{
			name:        "MinimumValid",
			duration:    "900",
			expectError: false,
		},
		{
			name:        "MaximumValid_36Hours",
			duration:    "129600",
			expectError: false,
		},
		{
			name:        "Default12Hours",
			duration:    "43200",
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest("POST", "/", nil)
			req.Form = url.Values{}
			req.Form.Set("Action", "GetFederationToken")
			req.Form.Set("Name", "TestUser")
			req.Form.Set("DurationSeconds", tt.duration)
			req.Form.Set("Version", "2011-06-15")

			rr := httptest.NewRecorder()
			stsHandlers.HandleSTSRequest(rr, req)

			if tt.expectError {
				assert.Equal(t, http.StatusBadRequest, rr.Code)
				assert.Contains(t, rr.Body.String(), tt.errContains)
			} else {
				// Valid duration should proceed past validation — will fail at SigV4
				assert.NotEqual(t, http.StatusBadRequest, rr.Code,
					"Valid duration should not produce a 400 for duration validation")
			}
		})
	}
}

// TestGetFederationToken_ResponseFormat tests the XML response structure
func TestGetFederationToken_ResponseFormat(t *testing.T) {
	// Verify the response XML structure matches AWS format
	response := GetFederationTokenResponse{
		Result: GetFederationTokenResult{
			Credentials: STSCredentials{
				AccessKeyId:     "ASIA1234567890",
				SecretAccessKey: "secret123",
				SessionToken:    "token123",
				Expiration:      "2026-04-02T12:00:00Z",
			},
			FederatedUser: FederatedUser{
				FederatedUserId: "000000000000:BobApp",
				Arn:             "arn:aws:sts::000000000000:federated-user/BobApp",
			},
		},
	}
	response.ResponseMetadata.RequestId = "test-request-id"

	data, err := xml.MarshalIndent(response, "", "  ")
	require.NoError(t, err)

	xmlStr := string(data)
	assert.Contains(t, xmlStr, "GetFederationTokenResponse")
	assert.Contains(t, xmlStr, "GetFederationTokenResult")
	assert.Contains(t, xmlStr, "FederatedUser")
	assert.Contains(t, xmlStr, "FederatedUserId")
	assert.Contains(t, xmlStr, "federated-user/BobApp")
	assert.Contains(t, xmlStr, "ASIA1234567890")
	assert.Contains(t, xmlStr, "test-request-id")

	// Verify it can be unmarshaled back
	var parsed GetFederationTokenResponse
	err = xml.Unmarshal(data, &parsed)
	require.NoError(t, err)
	assert.Equal(t, "ASIA1234567890", parsed.Result.Credentials.AccessKeyId)
	assert.Equal(t, "arn:aws:sts::000000000000:federated-user/BobApp", parsed.Result.FederatedUser.Arn)
	assert.Equal(t, "000000000000:BobApp", parsed.Result.FederatedUser.FederatedUserId)
}

// TestGetFederationToken_PolicyEmbedding tests that the caller's policies are embedded
// into the session token using the IAM integration manager
func TestGetFederationToken_PolicyEmbedding(t *testing.T) {
	ctx := context.Background()
	manager := newTestSTSIntegrationManager(t)

	// Create a policy that the user has attached
	userPolicy := &policy.PolicyDocument{
		Version: "2012-10-17",
		Statement: []policy.Statement{
			{
				Effect:   "Allow",
				Action:   []string{"s3:GetObject", "s3:PutObject"},
				Resource: []string{"arn:aws:s3:::user-bucket/*"},
			},
		},
	}
	require.NoError(t, manager.CreatePolicy(ctx, "", "UserS3Policy", userPolicy))

	stsService := manager.GetSTSService()

	// Simulate what handleGetFederationToken does for policy embedding
	name := "AppClient"
	callerPolicies := []string{"UserS3Policy"}

	sessionId, err := sts.GenerateSessionId()
	require.NoError(t, err)

	expiration := time.Now().Add(12 * time.Hour)
	accountID := defaultAccountID
	federatedUserArn := fmt.Sprintf("arn:aws:sts::%s:federated-user/%s", accountID, name)
	federatedUserId := fmt.Sprintf("%s:%s", accountID, name)

	claims := sts.NewSTSSessionClaims(sessionId, stsService.Config.Issuer, expiration).
		WithSessionName(name).
		WithRoleInfo("arn:aws:iam::000000000000:user/caller", federatedUserId, federatedUserArn).
		WithPolicies(callerPolicies)

	sessionToken, err := stsService.GetTokenGenerator().GenerateJWTWithClaims(claims)
	require.NoError(t, err)

	sessionInfo, err := stsService.ValidateSessionToken(ctx, sessionToken)
	require.NoError(t, err)
	require.NotNil(t, sessionInfo)

	// Verify the caller's policy names are embedded
	assert.Equal(t, []string{"UserS3Policy"}, sessionInfo.Policies)
}

// TestGetFederationToken_PolicyIntersection tests that both the caller's base policies
// and the restrictive session policy are embedded in the token, enabling the
// authorization layer to compute their intersection at request time.
func TestGetFederationToken_PolicyIntersection(t *testing.T) {
	ctx := context.Background()
	manager := newTestSTSIntegrationManager(t)

	// Create a broad policy for the caller
	broadPolicy := &policy.PolicyDocument{
		Version: "2012-10-17",
		Statement: []policy.Statement{
			{
				Effect:   "Allow",
				Action:   []string{"s3:*"},
				Resource: []string{"arn:aws:s3:::*", "arn:aws:s3:::*/*"},
			},
		},
	}
	require.NoError(t, manager.CreatePolicy(ctx, "", "S3FullAccess", broadPolicy))

	stsService := manager.GetSTSService()

	// Session policy restricts to one bucket and one action
	sessionPolicyJSON := `{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Action":["s3:GetObject"],"Resource":["arn:aws:s3:::restricted-bucket/*"]}]}`
	normalizedPolicy, err := sts.NormalizeSessionPolicy(sessionPolicyJSON)
	require.NoError(t, err)

	sessionId, err := sts.GenerateSessionId()
	require.NoError(t, err)

	expiration := time.Now().Add(12 * time.Hour)
	name := "RestrictedApp"
	accountID := defaultAccountID
	federatedUserArn := fmt.Sprintf("arn:aws:sts::%s:federated-user/%s", accountID, name)
	federatedUserId := fmt.Sprintf("%s:%s", accountID, name)

	claims := sts.NewSTSSessionClaims(sessionId, stsService.Config.Issuer, expiration).
		WithSessionName(name).
		WithRoleInfo("arn:aws:iam::000000000000:user/caller", federatedUserId, federatedUserArn).
		WithPolicies([]string{"S3FullAccess"}).
		WithSessionPolicy(normalizedPolicy)

	sessionToken, err := stsService.GetTokenGenerator().GenerateJWTWithClaims(claims)
	require.NoError(t, err)

	sessionInfo, err := stsService.ValidateSessionToken(ctx, sessionToken)
	require.NoError(t, err)
	require.NotNil(t, sessionInfo)

	// Verify both the broad base policies and the restrictive session policy are embedded
	// The authorization layer computes intersection at request time
	assert.Equal(t, []string{"S3FullAccess"}, sessionInfo.Policies,
		"Caller's base policies should be embedded in token")
	assert.Contains(t, sessionInfo.SessionPolicy, "restricted-bucket",
		"Session policy should restrict to specific bucket")
	assert.Contains(t, sessionInfo.SessionPolicy, "s3:GetObject",
		"Session policy should restrict to specific action")
}

// TestGetFederationToken_MalformedPolicy tests that invalid policy JSON is rejected
// by the session policy normalization used in the handler
func TestGetFederationToken_MalformedPolicy(t *testing.T) {
	tests := []struct {
		name      string
		policyStr string
		expectErr bool
	}{
		{
			name:      "InvalidJSON",
			policyStr: "not-valid-json",
			expectErr: true,
		},
		{
			name:      "EmptyObject",
			policyStr: "{}",
			expectErr: true,
		},
		{
			name:      "TooLarge",
			policyStr: `{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Action":["s3:GetObject"],"Resource":["` + strings.Repeat("a", 2048) + `"]}]}`,
			expectErr: true,
		},
		{
			name:      "ValidPolicy",
			policyStr: `{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Action":["s3:GetObject"],"Resource":["arn:aws:s3:::bucket/*"]}]}`,
			expectErr: false,
		},
		{
			name:      "EmptyString",
			policyStr: "",
			expectErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := sts.NormalizeSessionPolicy(tt.policyStr)
			if tt.expectErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

// TestGetFederationToken_STSNotReady tests that the handler returns 503 when STS is not initialized
func TestGetFederationToken_STSNotReady(t *testing.T) {
	// Create handlers with nil STS service
	stsHandlers := NewSTSHandlers(nil, &IdentityAccessManagement{
		iamIntegration: &MockIAMIntegration{},
	})

	req := httptest.NewRequest("POST", "/", nil)
	req.Form = url.Values{}
	req.Form.Set("Action", "GetFederationToken")
	req.Form.Set("Name", "TestUser")
	req.Form.Set("Version", "2011-06-15")

	rr := httptest.NewRecorder()
	stsHandlers.HandleSTSRequest(rr, req)

	assert.Equal(t, http.StatusServiceUnavailable, rr.Code)
	assert.Contains(t, rr.Body.String(), "ServiceUnavailable")
}

// TestGetFederationToken_DefaultDuration tests that the default duration is 12 hours
func TestGetFederationToken_DefaultDuration(t *testing.T) {
	assert.Equal(t, int64(43200), defaultFederationDurationSeconds, "Default duration should be 12 hours (43200 seconds)")
	assert.Equal(t, int64(129600), maxFederationDurationSeconds, "Max duration should be 36 hours (129600 seconds)")
}

// TestGetFederationToken_GetPoliciesForUser tests that GetPoliciesForUser
// correctly resolves user policies from the UserStore and returns errors
// when the store is unavailable.
func TestGetFederationToken_GetPoliciesForUser(t *testing.T) {
	ctx := context.Background()
	manager := newTestSTSIntegrationManager(t)

	t.Run("NoUserStore", func(t *testing.T) {
		// UserStore not set — should return error
		policies, err := manager.GetPoliciesForUser(ctx, "alice")
		assert.Error(t, err)
		assert.Nil(t, policies)
		assert.Contains(t, err.Error(), "user store not configured")
	})

	t.Run("UserNotFound", func(t *testing.T) {
		manager.SetUserStore(&mockUserStore{users: map[string]*iam_pb.Identity{}})
		policies, err := manager.GetPoliciesForUser(ctx, "nonexistent")
		assert.NoError(t, err)
		assert.Nil(t, policies)
	})

	t.Run("UserWithPolicies", func(t *testing.T) {
		manager.SetUserStore(&mockUserStore{
			users: map[string]*iam_pb.Identity{
				"alice": {
					Name:        "alice",
					PolicyNames: []string{"GroupReadPolicy", "GroupWritePolicy"},
				},
			},
		})
		policies, err := manager.GetPoliciesForUser(ctx, "alice")
		assert.NoError(t, err)
		assert.Equal(t, []string{"GroupReadPolicy", "GroupWritePolicy"}, policies)
	})

	t.Run("UserWithNoPolicies", func(t *testing.T) {
		manager.SetUserStore(&mockUserStore{
			users: map[string]*iam_pb.Identity{
				"bob": {Name: "bob"},
			},
		})
		policies, err := manager.GetPoliciesForUser(ctx, "bob")
		assert.NoError(t, err)
		assert.Empty(t, policies)
	})
}

// TestGetFederationToken_PolicyMergeAndDedup tests that the handler's policy
// merge logic correctly combines identity.PolicyNames with IAM-manager-resolved
// policies and deduplicates the result.
func TestGetFederationToken_PolicyMergeAndDedup(t *testing.T) {
	ctx := context.Background()
	manager := newTestSTSIntegrationManager(t)

	// Create policies so they exist in the engine
	for _, name := range []string{"DirectPolicy", "GroupPolicy", "SharedPolicy"} {
		require.NoError(t, manager.CreatePolicy(ctx, "", name, &policy.PolicyDocument{
			Version: "2012-10-17",
			Statement: []policy.Statement{
				{Effect: "Allow", Action: []string{"s3:GetObject"}, Resource: []string{"arn:aws:s3:::*/*"}},
			},
		}))
	}

	// Set up a user store that returns group-attached policies
	manager.SetUserStore(&mockUserStore{
		users: map[string]*iam_pb.Identity{
			"alice": {
				Name:        "alice",
				PolicyNames: []string{"GroupPolicy", "SharedPolicy"},
			},
		},
	})

	stsService := manager.GetSTSService()

	// Simulate what the handler does: merge identity.PolicyNames with GetPoliciesForUser
	identityPolicies := []string{"DirectPolicy", "SharedPolicy"} // SharedPolicy overlaps

	policySet := make(map[string]struct{})
	for _, p := range identityPolicies {
		policySet[p] = struct{}{}
	}

	userPolicies, err := manager.GetPoliciesForUser(ctx, "alice")
	require.NoError(t, err)
	for _, p := range userPolicies {
		policySet[p] = struct{}{}
	}

	merged := make([]string, 0, len(policySet))
	for p := range policySet {
		merged = append(merged, p)
	}
	sort.Strings(merged) // deterministic for assertion

	// Should contain all three unique policies, no duplicates
	assert.Equal(t, []string{"DirectPolicy", "GroupPolicy", "SharedPolicy"}, merged)

	// Verify the merged policies can be embedded in a token and recovered
	sessionId, err := sts.GenerateSessionId()
	require.NoError(t, err)

	expiration := time.Now().Add(time.Hour)
	claims := sts.NewSTSSessionClaims(sessionId, stsService.Config.Issuer, expiration).
		WithSessionName("test").
		WithRoleInfo("arn:aws:iam::000000000000:user/alice", "000000000000:test", "arn:aws:sts::000000000000:federated-user/test").
		WithPolicies(merged)

	token, err := stsService.GetTokenGenerator().GenerateJWTWithClaims(claims)
	require.NoError(t, err)

	sessionInfo, err := stsService.ValidateSessionToken(ctx, token)
	require.NoError(t, err)

	sort.Strings(sessionInfo.Policies)
	assert.Equal(t, []string{"DirectPolicy", "GroupPolicy", "SharedPolicy"}, sessionInfo.Policies,
		"Token should contain the deduplicated merge of identity and group policies")
}

// TestGetFederationToken_PolicyMergeNoManager tests that when the IAM manager
// is unavailable, identity.PolicyNames alone are still embedded.
func TestGetFederationToken_PolicyMergeNoManager(t *testing.T) {
	ctx := context.Background()
	stsService, _ := setupTestSTSService(t)

	// No IAM manager — only identity.PolicyNames should be used
	identityPolicies := []string{"UserDirectPolicy"}

	policySet := make(map[string]struct{})
	for _, p := range identityPolicies {
		policySet[p] = struct{}{}
	}

	// IAM manager is nil — skip GetPoliciesForUser (mirrors handler logic)
	var policyManager *integration.IAMManager // nil
	if policyManager != nil {
		t.Fatal("policyManager should be nil in this test")
	}

	merged := make([]string, 0, len(policySet))
	for p := range policySet {
		merged = append(merged, p)
	}

	sessionId, err := sts.GenerateSessionId()
	require.NoError(t, err)

	expiration := time.Now().Add(time.Hour)
	claims := sts.NewSTSSessionClaims(sessionId, stsService.Config.Issuer, expiration).
		WithSessionName("test").
		WithRoleInfo("arn:aws:iam::000000000000:user/alice", "000000000000:test", "arn:aws:sts::000000000000:federated-user/test").
		WithPolicies(merged)

	token, err := stsService.GetTokenGenerator().GenerateJWTWithClaims(claims)
	require.NoError(t, err)

	sessionInfo, err := stsService.ValidateSessionToken(ctx, token)
	require.NoError(t, err)

	assert.Equal(t, []string{"UserDirectPolicy"}, sessionInfo.Policies,
		"Without IAM manager, only identity policies should be embedded")
}