rdp-proxy/backend/internal/platform/authority/authority.go

package authority

import (
	"context"
	"crypto/ed25519"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"

	"github.com/example/remote-access-platform/backend/internal/platform/config"
	postgresplatform "github.com/example/remote-access-platform/backend/internal/platform/postgres"
)

const (
	ModeStrict = "strict"
	ModeCompat = "compat"

	ActivationSchemaVersion = "rap.installation.activation.v1"

	PlatformRoleUser          = "user"
	PlatformRoleAdmin         = "platform_admin"
	PlatformRoleRecoveryAdmin = "platform_recovery_admin"
)

var (
	ErrInvalidAuthorityMode = errors.New("invalid installation authority mode")
	ErrProductRootKeyNeeded = errors.New("product root public key is required")
	ErrInvalidActivation    = errors.New("invalid installation activation")
	ErrInvalidGrant         = errors.New("invalid platform role grant")
)

type ActivationPayload struct {
	SchemaVersion string     `json:"schema_version"`
	InstallID     string     `json:"install_id"`
	OwnerEmail    string     `json:"owner_email"`
	PlatformRole  string     `json:"platform_role"`
	IssuedAt      time.Time  `json:"issued_at"`
	ExpiresAt     *time.Time `json:"expires_at,omitempty"`
	Nonce         string     `json:"nonce,omitempty"`
	Environment   string     `json:"environment,omitempty"`
}

type Verifier struct {
	mode                   string
	rootPublicKey          ed25519.PublicKey
	rootFingerprint        string
	allowInsecureBootstrap bool
	now                    func() time.Time
}

func NewVerifier(cfg config.InstallationConfig) (*Verifier, error) {
	mode := strings.ToLower(strings.TrimSpace(cfg.AuthorityMode))
	if mode == "" {
		mode = ModeCompat
	}
	verifier := &Verifier{
		mode:                   mode,
		allowInsecureBootstrap: cfg.AllowInsecureBootstrap,
		now:                    time.Now,
	}

	switch mode {
	case ModeCompat:
		return verifier, nil
	case ModeStrict:
		publicKey, err := decodeEd25519PublicKey(cfg.ProductRootPublicKeyBase64)
		if err != nil {
			return nil, err
		}
		verifier.rootPublicKey = publicKey
		fingerprint := sha256.Sum256(publicKey)
		verifier.rootFingerprint = hex.EncodeToString(fingerprint[:])
		return verifier, nil
	default:
		return nil, fmt.Errorf("%w: %s", ErrInvalidAuthorityMode, mode)
	}
}

func (v *Verifier) Mode() string {
	if v == nil || v.mode == "" {
		return ModeCompat
	}
	return v.mode
}

func (v *Verifier) Strict() bool {
	return v != nil && v.mode == ModeStrict
}

func (v *Verifier) AllowInsecureBootstrap() bool {
	return v != nil && v.allowInsecureBootstrap
}

func (v *Verifier) RootFingerprint() string {
	if v == nil {
		return ""
	}
	return v.rootFingerprint
}

func (v *Verifier) VerifyActivation(payload json.RawMessage, signature string) (ActivationPayload, error) {
	if v == nil || !v.Strict() {
		return ActivationPayload{}, ErrProductRootKeyNeeded
	}
	activation, canonical, err := parseActivationPayload(payload)
	if err != nil {
		return ActivationPayload{}, err
	}
	if err := activation.validate(v.now().UTC()); err != nil {
		return ActivationPayload{}, err
	}
	if err := v.verifySignature(canonical, signature); err != nil {
		return ActivationPayload{}, fmt.Errorf("%w: %v", ErrInvalidActivation, err)
	}
	return activation, nil
}

func (v *Verifier) VerifyPlatformRoleGrant(payload json.RawMessage, signature, expectedInstallID, expectedEmail, expectedRole string) (ActivationPayload, error) {
	activation, err := v.VerifyActivation(payload, signature)
	if err != nil {
		return ActivationPayload{}, fmt.Errorf("%w: %v", ErrInvalidGrant, err)
	}
	if activation.InstallID != strings.TrimSpace(expectedInstallID) {
		return ActivationPayload{}, fmt.Errorf("%w: install_id mismatch", ErrInvalidGrant)
	}
	if !strings.EqualFold(activation.OwnerEmail, strings.TrimSpace(expectedEmail)) {
		return ActivationPayload{}, fmt.Errorf("%w: owner_email mismatch", ErrInvalidGrant)
	}
	if activation.PlatformRole != strings.TrimSpace(expectedRole) {
		return ActivationPayload{}, fmt.Errorf("%w: platform_role mismatch", ErrInvalidGrant)
	}
	return activation, nil
}

func CanonicalJSON(raw json.RawMessage) ([]byte, error) {
	if len(raw) == 0 {
		return nil, fmt.Errorf("%w: empty payload", ErrInvalidActivation)
	}
	var value any
	if err := json.Unmarshal(raw, &value); err != nil {
		return nil, fmt.Errorf("%w: invalid json: %v", ErrInvalidActivation, err)
	}
	canonical, err := json.Marshal(value)
	if err != nil {
		return nil, fmt.Errorf("%w: canonical json: %v", ErrInvalidActivation, err)
	}
	return canonical, nil
}

func EffectivePlatformRole(ctx context.Context, db postgresplatform.DBTX, verifier *Verifier, userID string) (string, error) {
	userID = strings.TrimSpace(userID)
	if userID == "" {
		return PlatformRoleUser, nil
	}
	if verifier == nil || !verifier.Strict() {
		return storedPlatformRole(ctx, db, userID)
	}

	var email string
	if err := db.QueryRow(ctx, `SELECT email FROM users WHERE id = $1::uuid`, userID).Scan(&email); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return PlatformRoleUser, nil
		}
		return "", fmt.Errorf("get user email for platform grant: %w", err)
	}

	rows, err := db.Query(ctx, `
SELECT prg.role, prg.install_id, prg.grant_payload, prg.grant_signature
FROM platform_role_grants prg
JOIN installation_authority ia
  ON ia.id = 1
 AND ia.install_id = prg.install_id
 AND ia.authority_state = 'active'
WHERE prg.user_id = $1::uuid
  AND prg.revoked_at IS NULL
  AND (prg.expires_at IS NULL OR prg.expires_at > NOW())
ORDER BY CASE prg.role
    WHEN 'platform_recovery_admin' THEN 0
    WHEN 'platform_admin' THEN 1
    ELSE 2
END, prg.granted_at DESC
`, userID)
	if err != nil {
		return "", fmt.Errorf("query platform role grants: %w", err)
	}
	defer rows.Close()

	bestRole := PlatformRoleUser
	for rows.Next() {
		var role, installID, signature string
		var payload []byte
		if err := rows.Scan(&role, &installID, &payload, &signature); err != nil {
			return "", fmt.Errorf("scan platform role grant: %w", err)
		}
		if _, err := verifier.VerifyPlatformRoleGrant(json.RawMessage(payload), signature, installID, email, role); err != nil {
			continue
		}
		if role == PlatformRoleRecoveryAdmin {
			return role, nil
		}
		if role == PlatformRoleAdmin {
			bestRole = role
		}
	}
	if err := rows.Err(); err != nil {
		return "", fmt.Errorf("iterate platform role grants: %w", err)
	}
	if bestRole == PlatformRoleUser {
		if role, ok, err := strictBootstrappedOwnerFallback(ctx, db, verifier, userID, email); err != nil {
			return "", err
		} else if ok {
			return role, nil
		}
		return storedPlatformRole(ctx, db, userID)
	}
	return bestRole, nil
}

func strictBootstrappedOwnerFallback(ctx context.Context, db postgresplatform.DBTX, verifier *Verifier, userID, email string) (string, bool, error) {
	var role string
	var bootstrappedOwnerEmail *string
	var authorityState string
	var rootFingerprint string
	err := db.QueryRow(ctx, `
SELECT u.platform_role, ia.bootstrapped_owner_email, ia.authority_state, ia.product_root_key_fingerprint
FROM users u
CROSS JOIN installation_authority ia
WHERE u.id = $1::uuid
  AND ia.id = 1
`, userID).Scan(&role, &bootstrappedOwnerEmail, &authorityState, &rootFingerprint)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return PlatformRoleUser, false, nil
		}
		return "", false, fmt.Errorf("query strict bootstrapped owner fallback: %w", err)
	}
	if bootstrappedOwnerEmail == nil ||
		!strings.EqualFold(*bootstrappedOwnerEmail, email) ||
		authorityState != "active" ||
		rootFingerprint != verifier.RootFingerprint() {
		return PlatformRoleUser, false, nil
	}
	switch role {
	case PlatformRoleAdmin, PlatformRoleRecoveryAdmin:
		return role, true, nil
	default:
		return PlatformRoleUser, false, nil
	}
}

func storedPlatformRole(ctx context.Context, db postgresplatform.DBTX, userID string) (string, error) {
	var role string
	if err := db.QueryRow(ctx, `SELECT platform_role FROM users WHERE id = $1::uuid`, userID).Scan(&role); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return PlatformRoleUser, nil
		}
		return "", fmt.Errorf("get platform role: %w", err)
	}
	if role == "" {
		return PlatformRoleUser, nil
	}
	return role, nil
}

func parseActivationPayload(raw json.RawMessage) (ActivationPayload, []byte, error) {
	canonical, err := CanonicalJSON(raw)
	if err != nil {
		return ActivationPayload{}, nil, err
	}
	var activation ActivationPayload
	if err := json.Unmarshal(canonical, &activation); err != nil {
		return ActivationPayload{}, nil, fmt.Errorf("%w: decode activation: %v", ErrInvalidActivation, err)
	}
	activation.SchemaVersion = strings.TrimSpace(activation.SchemaVersion)
	activation.InstallID = strings.TrimSpace(activation.InstallID)
	activation.OwnerEmail = strings.ToLower(strings.TrimSpace(activation.OwnerEmail))
	activation.PlatformRole = strings.TrimSpace(activation.PlatformRole)
	activation.Nonce = strings.TrimSpace(activation.Nonce)
	activation.Environment = strings.TrimSpace(activation.Environment)
	return activation, canonical, nil
}

func (p ActivationPayload) validate(now time.Time) error {
	if p.SchemaVersion != ActivationSchemaVersion {
		return fmt.Errorf("%w: schema_version must be %s", ErrInvalidActivation, ActivationSchemaVersion)
	}
	if p.InstallID == "" {
		return fmt.Errorf("%w: install_id is required", ErrInvalidActivation)
	}
	if p.OwnerEmail == "" || !strings.Contains(p.OwnerEmail, "@") {
		return fmt.Errorf("%w: owner_email is required", ErrInvalidActivation)
	}
	switch p.PlatformRole {
	case PlatformRoleAdmin, PlatformRoleRecoveryAdmin:
	default:
		return fmt.Errorf("%w: platform_role must be platform_admin or platform_recovery_admin", ErrInvalidActivation)
	}
	if p.IssuedAt.IsZero() {
		return fmt.Errorf("%w: issued_at is required", ErrInvalidActivation)
	}
	if p.IssuedAt.After(now.Add(5 * time.Minute)) {
		return fmt.Errorf("%w: issued_at is too far in the future", ErrInvalidActivation)
	}
	if p.ExpiresAt != nil && !p.ExpiresAt.After(now) {
		return fmt.Errorf("%w: activation expired", ErrInvalidActivation)
	}
	return nil
}

func (v *Verifier) verifySignature(payload []byte, signatureText string) error {
	signature, err := decodeBase64(strings.TrimSpace(signatureText))
	if err != nil {
		return fmt.Errorf("signature must be base64 encoded: %w", err)
	}
	if len(signature) != ed25519.SignatureSize {
		return fmt.Errorf("signature must decode to %d bytes", ed25519.SignatureSize)
	}
	if !ed25519.Verify(v.rootPublicKey, payload, signature) {
		return errors.New("signature verification failed")
	}
	return nil
}

func decodeEd25519PublicKey(value string) (ed25519.PublicKey, error) {
	value = strings.TrimSpace(value)
	if value == "" {
		return nil, ErrProductRootKeyNeeded
	}
	if block, _ := pem.Decode([]byte(value)); block != nil {
		parsed, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("parse product root public key PEM: %w", err)
		}
		publicKey, ok := parsed.(ed25519.PublicKey)
		if !ok {
			return nil, fmt.Errorf("product root public key PEM must contain an Ed25519 public key")
		}
		return publicKey, nil
	}
	decoded, err := decodeBase64(value)
	if err != nil {
		return nil, fmt.Errorf("product root public key must be base64 encoded: %w", err)
	}
	if len(decoded) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("product root public key must decode to %d bytes", ed25519.PublicKeySize)
	}
	return ed25519.PublicKey(decoded), nil
}

func decodeBase64(value string) ([]byte, error) {
	decoded, err := base64.StdEncoding.DecodeString(value)
	if err == nil {
		return decoded, nil
	}
	decoded, rawErr := base64.RawStdEncoding.DecodeString(value)
	if rawErr == nil {
		return decoded, nil
	}
	return nil, err
}