BaronSSO/baron-sso/backend/internal/handler/auth_handler_oidc_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"
)

func newOidcLoginTestApp(h *AuthHandler) *fiber.App {
	app := fiber.New()
	app.Post("/api/v1/auth/oidc/login/accept", h.AcceptOidcLoginRequest)
	return app
}

func TestAcceptOidcLoginRequest_CookieOnly(t *testing.T) {
	var gotSubject string
	var gotChallenge string
	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		switch r.URL.Host {
		case "kratos.test":
			if r.URL.Path != "/sessions/whoami" {
				return httpResponse(r, http.StatusNotFound, "not found"), nil
			}
			if r.Header.Get("X-Session-Token") != "" {
				return httpResponse(r, http.StatusUnauthorized, "invalid token"), nil
			}
			if r.Header.Get("Cookie") == "" {
				return httpResponse(r, http.StatusUnauthorized, "missing cookie"), nil
			}
			return httpJSONAny(r, http.StatusOK, map[string]any{
				"identity": map[string]any{
					"id":     "kratos-123",
					"traits": map[string]any{},
				},
			}), nil
		case "hydra.test":
			if r.URL.Path != "/oauth2/auth/requests/login/accept" {
				return httpResponse(r, http.StatusNotFound, "not found"), nil
			}
			gotChallenge = r.URL.Query().Get("login_challenge")
			body, _ := io.ReadAll(r.Body)
			var payload map[string]any
			_ = json.Unmarshal(body, &payload)
			if subject, ok := payload["subject"].(string); ok {
				gotSubject = subject
			}
			return httpResponse(r, http.StatusOK, `{"redirect_to":"http://rp/cb"}`), nil
		default:
			return httpResponse(r, http.StatusNotFound, "not found"), nil
		}
	})
	client := &http.Client{Transport: transport}
	origDefault := http.DefaultClient
	http.DefaultClient = client
	defer func() {
		http.DefaultClient = origDefault
	}()

	t.Setenv("KRATOS_PUBLIC_URL", "http://kratos.test")

	h := &AuthHandler{
		Hydra: &service.HydraAdminService{
			AdminURL:   "http://hydra.test",
			HTTPClient: client,
		},
	}
	app := newOidcLoginTestApp(h)

	body, _ := json.Marshal(map[string]string{
		"login_challenge": "challenge-123",
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/oidc/login/accept", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Cookie", "ory_kratos_session=abc123")

	resp, err := app.Test(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("expected 200, got %d", resp.StatusCode)
	}
	var got map[string]string
	if err := json.NewDecoder(resp.Body).Decode(&got); err != nil {
		t.Fatalf("failed to decode response: %v", err)
	}
	if got["redirectTo"] != "http://rp/cb" {
		t.Fatalf("unexpected redirectTo: %v", got["redirectTo"])
	}
	if gotSubject != "kratos-123" {
		t.Fatalf("unexpected subject: %v", gotSubject)
	}
	if gotChallenge != "challenge-123" {
		t.Fatalf("unexpected login_challenge: %v", gotChallenge)
	}
}

func TestAcceptOidcLoginRequest_TokenFallbackToCookie(t *testing.T) {
	var gotSubject string
	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		switch r.URL.Host {
		case "kratos.test":
			if r.URL.Path != "/sessions/whoami" {
				return httpResponse(r, http.StatusNotFound, "not found"), nil
			}
			if r.Header.Get("X-Session-Token") != "" {
				return httpResponse(r, http.StatusUnauthorized, "invalid token"), nil
			}
			if r.Header.Get("Cookie") == "" {
				return httpResponse(r, http.StatusUnauthorized, "missing cookie"), nil
			}
			return httpJSONAny(r, http.StatusOK, map[string]any{
				"identity": map[string]any{
					"id":     "kratos-456",
					"traits": map[string]any{},
				},
			}), nil
		case "hydra.test":
			if r.URL.Path != "/oauth2/auth/requests/login/accept" {
				return httpResponse(r, http.StatusNotFound, "not found"), nil
			}
			body, _ := io.ReadAll(r.Body)
			var payload map[string]any
			_ = json.Unmarshal(body, &payload)
			if subject, ok := payload["subject"].(string); ok {
				gotSubject = subject
			}
			return httpResponse(r, http.StatusOK, `{"redirect_to":"http://rp/cb"}`), nil
		default:
			return httpResponse(r, http.StatusNotFound, "not found"), nil
		}
	})
	client := &http.Client{Transport: transport}
	origDefault := http.DefaultClient
	http.DefaultClient = client
	defer func() {
		http.DefaultClient = origDefault
	}()

	t.Setenv("KRATOS_PUBLIC_URL", "http://kratos.test")

	h := &AuthHandler{
		Hydra: &service.HydraAdminService{
			AdminURL:   "http://hydra.test",
			HTTPClient: client,
		},
	}
	app := newOidcLoginTestApp(h)

	body, _ := json.Marshal(map[string]string{
		"login_challenge": "challenge-456",
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/oidc/login/accept", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer invalid-token")
	req.Header.Set("Cookie", "ory_kratos_session=def456")

	resp, err := app.Test(req)
	if err != nil {
		t.Fatalf("request failed: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("expected 200, got %d", resp.StatusCode)
	}
	if gotSubject != "kratos-456" {
		t.Fatalf("unexpected subject: %v", gotSubject)
	}
}