easyai-ai-gateway/apps/api/internal/store/oidc_platform_users.go

package store

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

const oidcPlatformTenantKey = "default"

func (s *Store) resolveOrProvisionOIDCPlatformUser(
	ctx context.Context,
	input ResolveOrProvisionOIDCUserInput,
) (ResolveOrProvisionOIDCUserResult, error) {
	if input.Issuer == "" || input.Subject == "" ||
		input.ContextType != "platform" ||
		input.TenantMode != "multi_tenant" ||
		uuid.Validate(input.ApplicationID) != nil ||
		input.TenantID != "" {
		return ResolveOrProvisionOIDCUserResult{},
			errors.New("invalid platform OIDC user projection input")
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return ResolveOrProvisionOIDCUserResult{}, err
	}
	defer rollbackTransaction(tx)

	userKey := deriveOIDCPlatformUserKey(
		input.Issuer,
		input.ApplicationID,
		input.Subject,
	)
	user, groupKey, err := loadOIDCPlatformUser(ctx, tx, userKey)
	created := false
	if errors.Is(err, pgx.ErrNoRows) {
		if !input.ProvisioningEnabled {
			return ResolveOrProvisionOIDCUserResult{},
				ErrOIDCUserNotProvisioned
		}
		user, groupKey, created, err = s.createOIDCPlatformUser(
			ctx,
			tx,
			userKey,
			input,
		)
	}
	if err != nil {
		return ResolveOrProvisionOIDCUserResult{}, err
	}
	if user.Status != "active" {
		return ResolveOrProvisionOIDCUserResult{}, ErrOIDCUserDisabled
	}
	rolesJSON, err := json.Marshal(input.Roles)
	if err != nil {
		return ResolveOrProvisionOIDCUserResult{}, err
	}
	user, err = scanUser(tx.QueryRow(ctx, `UPDATE gateway_users SET
		username=COALESCE(NULLIF($2,''),username),roles=$3::jsonb,
		last_login_at=now(),synced_at=now(),source_updated_at=now(),
		updated_at=now()
		WHERE id=$1::uuid AND source='oidc_v2_platform'
		  AND status='active' AND deleted_at IS NULL
		RETURNING `+userColumns,
		user.ID,
		input.Username,
		string(rolesJSON),
	))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return ResolveOrProvisionOIDCUserResult{}, ErrOIDCUserDisabled
		}
		return ResolveOrProvisionOIDCUserResult{}, err
	}
	var auditID string
	if created {
		subjectHash := sha256.Sum256([]byte(input.Subject))
		audit, auditErr := s.RecordAuditLogTx(ctx, tx, AuditLogInput{
			Category:              "identity",
			Action:                "identity.oidc_platform_user.provisioned",
			ActorGatewayUserID:    user.ID,
			ActorUsername:         user.Username,
			ActorSource:           "oidc",
			ActorRoles:            user.Roles,
			TargetType:            "gateway_user",
			TargetID:              user.ID,
			TargetGatewayUserID:   user.ID,
			TargetGatewayTenantID: user.GatewayTenantID,
			RequestIP:             input.RequestIP,
			UserAgent:             input.UserAgent,
			AfterState: map[string]any{
				"source":      "oidc_v2_platform",
				"contextType": "platform",
				"tenantKey":   user.TenantKey,
				"userGroupId": user.DefaultUserGroupID,
			},
			Metadata: map[string]any{
				"provisioningMode":    "oidc-platform-jit",
				"externalSubjectHash": hex.EncodeToString(subjectHash[:])[:16],
			},
		})
		if auditErr != nil {
			return ResolveOrProvisionOIDCUserResult{}, auditErr
		}
		auditID = audit.ID
	}
	if err := tx.Commit(ctx); err != nil {
		return ResolveOrProvisionOIDCUserResult{}, err
	}
	return ResolveOrProvisionOIDCUserResult{
		User:    platformAuthUser(user, groupKey, input),
		Created: created,
		AuditID: auditID,
	}, nil
}

func loadOIDCPlatformUser(
	ctx context.Context,
	tx pgx.Tx,
	userKey string,
) (GatewayUser, string, error) {
	var userID string
	if err := tx.QueryRow(ctx, `SELECT id::text
		FROM gateway_users
		WHERE user_key=$1 AND source='oidc_v2_platform'
		FOR UPDATE`, userKey).Scan(&userID); err != nil {
		return GatewayUser{}, "", err
	}
	user, err := scanUser(tx.QueryRow(ctx, `SELECT `+userColumns+`
		FROM gateway_users WHERE id=$1::uuid`, userID))
	if err != nil {
		return GatewayUser{}, "", err
	}
	var groupKey string
	if err := tx.QueryRow(ctx, `SELECT group_record.group_key
		FROM gateway_users user_record
		JOIN gateway_tenants tenant
		  ON tenant.id=user_record.gateway_tenant_id
		 AND tenant.status='active' AND tenant.deleted_at IS NULL
		JOIN gateway_user_groups group_record
		  ON group_record.id=user_record.default_user_group_id
		 AND group_record.status='active'
		WHERE user_record.id=$1::uuid`, userID).Scan(&groupKey); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return GatewayUser{}, "", ErrOIDCTenantUnavailable
		}
		return GatewayUser{}, "", err
	}
	return user, groupKey, nil
}

func (s *Store) createOIDCPlatformUser(
	ctx context.Context,
	tx pgx.Tx,
	userKey string,
	input ResolveOrProvisionOIDCUserInput,
) (GatewayUser, string, bool, error) {
	gatewayTenantID, groupID, groupKey, err :=
		loadOIDCProvisioningTenant(ctx, tx, oidcPlatformTenantKey)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return GatewayUser{}, "", false, ErrOIDCTenantUnavailable
		}
		return GatewayUser{}, "", false, err
	}
	rolesJSON, err := json.Marshal(input.Roles)
	if err != nil {
		return GatewayUser{}, "", false, err
	}
	metadataJSON, err := json.Marshal(map[string]any{
		"provisioningMode": "oidc-platform-jit",
		"contextType":      "platform",
		"applicationId":    input.ApplicationID,
	})
	if err != nil {
		return GatewayUser{}, "", false, err
	}
	username := input.Username
	if username == "" {
		username = "oidc-platform-" +
			strings.TrimPrefix(userKey, "oidc2-platform:")[:12]
	}
	user, err := scanUser(tx.QueryRow(ctx, `INSERT INTO gateway_users(
		user_key,source,external_user_id,username,gateway_tenant_id,
		tenant_id,tenant_key,default_user_group_id,roles,auth_profile,
		metadata,status,last_login_at,synced_at,source_updated_at
	) VALUES(
		$1,'oidc_v2_platform',NULL,$2,$3::uuid,NULL,$4,$5::uuid,
		$6::jsonb,'{}'::jsonb,$7::jsonb,'active',now(),now(),now()
	)
	ON CONFLICT(user_key) DO NOTHING
	RETURNING `+userColumns,
		userKey,
		username,
		gatewayTenantID,
		oidcPlatformTenantKey,
		groupID,
		string(rolesJSON),
		string(metadataJSON),
	))
	if errors.Is(err, pgx.ErrNoRows) {
		user, groupKey, err = loadOIDCPlatformUser(ctx, tx, userKey)
		return user, groupKey, false, err
	}
	if err != nil {
		return GatewayUser{}, "", false, err
	}
	if _, err := s.ensureWalletAccount(ctx, tx, user.ID, "resource"); err != nil {
		return GatewayUser{}, "", false, err
	}
	return user, groupKey, true, nil
}

func deriveOIDCPlatformUserKey(
	issuer string,
	applicationID string,
	subject string,
) string {
	sum := sha256.Sum256([]byte(
		strings.TrimRight(issuer, "/") + "\x00" +
			applicationID + "\x00" + subject,
	))
	return fmt.Sprintf("oidc2-platform:%x", sum)
}

func platformAuthUser(
	user GatewayUser,
	groupKey string,
	input ResolveOrProvisionOIDCUserInput,
) *auth.User {
	result := authUserFromOIDCProjection(user, groupKey)
	result.ID = input.Subject
	result.ContextType = "platform"
	result.TenantID = ""
	result.Issuer = input.Issuer
	result.ApplicationID = input.ApplicationID
	result.OIDCClientID = input.OIDCClientID
	return result
}