easyai-ai-gateway/apps/api/internal/identityruntime/builder.go

package identityruntime

import (
	"context"
	"errors"
	"net/http"
	"net/url"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/identity"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/oidcsession"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/securityevents"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type RuntimeBuilderConfig struct {
	AppEnv            string
	JWKSCacheTTL      time.Duration
	HeartbeatInterval time.Duration
	StaleAfter        time.Duration
	ClockSkew         time.Duration
}

type preparedSecurityRuntime struct {
	manager         *securityevents.ConnectionManager
	cancel          context.CancelFunc
	receiverReady   bool
	blockedCategory string
}

type safeRuntimeError struct {
	category string
	cause    error
}

func (err safeRuntimeError) Error() string {
	return "identity runtime operation failed: " + err.category
}
func (err safeRuntimeError) Unwrap() error             { return err.cause }
func (err safeRuntimeError) SafeErrorCategory() string { return err.category }

type RuntimeBuilder struct {
	ctx      context.Context
	store    *store.Store
	secrets  PairingSecretStore
	config   RuntimeBuilderConfig
	metrics  *securityevents.Metrics
	mutex    sync.Mutex
	prepared map[string]preparedSecurityRuntime
}

type PairingSecretStore interface {
	Put(context.Context, string, []byte) error
	Get(context.Context, string) ([]byte, error)
	Delete(context.Context, string) error
}

func NewRuntimeBuilder(ctx context.Context, data *store.Store, secrets PairingSecretStore, config RuntimeBuilderConfig, metrics *securityevents.Metrics) *RuntimeBuilder {
	if metrics == nil {
		metrics = &securityevents.Metrics{}
	}
	if config.JWKSCacheTTL <= 0 {
		config.JWKSCacheTTL = 5 * time.Minute
	}
	return &RuntimeBuilder{ctx: ctx, store: data, secrets: secrets, config: config, metrics: metrics, prepared: map[string]preparedSecurityRuntime{}}
}

func (builder *RuntimeBuilder) PreparedSecurityEventReceiver() http.Handler {
	builder.mutex.Lock()
	defer builder.mutex.Unlock()
	if len(builder.prepared) != 1 {
		return nil
	}
	for _, prepared := range builder.prepared {
		if prepared.receiverReady {
			return prepared.manager
		}
	}
	return nil
}

func (builder *RuntimeBuilder) PreparedSecurityEventManager() *securityevents.ConnectionManager {
	builder.mutex.Lock()
	defer builder.mutex.Unlock()
	if len(builder.prepared) != 1 {
		return nil
	}
	for _, prepared := range builder.prepared {
		return prepared.manager
	}
	return nil
}

func (builder *RuntimeBuilder) Build(ctx context.Context, revision identity.Revision) (*Runtime, error) {
	if err := identity.ValidateRevisionURLs(revision, builder.config.AppEnv); err != nil {
		return nil, err
	}
	if builder.store == nil || builder.secrets == nil || revision.Issuer == "" || revision.TenantID == "" ||
		revision.Audience == "" || revision.RolePrefix == "" {
		return nil, errors.New("identity runtime configuration is incomplete")
	}
	if exists, err := builder.store.HasActiveTenantKey(ctx, revision.LocalTenantKey); err != nil {
		return nil, err
	} else if !exists {
		return nil, identity.ErrLocalTenantInvalid
	}
	runtimeCtx, cancel := context.WithCancel(builder.ctx)
	runtime := &Runtime{Revision: revision, CookieSecure: secureCookieFor(revision.PublicBaseURL), close: cancel}

	var securityManager *securityevents.ConnectionManager
	if revision.SessionRevocation {
		prepared := builder.preparedSecurityRuntime(revision.ID)
		securityManager = prepared.manager
		if securityManager != nil {
			if err := securityManager.ValidateConfiguredConnectionBinding(ctx); err != nil {
				cancel()
				return nil, err
			}
			if !prepared.receiverReady {
				cancel()
				return nil, safeRuntimeError{category: "security_event_not_ready", cause: store.ErrSecurityEventConnectionConflict}
			}
		}
		if securityManager == nil {
			var err error
			securityManager, err = builder.newSecurityEventManager(runtimeCtx, revision, true)
			if err != nil {
				cancel()
				return nil, err
			}
		}
		runtime.SecurityEvents = securityManager
		runtime.securityEventDisconnector = securityManager
	}

	var evaluator func(context.Context, auth.OIDCSecurityEventIdentity) (auth.OIDCSecurityEventEvaluation, error)
	if securityManager != nil {
		evaluator = securityManager.Evaluate
	}
	credentialProvider := func(ctx context.Context) (string, []byte, error) {
		if revision.MachineClientID == "" || revision.MachineCredentialRef == "" {
			return "", nil, errors.New("managed machine credential is unavailable")
		}
		secret, err := builder.secrets.Get(ctx, revision.MachineCredentialRef)
		return revision.MachineClientID, secret, err
	}
	verifier, err := auth.NewOIDCVerifier(auth.OIDCConfig{
		AppEnv: builder.config.AppEnv,
		Issuer: revision.Issuer, Audience: revision.Audience, TenantID: revision.TenantID,
		RolePrefix: revision.RolePrefix, RequiredScopes: append([]string(nil), revision.Scopes...),
		JWKSCacheTTL: builder.config.JWKSCacheTTL, IntrospectionEnabled: revision.TokenIntrospection,
		IntrospectionCredentialProvider: credentialProvider, SecurityEventEvaluator: evaluator,
		IntrospectionObserver:      builder.metrics.ObserveIntrospection,
		JWKSRefreshFailureObserver: func() { builder.metrics.ObserveJWKSRefreshFailure("oidc") },
	})
	if err != nil {
		cancel()
		return nil, err
	}
	if err := verifier.ValidateConfiguration(ctx); err != nil {
		cancel()
		return nil, err
	}
	runtime.Verifier = verifier

	if slices.Contains(revision.Capabilities, "oidc_login") {
		if revision.BrowserClientID == "" || revision.SessionEncryptionKeyRef == "" {
			cancel()
			return nil, errors.New("OIDC browser session configuration is incomplete")
		}
		key, err := builder.secrets.Get(ctx, revision.SessionEncryptionKeyRef)
		if err != nil {
			cancel()
			return nil, err
		}
		cipher, err := oidcsession.NewCipher(key)
		clear(key)
		if err != nil {
			cancel()
			return nil, err
		}
		client, err := auth.NewOIDCPublicClient(auth.OIDCPublicClientConfig{
			AppEnv: builder.config.AppEnv,
			Issuer: revision.Issuer, ClientID: revision.BrowserClientID,
			RedirectURI:           revision.PublicBaseURL + "/api/v1/auth/oidc/callback",
			PostLogoutRedirectURI: revision.WebBaseURL + "/", Scopes: append([]string{"openid", "profile"}, revision.Scopes...),
		})
		if err != nil {
			cancel()
			return nil, err
		}
		if err := client.ValidateConfiguration(ctx); err != nil {
			cancel()
			return nil, err
		}
		sessions, err := oidcsession.NewService(builder.store, cipher, verifier, client, oidcsession.Config{
			IdleTTL:       time.Duration(revision.SessionIdleSeconds) * time.Second,
			AbsoluteTTL:   time.Duration(revision.SessionAbsoluteSeconds) * time.Second,
			RefreshBefore: time.Duration(revision.SessionRefreshSeconds) * time.Second,
		})
		if err != nil {
			cancel()
			return nil, err
		}
		runtime.PublicClient, runtime.Sessions, runtime.SessionCipher = client, sessions, cipher
	}
	return runtime, nil
}

func (builder *RuntimeBuilder) PrepareSecurityEvents(ctx context.Context, revision identity.Revision, managementSecret []byte) error {
	if err := identity.ValidateRevisionURLs(revision, builder.config.AppEnv); err != nil {
		return safeRuntimeError{category: "configuration_invalid", cause: err}
	}
	if !revision.SessionRevocation || revision.SecurityEventIssuer == "" || revision.SecurityEventAudience == "" || revision.MachineClientID == "" {
		return errors.New("security event configuration is incomplete")
	}
	if prepared := builder.preparedSecurityRuntime(revision.ID); prepared.manager != nil && prepared.blockedCategory != "" {
		view, err := prepared.manager.Get(ctx)
		if err == nil {
			category := prepared.blockedCategory
			if view.LifecycleStatus == "retiring" || view.LifecycleStatus == "disconnect_pending" {
				category = "retirement_pending"
			}
			return safeRuntimeError{category: category, cause: store.ErrSecurityEventConnectionConflict}
		}
		if !errors.Is(err, store.ErrSecurityEventConnectionNotFound) {
			return safeSecurityEventRuntimeError(err)
		}
		builder.removePreparedSecurityRuntime(revision.ID, prepared)
	}
	runtimeCtx, cancel := context.WithCancel(builder.ctx)
	manager, err := builder.newSecurityEventManager(runtimeCtx, revision, false)
	if err != nil {
		cancel()
		return safeRuntimeError{category: "configuration_invalid", cause: err}
	}
	secretCopy := append([]byte(nil), managementSecret...)
	_, err = manager.Connect(ctx, revision.SecurityEventIssuer, revision.MachineClientID, secretCopy, "identity-pairing-ssf-"+revision.ID)
	clear(secretCopy)
	if err != nil {
		// A conflicting persisted connection is itself the resource an
		// administrator must safely retire. Keep this manager reachable by the
		// recovery API even when there is no Active identity Runtime yet.
		if errors.Is(err, store.ErrSecurityEventConnectionConflict) {
			builder.replacePreparedSecurityRuntime(revision.ID, preparedSecurityRuntime{
				manager: manager, cancel: cancel, blockedCategory: safeSecurityEventCategory(err, "connection_conflict"),
			})
			return safeSecurityEventRuntimeError(err)
		}
		cancel()
		return safeSecurityEventRuntimeError(err)
	}
	if err := manager.ValidateConfiguredConnectionBinding(ctx); err != nil {
		builder.replacePreparedSecurityRuntime(revision.ID, preparedSecurityRuntime{
			manager: manager, cancel: cancel, blockedCategory: safeSecurityEventCategory(err, "connection_binding_mismatch"),
		})
		return safeSecurityEventRuntimeError(err)
	}
	builder.replacePreparedSecurityRuntime(revision.ID, preparedSecurityRuntime{manager: manager, cancel: cancel, receiverReady: true})
	return nil
}

func (builder *RuntimeBuilder) removePreparedSecurityRuntime(revisionID string, prepared preparedSecurityRuntime) {
	builder.mutex.Lock()
	current := builder.prepared[revisionID]
	if current.manager == prepared.manager {
		delete(builder.prepared, revisionID)
	}
	builder.mutex.Unlock()
	if prepared.cancel != nil {
		prepared.cancel()
	}
}

func (builder *RuntimeBuilder) replacePreparedSecurityRuntime(revisionID string, replacement preparedSecurityRuntime) {
	builder.mutex.Lock()
	previous := builder.prepared[revisionID]
	builder.prepared[revisionID] = replacement
	builder.mutex.Unlock()
	if previous.cancel != nil {
		previous.cancel()
	}
}

func (builder *RuntimeBuilder) CleanupPreparedSecurityEvents(ctx context.Context, revision identity.Revision) error {
	builder.mutex.Lock()
	prepared, exists := builder.prepared[revision.ID]
	builder.mutex.Unlock()
	if !exists {
		runtimeCtx, cancel := context.WithCancel(builder.ctx)
		manager, err := builder.newSecurityEventManager(runtimeCtx, revision, false)
		if err != nil {
			cancel()
			return safeRuntimeError{category: "configuration_invalid", cause: err}
		}
		prepared = preparedSecurityRuntime{manager: manager, cancel: cancel}
	}
	err := prepared.manager.DiscardPreparedConnection(ctx, "identity-pairing-ssf-"+revision.ID)
	if err != nil {
		builder.mutex.Lock()
		if _, alreadyStored := builder.prepared[revision.ID]; !alreadyStored {
			builder.prepared[revision.ID] = prepared
		}
		builder.mutex.Unlock()
		return safeSecurityEventRuntimeError(err)
	}
	builder.mutex.Lock()
	current := builder.prepared[revision.ID]
	if current.manager == prepared.manager {
		delete(builder.prepared, revision.ID)
	}
	builder.mutex.Unlock()
	if prepared.cancel != nil {
		prepared.cancel()
	}
	return nil
}

func (builder *RuntimeBuilder) RetireConflictingSecurityEvents(ctx context.Context, revision identity.Revision) error {
	prepared := builder.preparedSecurityRuntime(revision.ID)
	if prepared.manager == nil {
		runtimeCtx, cancel := context.WithCancel(builder.ctx)
		manager, err := builder.newSecurityEventManager(runtimeCtx, revision, false)
		if err != nil {
			cancel()
			return safeRuntimeError{category: "configuration_invalid", cause: err}
		}
		prepared = preparedSecurityRuntime{manager: manager, cancel: cancel}
		builder.replacePreparedSecurityRuntime(revision.ID, prepared)
	}
	if err := prepared.manager.RetireConflictingConnection(ctx, "identity-pairing-ssf-"+revision.ID); err != nil {
		return safeSecurityEventRuntimeError(err)
	}
	return nil
}

// RecoverSecurityEventDisconnector reconstructs only the SSF management
// surface needed to disable a fail-closed Active Revision. It deliberately
// avoids OIDC discovery, JIT, BFF Session and tenant Runtime construction.
func (builder *RuntimeBuilder) RecoverSecurityEventDisconnector(_ context.Context, revision identity.Revision) (SecurityEventDisconnector, error) {
	prepared := builder.preparedSecurityRuntime(revision.ID)
	if prepared.manager != nil {
		if err := prepared.manager.ValidateConfiguredConnectionBinding(builder.ctx); err != nil {
			return nil, err
		}
		return prepared.manager, nil
	}
	manager, err := builder.newSecurityEventManager(builder.ctx, revision, true)
	if errors.Is(err, store.ErrSecurityEventConnectionNotFound) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return manager, nil
}

// AdoptPreparedSecurityEvents transfers the prepared manager lifetime to an
// activated Runtime. Validation only peeks at prepared state, so Verification
// callbacks remain available between validation and activation.
func (builder *RuntimeBuilder) AdoptPreparedSecurityEvents(revisionID string) context.CancelFunc {
	builder.mutex.Lock()
	defer builder.mutex.Unlock()
	prepared := builder.prepared[revisionID]
	delete(builder.prepared, revisionID)
	return prepared.cancel
}

func safeSecurityEventRuntimeError(err error) error {
	var categorized interface{ SafeErrorCategory() string }
	if errors.As(err, &categorized) {
		return err
	}
	if errors.Is(err, store.ErrSecurityEventConnectionConflict) {
		return safeRuntimeError{category: "connection_conflict", cause: err}
	}
	return safeRuntimeError{category: "preparation_failed", cause: err}
}

func safeSecurityEventCategory(err error, fallback string) string {
	var categorized interface{ SafeErrorCategory() string }
	if errors.As(err, &categorized) && categorized.SafeErrorCategory() != "" {
		return categorized.SafeErrorCategory()
	}
	return fallback
}

func (builder *RuntimeBuilder) newSecurityEventManager(ctx context.Context, revision identity.Revision, strictRevisionBinding bool) (*securityevents.ConnectionManager, error) {
	return securityevents.NewConnectionManager(ctx, builder.store, builder.secrets, builder.securityEventManagerConfig(revision, strictRevisionBinding), builder.metrics)
}

func (builder *RuntimeBuilder) securityEventManagerConfig(revision identity.Revision, strictRevisionBinding bool) securityevents.ConnectionManagerConfig {
	return securityevents.ConnectionManagerConfig{
		AppEnv: builder.config.AppEnv, OIDCEnabled: true, OIDCIssuer: revision.Issuer, OIDCTenantID: revision.TenantID,
		ManagementClientID: revision.MachineClientID, PublicBaseURL: revision.PublicBaseURL,
		ExpectedTransmitterIssuer: strings.TrimRight(strings.TrimSpace(revision.SecurityEventIssuer), "/"),
		ExpectedAudience:          revision.SecurityEventAudience,
		ExpectedOwnerKey:          "identity-pairing-ssf-" + revision.ID,
		StrictRevisionBinding:     strictRevisionBinding,
		HeartbeatInterval:         builder.config.HeartbeatInterval,
		StaleAfter:                builder.config.StaleAfter,
		ClockSkew:                 builder.config.ClockSkew,
	}
}

func (builder *RuntimeBuilder) preparedSecurityRuntime(revisionID string) preparedSecurityRuntime {
	builder.mutex.Lock()
	defer builder.mutex.Unlock()
	return builder.prepared[revisionID]
}

func secureCookieFor(baseURL string) bool {
	parsed, err := url.Parse(baseURL)
	return err == nil && parsed.Scheme == "https"
}