easyai-ai-gateway/apps/api/internal/identity/onboarding_client.go

package identity

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/google/uuid"
)

const maxOnboardingResponseBytes = 1 << 20

type ExchangeStatus string

const (
	ExchangeMetadataPending     ExchangeStatus = "metadata_pending"
	ExchangePreparing           ExchangeStatus = "preparing"
	ExchangeReady               ExchangeStatus = "ready"
	ExchangeCredentialDelivered ExchangeStatus = "credential_delivered"
	ExchangeCompleted           ExchangeStatus = "completed"
	ExchangeFailed              ExchangeStatus = "failed"
	ExchangeExpired             ExchangeStatus = "expired"
)

type ClaimedExchange struct {
	ExchangeID    string    `json:"exchange_id"`
	ExchangeToken string    `json:"exchange_token"`
	ExpiresAt     time.Time `json:"expires_at"`
	Version       int64     `json:"version"`
}

type Exchange struct {
	ExchangeID        string         `json:"exchange_id"`
	ApplicationID     string         `json:"application_id"`
	Status            ExchangeStatus `json:"status"`
	Version           int64          `json:"version"`
	ExpiresAt         time.Time      `json:"expires_at"`
	LastErrorCategory string         `json:"last_error_category,omitempty"`
	AuditID           string         `json:"audit_id,omitempty"`
}

type ManifestClient struct {
	ClientID string `json:"client_id"`
}

type ManifestClients struct {
	BrowserLogin     *ManifestClient `json:"browser_login,omitempty"`
	MachineToMachine *ManifestClient `json:"machine_to_machine,omitempty"`
}

type ManifestSecurityEvents struct {
	TransmitterIssuer     string `json:"transmitter_issuer"`
	ConfigurationEndpoint string `json:"configuration_endpoint"`
	Audience              string `json:"audience"`
}

type ManifestV1 struct {
	SchemaVersion  int                     `json:"schema_version"`
	Issuer         string                  `json:"issuer"`
	TenantID       string                  `json:"tenant_id"`
	ApplicationID  string                  `json:"application_id"`
	Capabilities   []string                `json:"capabilities"`
	Audience       string                  `json:"audience,omitempty"`
	Scopes         []string                `json:"scopes"`
	Clients        ManifestClients         `json:"clients"`
	SecurityEvents *ManifestSecurityEvents `json:"security_events,omitempty"`
}

type MachineCredential struct {
	ClientID     string    `json:"client_id"`
	ClientSecret string    `json:"client_secret"`
	IssuedAt     time.Time `json:"issued_at"`
}

type CredentialDelivery struct {
	Manifest          ManifestV1         `json:"manifest"`
	MachineCredential *MachineCredential `json:"machine_credential,omitempty"`
	Version           int64              `json:"-"`
}

func (manifest ManifestV1) Validate() error {
	if manifest.SchemaVersion != 1 || validatePublicIdentityURL(manifest.Issuer) != nil {
		return errors.New("application manifest identity metadata is invalid")
	}
	if _, err := uuid.Parse(manifest.TenantID); err != nil {
		return errors.New("application manifest tenant is invalid")
	}
	if _, err := uuid.Parse(manifest.ApplicationID); err != nil {
		return errors.New("application manifest application is invalid")
	}
	allowed := map[string]bool{"oidc_login": true, "api_access": true, "machine_to_machine": true, "token_introspection": true, "session_revocation": true}
	capabilities := make(map[string]bool, len(manifest.Capabilities))
	for _, capability := range manifest.Capabilities {
		if !allowed[capability] || capabilities[capability] {
			return errors.New("application manifest capabilities are invalid")
		}
		capabilities[capability] = true
	}
	if capabilities["session_revocation"] && !capabilities["token_introspection"] ||
		capabilities["token_introspection"] && !capabilities["machine_to_machine"] {
		return errors.New("application manifest capability dependencies are invalid")
	}
	if capabilities["oidc_login"] && (manifest.Clients.BrowserLogin == nil || strings.TrimSpace(manifest.Clients.BrowserLogin.ClientID) == "") {
		return errors.New("application manifest browser client is missing")
	}
	if capabilities["machine_to_machine"] && (manifest.Clients.MachineToMachine == nil || strings.TrimSpace(manifest.Clients.MachineToMachine.ClientID) == "") {
		return errors.New("application manifest machine client is missing")
	}
	if capabilities["api_access"] && strings.TrimSpace(manifest.Audience) == "" {
		return errors.New("application manifest audience is missing")
	}
	if capabilities["session_revocation"] {
		if manifest.SecurityEvents == nil || validatePublicIdentityURL(manifest.SecurityEvents.TransmitterIssuer) != nil ||
			validatePublicIdentityURL(manifest.SecurityEvents.ConfigurationEndpoint) != nil || strings.TrimSpace(manifest.SecurityEvents.Audience) == "" {
			return errors.New("application manifest security event metadata is invalid")
		}
	}
	seenScopes := make(map[string]bool, len(manifest.Scopes))
	for _, scope := range manifest.Scopes {
		scope = strings.TrimSpace(scope)
		if scope == "" || seenScopes[scope] {
			return errors.New("application manifest scopes are invalid")
		}
		seenScopes[scope] = true
	}
	return nil
}

func validatePublicIdentityURL(raw string) error {
	_, err := exactBaseURL(raw)
	if err == nil {
		return nil
	}
	// Discovery endpoints may contain a path, but must retain the same URL
	// safety constraints as a base URL.
	parsed, parseErr := http.NewRequest(http.MethodGet, strings.TrimSpace(raw), nil)
	if parseErr != nil || parsed.URL.User != nil || parsed.URL.Host == "" || parsed.URL.RawQuery != "" || parsed.URL.Fragment != "" {
		return errors.New("identity URL is invalid")
	}
	scheme := strings.ToLower(parsed.URL.Scheme)
	if scheme == "https" || scheme == "http" && isLoopbackHost(strings.ToLower(parsed.URL.Hostname())) {
		return nil
	}
	return errors.New("identity URL must use HTTPS")
}

type OnboardingClient struct {
	baseURL string
	client  *http.Client
}

func NewOnboardingClient(baseURL string, base *http.Client) (*OnboardingClient, error) {
	normalized, err := exactBaseURL(baseURL)
	if err != nil {
		return nil, errors.New("auth center URL is invalid")
	}
	if base == nil {
		base = &http.Client{Timeout: 10 * time.Second}
	}
	client := *base
	if client.Timeout <= 0 {
		client.Timeout = 10 * time.Second
	}
	client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { return http.ErrUseLastResponse }
	return &OnboardingClient{baseURL: normalized, client: &client}, nil
}

func (client *OnboardingClient) Claim(ctx context.Context, code string) (ClaimedExchange, error) {
	var output ClaimedExchange
	status, _, err := client.request(ctx, http.MethodPost, "/api/v1/onboarding-exchanges", "", "", 0,
		map[string]string{"onboarding_code": strings.TrimSpace(code)}, &output)
	if err != nil || status != http.StatusCreated || output.Version < 1 || output.ExchangeID == "" || output.ExchangeToken == "" {
		return ClaimedExchange{}, onboardingProtocolError(err)
	}
	return output, nil
}

func (client *OnboardingClient) SubmitMetadata(ctx context.Context, claimed ClaimedExchange, metadata ConsumerMetadata, idempotencyKey string) (Exchange, error) {
	var output Exchange
	status, _, err := client.request(ctx, http.MethodPut, "/api/v1/onboarding-exchanges/"+claimed.ExchangeID+"/metadata",
		claimed.ExchangeToken, idempotencyKey, claimed.Version, metadata, &output)
	if err != nil || status != http.StatusAccepted || output.Version < 1 {
		return Exchange{}, onboardingProtocolError(err)
	}
	return output, nil
}

func (client *OnboardingClient) Get(ctx context.Context, exchangeID, exchangeToken string) (Exchange, error) {
	var output Exchange
	status, _, err := client.request(ctx, http.MethodGet, "/api/v1/onboarding-exchanges/"+exchangeID, exchangeToken, "", 0, nil, &output)
	if err != nil || status != http.StatusOK || output.Version < 1 {
		return Exchange{}, onboardingProtocolError(err)
	}
	return output, nil
}

func (client *OnboardingClient) DeliverCredential(ctx context.Context, exchange Exchange, exchangeToken, idempotencyKey string) (CredentialDelivery, error) {
	var output CredentialDelivery
	status, etag, err := client.request(ctx, http.MethodPost, "/api/v1/onboarding-exchanges/"+exchange.ExchangeID+"/credential-deliveries",
		exchangeToken, idempotencyKey, exchange.Version, nil, &output)
	if err != nil || status != http.StatusOK {
		return CredentialDelivery{}, onboardingProtocolError(err)
	}
	if err := output.Manifest.Validate(); err != nil {
		return CredentialDelivery{}, err
	}
	output.Version, err = parseWeakETag(etag)
	if err != nil {
		return CredentialDelivery{}, errors.New("onboarding response ETag is invalid")
	}
	return output, nil
}

func (client *OnboardingClient) Complete(ctx context.Context, exchangeID, exchangeToken, idempotencyKey string, version int64) error {
	status, _, err := client.request(ctx, http.MethodPost, "/api/v1/onboarding-exchanges/"+exchangeID+"/completion",
		exchangeToken, idempotencyKey, version, nil, nil)
	if err != nil || status != http.StatusNoContent {
		return onboardingProtocolError(err)
	}
	return nil
}

func (client *OnboardingClient) request(ctx context.Context, method, path, token, idempotencyKey string, version int64, body, output any) (int, string, error) {
	var reader io.Reader
	if body != nil {
		payload, err := json.Marshal(body)
		if err != nil {
			return 0, "", err
		}
		reader = bytes.NewReader(payload)
	}
	request, err := http.NewRequestWithContext(ctx, method, client.baseURL+path, reader)
	if err != nil {
		return 0, "", err
	}
	request.Header.Set("Accept", "application/json")
	if body != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	if token != "" {
		request.Header.Set("Authorization", "Bearer "+token)
	}
	if idempotencyKey != "" {
		request.Header.Set("Idempotency-Key", idempotencyKey)
	}
	if version > 0 {
		request.Header.Set("If-Match", fmt.Sprintf(`W/"%d"`, version))
	}
	response, err := client.client.Do(request)
	if err != nil {
		return 0, "", errors.New("onboarding endpoint request failed")
	}
	defer response.Body.Close()
	if response.StatusCode >= 300 && response.StatusCode < 400 {
		return response.StatusCode, "", errors.New("onboarding endpoint redirect is forbidden")
	}
	if output == nil || response.StatusCode == http.StatusNoContent {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, maxOnboardingResponseBytes))
		return response.StatusCode, response.Header.Get("ETag"), nil
	}
	payload, err := io.ReadAll(io.LimitReader(response.Body, maxOnboardingResponseBytes+1))
	if err != nil || len(payload) > maxOnboardingResponseBytes {
		return response.StatusCode, "", errors.New("onboarding endpoint response is invalid")
	}
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		return response.StatusCode, "", errors.New("onboarding endpoint rejected the request")
	}
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(output); err != nil {
		return response.StatusCode, "", errors.New("onboarding endpoint response is invalid")
	}
	return response.StatusCode, response.Header.Get("ETag"), nil
}

func onboardingProtocolError(cause error) error {
	if cause == nil {
		return errors.New("onboarding endpoint returned an unexpected status")
	}
	return cause
}

func parseWeakETag(value string) (int64, error) {
	value = strings.TrimSpace(value)
	if !strings.HasPrefix(value, `W/"`) || !strings.HasSuffix(value, `"`) {
		return 0, errors.New("invalid ETag")
	}
	version, err := strconv.ParseInt(strings.TrimSuffix(strings.TrimPrefix(value, `W/"`), `"`), 10, 64)
	if err != nil || version < 1 {
		return 0, errors.New("invalid ETag")
	}
	return version, nil
}