From 7eb587f956f02c8d8b29c27dd17a79e7bda67ff8 Mon Sep 17 00:00:00 2001
From: chrislu <chris.lu@gmail.com>
Date: Tue, 26 Aug 2025 22:28:03 -0700
Subject: [PATCH] add MockTrustPolicyValidator

---
 weed/iam/sts/cross_instance_token_test.go | 12 +++++++++
 weed/iam/sts/sts_service_test.go          |  4 +++
 weed/iam/sts/test_utils.go                | 33 +++++++++++++++++++++++
 3 files changed, 49 insertions(+)
 create mode 100644 weed/iam/sts/test_utils.go

diff --git a/weed/iam/sts/cross_instance_token_test.go b/weed/iam/sts/cross_instance_token_test.go
index 713ba9ec1..c9359facf 100644
--- a/weed/iam/sts/cross_instance_token_test.go
+++ b/weed/iam/sts/cross_instance_token_test.go
@@ -59,6 +59,12 @@ func TestCrossInstanceTokenUsage(t *testing.T) {
 	err = instanceC.Initialize(sharedConfig)
 	require.NoError(t, err, "Instance C should initialize")
 
+	// Set up mock trust policy validator for all instances (required for STS testing)
+	mockValidator := &MockTrustPolicyValidator{}
+	instanceA.SetTrustPolicyValidator(mockValidator)
+	instanceB.SetTrustPolicyValidator(mockValidator)
+	instanceC.SetTrustPolicyValidator(mockValidator)
+
 	// Test 1: Token generated on Instance A can be validated on Instance B & C
 	t.Run("cross_instance_token_validation", func(t *testing.T) {
 		// Generate session token on Instance A
@@ -368,6 +374,12 @@ func TestSTSRealWorldDistributedScenarios(t *testing.T) {
 		err = gateway3.Initialize(productionConfig)
 		require.NoError(t, err)
 
+		// Set up mock trust policy validator for all gateway instances
+		mockValidator := &MockTrustPolicyValidator{}
+		gateway1.SetTrustPolicyValidator(mockValidator)
+		gateway2.SetTrustPolicyValidator(mockValidator)
+		gateway3.SetTrustPolicyValidator(mockValidator)
+
 		// Step 1: User authenticates and hits Gateway 1 for AssumeRole
 		assumeRequest := &AssumeRoleWithWebIdentityRequest{
 			RoleArn:          "arn:seaweed:iam::role/ProductionS3User",
diff --git a/weed/iam/sts/sts_service_test.go b/weed/iam/sts/sts_service_test.go
index a45536a16..7c1eee98e 100644
--- a/weed/iam/sts/sts_service_test.go
+++ b/weed/iam/sts/sts_service_test.go
@@ -312,6 +312,10 @@ func setupTestSTSService(t *testing.T) *STSService {
 	err := service.Initialize(config)
 	require.NoError(t, err)
 
+	// Set up mock trust policy validator (required for STS testing)
+	mockValidator := &MockTrustPolicyValidator{}
+	service.SetTrustPolicyValidator(mockValidator)
+
 	// Register test providers
 	mockOIDCProvider := &MockIdentityProvider{
 		name: "test-oidc",
diff --git a/weed/iam/sts/test_utils.go b/weed/iam/sts/test_utils.go
new file mode 100644
index 000000000..8ea38573f
--- /dev/null
+++ b/weed/iam/sts/test_utils.go
@@ -0,0 +1,33 @@
+package sts
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/seaweedfs/seaweedfs/weed/iam/providers"
+)
+
+// MockTrustPolicyValidator is a simple mock for testing STS functionality
+type MockTrustPolicyValidator struct{}
+
+// ValidateTrustPolicyForWebIdentity allows valid test tokens for STS testing
+func (m *MockTrustPolicyValidator) ValidateTrustPolicyForWebIdentity(ctx context.Context, roleArn string, webIdentityToken string) error {
+	// For STS unit tests, allow valid test tokens
+	if webIdentityToken == "valid_test_token" || webIdentityToken == "valid-oidc-token" {
+		return nil
+	}
+	// Reject invalid tokens
+	if webIdentityToken == "invalid_token" || webIdentityToken == "expired_token" || webIdentityToken == "invalid-token" {
+		return fmt.Errorf("trust policy denies token")
+	}
+	return nil
+}
+
+// ValidateTrustPolicyForCredentials allows valid test identities for STS testing
+func (m *MockTrustPolicyValidator) ValidateTrustPolicyForCredentials(ctx context.Context, roleArn string, identity *providers.ExternalIdentity) error {
+	// For STS unit tests, allow test identities
+	if identity != nil && identity.UserID != "" {
+		return nil
+	}
+	return fmt.Errorf("invalid identity for role assumption")
+}