easyai-ai-gateway/apps/api/internal/oidcsession/login_transaction.go

package oidcsession

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"net/url"
	"strings"
	"time"

	"github.com/google/uuid"
)

const LoginTransactionCookieName = "easyai_gateway_oidc_login"

var loginTransactionAAD = []byte("easyai-gateway/oidc-login-transaction/v1")

type LoginTransaction struct {
	State        string    `json:"state"`
	Nonce        string    `json:"nonce"`
	PKCEVerifier string    `json:"pkceVerifier"`
	ReturnTo     string    `json:"returnTo"`
	ContextType  string    `json:"contextType"`
	TenantHint   string    `json:"tenantHint,omitempty"`
	CreatedAt    time.Time `json:"createdAt"`
}

func NewLoginTransaction(returnTo string, now time.Time) (LoginTransaction, error) {
	return NewLoginTransactionWithContext(returnTo, "tenant", "", now)
}

func NewLoginTransactionWithTenantHint(returnTo, tenantHint string, now time.Time) (LoginTransaction, error) {
	return NewLoginTransactionWithContext(returnTo, "tenant", tenantHint, now)
}

func NewLoginTransactionWithContext(
	returnTo, contextType, tenantHint string,
	now time.Time,
) (LoginTransaction, error) {
	if !ValidReturnTo(returnTo) {
		return LoginTransaction{}, errors.New("returnTo must be a same-origin relative path")
	}
	contextType = strings.TrimSpace(contextType)
	tenantHint = strings.TrimSpace(tenantHint)
	if contextType != "platform" && contextType != "tenant" {
		return LoginTransaction{}, errors.New("contextType must be platform or tenant")
	}
	if contextType == "platform" && tenantHint != "" {
		return LoginTransaction{}, errors.New("platform context cannot bind a tenant hint")
	}
	if tenantHint != "" && uuid.Validate(tenantHint) != nil {
		return LoginTransaction{}, errors.New("tenantHint must be a UUID")
	}
	state, err := randomBase64URL(32)
	if err != nil {
		return LoginTransaction{}, err
	}
	nonce, err := randomBase64URL(32)
	if err != nil {
		return LoginTransaction{}, err
	}
	verifier, err := randomBase64URL(32)
	if err != nil {
		return LoginTransaction{}, err
	}
	return LoginTransaction{
		State: state, Nonce: nonce, PKCEVerifier: verifier,
		ReturnTo: returnTo, ContextType: contextType,
		TenantHint: tenantHint, CreatedAt: now.UTC(),
	}, nil
}

func (c *Cipher) EncodeLoginTransaction(transaction LoginTransaction) (string, error) {
	payload, err := c.SealJSON(transaction, loginTransactionAAD)
	if err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(payload), nil
}

func (c *Cipher) DecodeLoginTransaction(encoded string, now time.Time) (LoginTransaction, error) {
	payload, err := base64.RawURLEncoding.DecodeString(strings.TrimSpace(encoded))
	if err != nil {
		return LoginTransaction{}, errors.New("OIDC login transaction is invalid")
	}
	var transaction LoginTransaction
	if err := c.OpenJSON(payload, loginTransactionAAD, &transaction); err != nil {
		return LoginTransaction{}, err
	}
	if transaction.State == "" || transaction.Nonce == "" || transaction.PKCEVerifier == "" || !ValidReturnTo(transaction.ReturnTo) ||
		transaction.ContextType != "platform" && transaction.ContextType != "tenant" ||
		transaction.ContextType == "platform" && transaction.TenantHint != "" ||
		transaction.TenantHint != "" && uuid.Validate(transaction.TenantHint) != nil ||
		transaction.CreatedAt.IsZero() || now.Before(transaction.CreatedAt.Add(-time.Minute)) || !now.Before(transaction.CreatedAt.Add(10*time.Minute)) {
		return LoginTransaction{}, errors.New("OIDC login transaction has expired or is invalid")
	}
	return transaction, nil
}

func ValidReturnTo(value string) bool {
	value = strings.TrimSpace(value)
	if value == "" || !strings.HasPrefix(value, "/") || strings.HasPrefix(value, "//") || strings.Contains(value, "\\") {
		return false
	}
	parsed, err := url.Parse(value)
	return err == nil && !parsed.IsAbs() && parsed.Host == ""
}

func randomBase64URL(size int) (string, error) {
	value := make([]byte, size)
	if _, err := rand.Read(value); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(value), nil
}