iris/auth/oidc.go

//go:build js && wasm

package auth

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/url"
	"strings"
	"syscall/js"
	"time"
)

// OIDC types are defined in types.go

// OIDCClient handles OIDC authentication flows
type OIDCClient struct {
	config      *OIDCConfig
	clientID    string
	redirectURI string
	scopes      []string

	// Browser storage
	localStorage   js.Value
	sessionStorage js.Value

	// HTTP client for WASM
	httpClient *WASMHTTPClient
}

// Token and user info types are defined in types.go

// NewOIDCClient creates a new OIDC client
func NewOIDCClient(issuer, clientID, redirectURI string) *OIDCClient {
	client := &OIDCClient{
		clientID:       clientID,
		redirectURI:    redirectURI,
		scopes:         []string{"openid", "email", "profile"},
		localStorage:   js.Global().Get("localStorage"),
		sessionStorage: js.Global().Get("sessionStorage"),
		httpClient:     &WASMHTTPClient{},
	}

	return client
}

// DiscoverConfig fetches the OIDC discovery configuration synchronously
func (c *OIDCClient) DiscoverConfig(issuer string) error {
	discoveryURL := strings.TrimSuffix(issuer, "/") + "/.well-known/openid-configuration"

	// Use fetch to get the configuration
	return c.fetchJSON(discoveryURL, &c.config)
}

// DiscoverConfigAsync fetches the OIDC discovery configuration asynchronously
func (c *OIDCClient) DiscoverConfigAsync(issuer string, callback func(error)) {
	discoveryURL := strings.TrimSuffix(issuer, "/") + "/.well-known/openid-configuration"

	c.httpClient.FetchJSONAsync(discoveryURL, func(result HTTPResult) {
		if result.Error != nil {
			callback(result.Error)
			return
		}

		// Parse the configuration
		var config OIDCConfig
		if err := json.Unmarshal(result.Data, &config); err != nil {
			callback(fmt.Errorf("failed to parse OIDC config: %w", err))
			return
		}

		c.config = &config
		callback(nil)
	})
}

// EnsureConfigLoaded ensures OIDC configuration is loaded before proceeding
func (c *OIDCClient) EnsureConfigLoaded(issuer string) error {
	if c.config != nil {
		return nil
	}
	return c.DiscoverConfig(issuer)
}

// StartAuthFlow initiates the OIDC authentication flow
func (c *OIDCClient) StartAuthFlow() error {
	if c.config == nil {
		return fmt.Errorf("OIDC configuration not loaded - call DiscoverConfig first")
	}

	// Generate PKCE parameters
	verifier, err := c.generateCodeVerifier()
	if err != nil {
		return fmt.Errorf("failed to generate code verifier: %w", err)
	}

	challenge := c.generateCodeChallenge(verifier)
	state := c.generateState()

	// Store PKCE verifier and state in session storage
	c.sessionStorage.Call("setItem", "pkce_verifier", verifier)
	c.sessionStorage.Call("setItem", "auth_state", state)

	// Build authorization URL
	authURL, err := c.buildAuthURL(challenge, state)
	if err != nil {
		return fmt.Errorf("failed to build auth URL: %w", err)
	}

	// Redirect to Dex
	js.Global().Get("window").Get("location").Set("href", authURL)
	return nil
}

// HandleCallback processes the OAuth callback
func (c *OIDCClient) HandleCallback(issuer string) (*TokenResponse, error) {
	// Ensure configuration is loaded before processing
	if err := c.EnsureConfigLoaded(issuer); err != nil {
		return nil, fmt.Errorf("failed to load OIDC configuration: %w", err)
	}

	// Get current URL parameters
	urlParams := c.getURLParams()

	code := urlParams.Get("code")
	state := urlParams.Get("state")
	errorParam := urlParams.Get("error")

	if errorParam != "" {
		return nil, fmt.Errorf("OAuth error: %s", errorParam)
	}

	if code == "" {
		return nil, fmt.Errorf("authorization code not found in callback")
	}

	// Verify state parameter
	storedState := c.sessionStorage.Call("getItem", "auth_state").String()
	if state != storedState {
		return nil, fmt.Errorf("invalid state parameter")
	}

	// Get PKCE verifier
	verifier := c.sessionStorage.Call("getItem", "pkce_verifier").String()
	if verifier == "" {
		return nil, fmt.Errorf("PKCE verifier not found")
	}

	// Exchange code for tokens
	return c.exchangeCodeForTokens(code, verifier)
}

// GetStoredTokens retrieves tokens from local storage
func (c *OIDCClient) GetStoredTokens() *TokenResponse {
	accessToken := c.localStorage.Call("getItem", "access_token")
	idToken := c.localStorage.Call("getItem", "id_token")
	refreshToken := c.localStorage.Call("getItem", "refresh_token")

	if accessToken.IsNull() || idToken.IsNull() {
		return nil
	}

	return &TokenResponse{
		AccessToken:  accessToken.String(),
		IDToken:      idToken.String(),
		RefreshToken: refreshToken.String(),
		TokenType:    "Bearer",
	}
}

// StoreTokens saves tokens to local storage
func (c *OIDCClient) StoreTokens(tokens *TokenResponse) {
	c.localStorage.Call("setItem", "access_token", tokens.AccessToken)
	c.localStorage.Call("setItem", "id_token", tokens.IDToken)
	if tokens.RefreshToken != "" {
		c.localStorage.Call("setItem", "refresh_token", tokens.RefreshToken)
	}

	// Store expiration time
	if tokens.ExpiresIn > 0 {
		expiresAt := time.Now().Add(time.Duration(tokens.ExpiresIn) * time.Second)
		c.localStorage.Call("setItem", "token_expires_at", expiresAt.Unix())
	}
}

// ClearTokens removes all stored tokens
func (c *OIDCClient) ClearTokens() {
	c.localStorage.Call("removeItem", "access_token")
	c.localStorage.Call("removeItem", "id_token")
	c.localStorage.Call("removeItem", "refresh_token")
	c.localStorage.Call("removeItem", "token_expires_at")
}

// IsAuthenticated checks if the user has valid tokens
func (c *OIDCClient) IsAuthenticated() bool {
	tokens := c.GetStoredTokens()
	if tokens == nil {
		return false
	}

	// Check if token is expired
	expiresAtStr := c.localStorage.Call("getItem", "token_expires_at")
	if !expiresAtStr.IsNull() {
		// localStorage returns strings, need to convert to int
		expiresAtString := expiresAtStr.String()
		var expiresAt int64
		if _, err := fmt.Sscanf(expiresAtString, "%d", &expiresAt); err == nil {
			if time.Now().Unix() > expiresAt {
				return false
			}
		}
	}

	return tokens.AccessToken != "" && tokens.IDToken != ""
}

// Logout clears all authentication data
func (c *OIDCClient) Logout() {
	c.ClearTokens()
	c.sessionStorage.Call("removeItem", "pkce_verifier")
	c.sessionStorage.Call("removeItem", "auth_state")
}

// GetAuthHeader returns the Authorization header value
func (c *OIDCClient) GetAuthHeader() string {
	tokens := c.GetStoredTokens()
	if tokens == nil {
		return ""
	}
	return "Bearer " + tokens.IDToken
}

// Helper methods

func (c *OIDCClient) generateCodeVerifier() (string, error) {
	bytes := make([]byte, 32)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(bytes), nil
}

func (c *OIDCClient) generateCodeChallenge(verifier string) string {
	hash := sha256.Sum256([]byte(verifier))
	return base64.RawURLEncoding.EncodeToString(hash[:])
}

func (c *OIDCClient) generateState() string {
	bytes := make([]byte, 16)
	rand.Read(bytes)
	return base64.RawURLEncoding.EncodeToString(bytes)
}

func (c *OIDCClient) buildAuthURL(challenge, state string) (string, error) {
	params := url.Values{
		"client_id":             {c.clientID},
		"response_type":         {"code"},
		"scope":                 {strings.Join(c.scopes, " ")},
		"redirect_uri":          {c.redirectURI},
		"code_challenge":        {challenge},
		"code_challenge_method": {"S256"},
		"state":                 {state},
	}

	authURL := c.config.AuthURL + "?" + params.Encode()
	return authURL, nil
}

func (c *OIDCClient) getURLParams() url.Values {
	location := js.Global().Get("window").Get("location")
	search := location.Get("search").String()
	params, _ := url.ParseQuery(strings.TrimPrefix(search, "?"))
	return params
}

func (c *OIDCClient) exchangeCodeForTokens(code, verifier string) (*TokenResponse, error) {
	// Prepare form data
	data := url.Values{
		"grant_type":    {"authorization_code"},
		"client_id":     {c.clientID},
		"code":          {code},
		"redirect_uri":  {c.redirectURI},
		"code_verifier": {verifier},
	}

	// Make POST request to token endpoint
	var tokens TokenResponse
	if err := c.postForm(c.config.TokenURL, data, &tokens); err != nil {
		return nil, fmt.Errorf("token exchange failed: %w", err)
	}

	return &tokens, nil
}

// fetchJSON performs a GET request and unmarshals JSON response
func (c *OIDCClient) fetchJSON(url string, dest interface{}) error {
	return c.httpClient.FetchJSON(url, dest)
}

// postForm performs a POST request with form data
func (c *OIDCClient) postForm(url string, data url.Values, dest interface{}) error {
	return c.httpClient.PostForm(url, data, dest)
}