baron-sso/backend/internal/handler/auth_handler_dynamic_claims_test.go

package handler

import (
	"baron-sso-backend/internal/service"
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gofiber/fiber/v2"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

func TestBuildOidcClaimsFromTraits_DynamicClaims(t *testing.T) {
	traits := map[string]interface{}{
		"email":     "user@baron.com",
		"name":      "홍길동",
		"tenant_id": "primary-tenant-999", // Added primary tenant
		"tenant-1": map[string]interface{}{
			"department": "개발팀",
			"grade":      "선임",
		},
		"tenant-2": map[string]interface{}{
			"department": "재무팀",
			"grade":      "팀장",
		},
	}
	scopes := []string{"openid", "profile"}

	t.Run("No tenantID", func(t *testing.T) {
		claims := buildOidcClaimsFromTraits(traits, scopes, "")
		assert.Equal(t, "user@baron.com", claims["email"])
		assert.Equal(t, "홍길동", claims["name"])
		assert.Nil(t, claims["department"])
		assert.Nil(t, claims["grade"])

		assert.NotNil(t, claims["tenants"])
		assert.Contains(t, claims["joined_tenants"], "tenant-1")
		assert.Contains(t, claims["joined_tenants"], "tenant-2")
		assert.Contains(t, claims["joined_tenants"], "primary-tenant-999") // Should contain primary
	})

	t.Run("With tenant-1", func(t *testing.T) {
		claims := buildOidcClaimsFromTraits(traits, scopes, "tenant-1")
		assert.Equal(t, "user@baron.com", claims["email"])
		assert.Equal(t, "홍길동", claims["name"])
		assert.Equal(t, "tenant-1", claims["tenant_id"]) // Dynamic tenant injection overwrites top-level for this context
		assert.Equal(t, "개발팀", claims["department"])
		assert.Equal(t, "선임", claims["grade"])

		assert.NotNil(t, claims["tenants"])
		assert.Contains(t, claims["joined_tenants"], "tenant-1")
		assert.Contains(t, claims["joined_tenants"], "tenant-2")
		assert.Contains(t, claims["joined_tenants"], "primary-tenant-999")
	})

	t.Run("With tenant-2", func(t *testing.T) {
		claims := buildOidcClaimsFromTraits(traits, scopes, "tenant-2")
		assert.Equal(t, "user@baron.com", claims["email"])
		assert.Equal(t, "홍길동", claims["name"])
		assert.Equal(t, "tenant-2", claims["tenant_id"])
		assert.Equal(t, "재무팀", claims["department"])
		assert.Equal(t, "팀장", claims["grade"])

		assert.NotNil(t, claims["tenants"])
		assert.Contains(t, claims["joined_tenants"], "primary-tenant-999")
	})

	t.Run("With non-existent tenant", func(t *testing.T) {
		claims := buildOidcClaimsFromTraits(traits, scopes, "tenant-3")
		assert.Equal(t, "user@baron.com", claims["email"])
		assert.Equal(t, "홍길동", claims["name"])
		assert.Equal(t, "tenant-3", claims["tenant_id"])
		assert.Nil(t, claims["department"])
		assert.Nil(t, claims["grade"])

		assert.NotNil(t, claims["tenants"])
		assert.Contains(t, claims["joined_tenants"], "tenant-1")
		assert.Contains(t, claims["joined_tenants"], "primary-tenant-999")
	})
}

func TestAcceptConsentRequest_DynamicClaims(t *testing.T) {
	var capturedClaims map[string]interface{}

	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		// Hydra: Get Consent Request
		if r.URL.Path == "/oauth2/auth/requests/consent" && r.URL.Query().Get("consent_challenge") == "challenge-dynamic" {
			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"challenge":       "challenge-dynamic",
				"requested_scope": []string{"openid", "profile"},
				"subject":         "user-123",
				"client": map[string]interface{}{
					"client_id": "client-app",
					"metadata": map[string]interface{}{
						"tenant_id": "tenant-abc",
					},
				},
			}), nil
		}
		// Kratos: Get Identity
		if r.URL.Path == "/admin/identities/user-123" {
			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"id": "user-123",
				"traits": map[string]interface{}{
					"email": "user@test.com",
					"name":  "Test User",
					"tenant-abc": map[string]interface{}{
						"department": "Innovation",
						"position":   "Architect",
					},
				},
			}), nil
		}
		// Hydra: Accept Consent Request
		if r.URL.Path == "/oauth2/auth/requests/consent/accept" && r.URL.Query().Get("consent_challenge") == "challenge-dynamic" {
			// Capture the claims sent to Hydra
			body, _ := io.ReadAll(r.Body)
			var acceptReq map[string]interface{}
			json.Unmarshal(body, &acceptReq)
			if session, ok := acceptReq["session"].(map[string]interface{}); ok {
				capturedClaims = session["id_token"].(map[string]interface{})
			}

			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"redirect_to": "http://rp/cb",
			}), nil
		}
		return httpResponse(r, http.StatusNotFound, "not found"), nil
	})

	client := &http.Client{Transport: transport}
	origDefault := http.DefaultClient
	http.DefaultClient = client
	defer func() { http.DefaultClient = origDefault }()

	h := &AuthHandler{
		Hydra: &service.HydraAdminService{
			AdminURL:   "http://hydra.test",
			HTTPClient: client,
		},
		KratosAdmin: new(MockKratosAdminService),
	}
	h.KratosAdmin.(*MockKratosAdminService).On("GetIdentity", mock.Anything, "user-123").Return(&service.KratosIdentity{
		ID: "user-123",
		Traits: map[string]interface{}{
			"email": "user@test.com",
			"name":  "Test User",
			"tenant-abc": map[string]interface{}{
				"department": "Innovation",
				"position":   "Architect",
			},
		},
	}, nil)

	app := fiber.New()
	app.Post("/api/v1/auth/consent/accept", h.AcceptConsentRequest)

	reqBody, _ := json.Marshal(map[string]interface{}{
		"consent_challenge": "challenge-dynamic",
		"grant_scope":       []string{"openid", "profile"},
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/consent/accept", bytes.NewReader(reqBody))
	req.Header.Set("Content-Type", "application/json")

	resp, err := app.Test(req)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusOK, resp.StatusCode)

	// Verify captured claims
	assert.NotNil(t, capturedClaims)
	assert.Equal(t, "user@test.com", capturedClaims["email"])
	assert.Equal(t, "tenant-abc", capturedClaims["tenant_id"])
	assert.Equal(t, "Innovation", capturedClaims["department"])
	assert.Equal(t, "Architect", capturedClaims["position"])
}

func TestGetConsentRequest_Skip_DynamicClaims(t *testing.T) {
	var capturedClaims map[string]interface{}

	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		// Hydra: Get Consent Request
		if r.URL.Path == "/oauth2/auth/requests/consent" && r.URL.Query().Get("consent_challenge") == "challenge-skip-dynamic" {
			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"challenge":       "challenge-skip-dynamic",
				"requested_scope": []string{"openid", "profile"},
				"skip":            true,
				"subject":         "user-456",
				"client": map[string]interface{}{
					"client_id": "skip-app",
					"metadata": map[string]interface{}{
						"tenant_id": "tenant-xyz",
					},
				},
			}), nil
		}
		// Kratos: Get Identity
		if r.URL.Path == "/admin/identities/user-456" {
			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"id": "user-456",
				"traits": map[string]interface{}{
					"email": "skip@test.com",
					"tenant-xyz": map[string]interface{}{
						"department": "Security",
						"position":   "Officer",
					},
				},
			}), nil
		}
		// Hydra: Accept Consent Request
		if r.URL.Path == "/oauth2/auth/requests/consent/accept" && r.URL.Query().Get("consent_challenge") == "challenge-skip-dynamic" {
			// Capture the claims sent to Hydra
			body, _ := io.ReadAll(r.Body)
			var acceptReq map[string]interface{}
			json.Unmarshal(body, &acceptReq)
			if session, ok := acceptReq["session"].(map[string]interface{}); ok {
				capturedClaims = session["id_token"].(map[string]interface{})
			}

			return httpJSONAny(r, http.StatusOK, map[string]interface{}{
				"redirect_to": "http://rp/cb",
			}), nil
		}
		return httpResponse(r, http.StatusNotFound, "not found"), nil
	})

	client := &http.Client{Transport: transport}
	origDefault := http.DefaultClient
	http.DefaultClient = client
	defer func() { http.DefaultClient = origDefault }()

	h := &AuthHandler{
		Hydra: &service.HydraAdminService{
			AdminURL:   "http://hydra.test",
			HTTPClient: client,
		},
		KratosAdmin: new(MockKratosAdminService),
	}
	h.KratosAdmin.(*MockKratosAdminService).On("GetIdentity", mock.Anything, "user-456").Return(&service.KratosIdentity{
		ID: "user-456",
		Traits: map[string]interface{}{
			"email": "skip@test.com",
			"tenant-xyz": map[string]interface{}{
				"department": "Security",
				"position":   "Officer",
			},
		},
	}, nil)

	app := fiber.New()
	app.Get("/api/v1/auth/consent", h.GetConsentRequest)

	req := httptest.NewRequest(http.MethodGet, "/api/v1/auth/consent?consent_challenge=challenge-skip-dynamic", nil)
	resp, err := app.Test(req)
	assert.NoError(t, err)
	assert.Equal(t, http.StatusOK, resp.StatusCode)

	// Verify captured claims
	assert.NotNil(t, capturedClaims)
	assert.Equal(t, "skip@test.com", capturedClaims["email"])
	assert.Equal(t, "tenant-xyz", capturedClaims["tenant_id"])
	assert.Equal(t, "Security", capturedClaims["department"])
	assert.Equal(t, "Officer", capturedClaims["position"])
}