story/internal/deviceauth/database_secret.go

package deviceauth

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"strings"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const databaseCredentialKeyBytes = 32

var ErrCredentialAlreadyExists = errors.New("active device credential already exists")

type DatabaseSecretResolver struct {
	pool *pgxpool.Pool
	aead cipher.AEAD
}

func NewDatabaseSecretResolver(pool *pgxpool.Pool, masterKeyHex string) (*DatabaseSecretResolver, error) {
	if pool == nil {
		return nil, errors.New("device credential database is required")
	}
	key, err := hex.DecodeString(strings.TrimSpace(masterKeyHex))
	if err != nil || len(key) != databaseCredentialKeyBytes {
		return nil, errors.New("device credential master key must be 32 bytes encoded as 64 hexadecimal characters")
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("create device credential cipher: %w", err)
	}
	aead, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("create device credential AEAD: %w", err)
	}
	return &DatabaseSecretResolver{pool: pool, aead: aead}, nil
}

func (r *DatabaseSecretResolver) Resolve(ctx context.Context, reference string) ([]byte, error) {
	var encrypted, nonce []byte
	var version int64
	err := r.pool.QueryRow(ctx, `
		SELECT encrypted_secret,nonce,key_version
		FROM device_credential
		WHERE credential_ref=$1 AND status='active'`, strings.TrimSpace(reference),
	).Scan(&encrypted, &nonce, &version)
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, ErrCredentialInvalid
	}
	if err != nil {
		return nil, fmt.Errorf("read device credential: %w", err)
	}
	secret, err := r.aead.Open(nil, nonce, encrypted, credentialAAD(reference, version))
	if err != nil || len(secret) < 32 {
		return nil, ErrCredentialInvalid
	}
	return secret, nil
}

// ImportMissingLegacyCredentials moves existing deployment credentials into
// the database without replacing an already imported or rotated credential.
func (r *DatabaseSecretResolver) ImportMissingLegacyCredentials(ctx context.Context, credentials map[string][]byte) error {
	for reference, secret := range credentials {
		if len(secret) < 32 {
			return ErrCredentialInvalid
		}
		var deviceID, tenantID string
		err := r.pool.QueryRow(ctx, `SELECT id::text,tenant_id::text FROM device WHERE credential_ref=$1`, reference).Scan(&deviceID, &tenantID)
		if errors.Is(err, pgx.ErrNoRows) {
			continue
		}
		if err != nil {
			return fmt.Errorf("resolve legacy device credential %s: %w", reference, err)
		}
		nonce := make([]byte, r.aead.NonceSize())
		if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
			return fmt.Errorf("generate device credential nonce: %w", err)
		}
		const version int64 = 1
		encrypted := r.aead.Seal(nil, nonce, secret, credentialAAD(reference, version))
		if _, err := r.pool.Exec(ctx, `
			INSERT INTO device_credential(device_id,tenant_id,credential_ref,key_version,encrypted_secret,nonce)
			VALUES($1::uuid,$2::uuid,$3,$4,$5,$6)
			ON CONFLICT (device_id) DO NOTHING`, deviceID, tenantID, reference, version, encrypted, nonce); err != nil {
			return fmt.Errorf("import legacy device credential %s: %w", reference, err)
		}
	}
	return nil
}

func (r *DatabaseSecretResolver) Provision(ctx context.Context, devicePublicID string, secret []byte, rotate bool) (int64, error) {
	if len(secret) < 32 {
		return 0, ErrCredentialInvalid
	}
	tx, err := r.pool.Begin(ctx)
	if err != nil {
		return 0, fmt.Errorf("begin device credential provisioning: %w", err)
	}
	defer tx.Rollback(ctx)

	version, err := r.ProvisionTx(ctx, tx, devicePublicID, secret, rotate)
	if err != nil {
		return 0, err
	}
	if err := tx.Commit(ctx); err != nil {
		return 0, fmt.Errorf("commit device credential provisioning: %w", err)
	}
	return version, nil
}

// ProvisionTx stores a credential inside a caller-owned transaction so device
// registration and first credential issuance can commit atomically.
func (r *DatabaseSecretResolver) ProvisionTx(ctx context.Context, tx pgx.Tx, devicePublicID string, secret []byte, rotate bool) (int64, error) {
	if len(secret) < 32 {
		return 0, ErrCredentialInvalid
	}
	var deviceID, tenantID, reference string
	if err := tx.QueryRow(ctx, `
		SELECT id::text,tenant_id::text,credential_ref FROM device
		WHERE public_id=$1::uuid AND status <> 'revoked' FOR UPDATE`, strings.TrimSpace(devicePublicID),
	).Scan(&deviceID, &tenantID, &reference); errors.Is(err, pgx.ErrNoRows) {
		return 0, ErrDeviceNotFound
	} else if err != nil {
		return 0, fmt.Errorf("resolve device for credential provisioning: %w", err)
	}
	var currentVersion int64
	err := tx.QueryRow(ctx, `SELECT key_version FROM device_credential WHERE device_id=$1::uuid FOR UPDATE`, deviceID).Scan(&currentVersion)
	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return 0, fmt.Errorf("read current device credential: %w", err)
	}
	if err == nil && !rotate {
		return 0, ErrCredentialAlreadyExists
	}
	version := currentVersion + 1
	nonce := make([]byte, r.aead.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return 0, fmt.Errorf("generate device credential nonce: %w", err)
	}
	encrypted := r.aead.Seal(nil, nonce, secret, credentialAAD(reference, version))
	if currentVersion == 0 {
		_, err = tx.Exec(ctx, `
			INSERT INTO device_credential(device_id,tenant_id,credential_ref,key_version,encrypted_secret,nonce)
			VALUES($1::uuid,$2::uuid,$3,$4,$5,$6)`, deviceID, tenantID, reference, version, encrypted, nonce)
	} else {
		_, err = tx.Exec(ctx, `
			UPDATE device_credential SET key_version=$2,encrypted_secret=$3,nonce=$4,status='active',
				rotated_at=now(),revoked_at=NULL
			WHERE device_id=$1::uuid`, deviceID, version, encrypted, nonce)
	}
	if err != nil {
		return 0, fmt.Errorf("store device credential: %w", err)
	}
	return version, nil
}

func (r *DatabaseSecretResolver) Revoke(ctx context.Context, devicePublicID string) error {
	result, err := r.pool.Exec(ctx, `
		UPDATE device_credential dc SET status='revoked',revoked_at=now()
		FROM device d WHERE d.id=dc.device_id AND d.public_id=$1::uuid AND dc.status='active'`, strings.TrimSpace(devicePublicID))
	if err != nil {
		return fmt.Errorf("revoke device credential: %w", err)
	}
	if result.RowsAffected() != 1 {
		return ErrDeviceNotFound
	}
	return nil
}

func credentialAAD(reference string, version int64) []byte {
	return []byte(fmt.Sprintf("device-credential-v1\n%s\n%d", strings.TrimSpace(reference), version))
}