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

package identityruntime

import (
	"context"
	"errors"
	"net/url"
	"slices"
	"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
}

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) Build(ctx context.Context, revision identity.Revision) (*Runtime, error) {
	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.takePreparedSecurityRuntime(revision.ID)
		securityManager = prepared.manager
		if prepared.cancel != nil {
			runtime.close = func() {
				cancel()
				prepared.cancel()
			}
		}
		if securityManager == nil {
			var err error
			securityManager, err = builder.newSecurityEventManager(runtimeCtx, revision)
			if err != nil {
				cancel()
				return nil, err
			}
		}
		runtime.SecurityEvents = 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{
		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{
			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 !revision.SessionRevocation || revision.SecurityEventIssuer == "" || revision.MachineClientID == "" {
		return errors.New("security event configuration is incomplete")
	}
	runtimeCtx, cancel := context.WithCancel(builder.ctx)
	manager, err := builder.newSecurityEventManager(runtimeCtx, revision)
	if err != nil {
		cancel()
		return 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 {
		cancel()
		return err
	}
	builder.mutex.Lock()
	previous := builder.prepared[revision.ID]
	builder.prepared[revision.ID] = preparedSecurityRuntime{manager: manager, cancel: cancel}
	builder.mutex.Unlock()
	if previous.cancel != nil {
		previous.cancel()
	}
	return nil
}

func (builder *RuntimeBuilder) newSecurityEventManager(ctx context.Context, revision identity.Revision) (*securityevents.ConnectionManager, error) {
	return securityevents.NewConnectionManager(ctx, builder.store, builder.secrets, securityevents.ConnectionManagerConfig{
		AppEnv: builder.config.AppEnv, OIDCEnabled: true, OIDCIssuer: revision.Issuer, OIDCTenantID: revision.TenantID,
		ManagementClientID: revision.MachineClientID, PublicBaseURL: revision.PublicBaseURL,
		HeartbeatInterval: builder.config.HeartbeatInterval, StaleAfter: builder.config.StaleAfter, ClockSkew: builder.config.ClockSkew,
	}, builder.metrics)
}

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

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