ultisuite-backend/internal/apitokens/tokens.go

package apitokens

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

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

const tokenPrefix = "ulti_"

func TokenPrefix() string {
	return tokenPrefix
}

var (
	ErrNotFound = errors.New("api token not found")
	ErrRevoked  = errors.New("api token revoked")
	ErrExpired  = errors.New("api token expired")
)

type PermissionGrant struct {
	Resource string `json:"resource"`
	Read     bool   `json:"read"`
	Write    bool   `json:"write"`
}

type MailScope struct {
	AllAccounts bool     `json:"all_accounts"`
	AccountIDs  []string `json:"account_ids"`
}

type DriveScope struct {
	AllFolders  bool     `json:"all_folders"`
	FolderPaths []string `json:"folder_paths"`
}

type Token struct {
	ID           string            `json:"id"`
	Name         string            `json:"name"`
	TokenPrefix  string            `json:"token_prefix"`
	Permissions  []PermissionGrant `json:"permissions"`
	MailScope    MailScope         `json:"mail_scope"`
	DriveScope   DriveScope        `json:"drive_scope"`
	CreatedAt    time.Time         `json:"created_at"`
	LastUsedAt   *time.Time        `json:"last_used_at,omitempty"`
	ExpiresAt    *time.Time        `json:"expires_at,omitempty"`
}

type CreatedToken struct {
	Token
	TokenSecret string `json:"token"`
}

type AuthContext struct {
	TokenID    string
	UserID     string
	ExternalID string
	Email      string
	Name       string
	Permissions []PermissionGrant
	MailScope  MailScope
	DriveScope DriveScope
}

func HashSecret(secret string) []byte {
	sum := sha256.Sum256([]byte(secret))
	return sum[:]
}

func generateSecret() (string, string, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", "", err
	}
	encoded := base64.RawURLEncoding.EncodeToString(raw)
	full := tokenPrefix + encoded
	visible := tokenPrefix + encoded[:8]
	return full, visible, nil
}

func List(ctx context.Context, db *pgxpool.Pool, externalID string) ([]Token, error) {
	rows, err := db.Query(ctx, `
		SELECT t.id, t.name, t.token_prefix, t.permissions, t.mail_scope, t.drive_scope,
		       t.created_at, t.last_used_at, t.expires_at
		FROM api_tokens t
		JOIN users u ON u.id = t.user_id
		WHERE u.external_id = $1 AND t.revoked_at IS NULL
		ORDER BY t.created_at DESC
	`, externalID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := make([]Token, 0)
	for rows.Next() {
		item, err := scanToken(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, item)
	}
	return out, rows.Err()
}

func Create(ctx context.Context, db *pgxpool.Pool, externalID string, name string, permissions []PermissionGrant, mailScope MailScope, driveScope DriveScope, expiresAt *time.Time) (CreatedToken, error) {
	secret, prefix, err := generateSecret()
	if err != nil {
		return CreatedToken{}, err
	}

	permJSON, err := json.Marshal(permissions)
	if err != nil {
		return CreatedToken{}, err
	}
	mailJSON, err := json.Marshal(mailScope)
	if err != nil {
		return CreatedToken{}, err
	}
	driveJSON, err := json.Marshal(driveScope)
	if err != nil {
		return CreatedToken{}, err
	}

	var item Token
	err = db.QueryRow(ctx, `
		INSERT INTO api_tokens (
			user_id, name, token_prefix, secret_hash, permissions, mail_scope, drive_scope, expires_at
		)
		VALUES (
			(SELECT id FROM users WHERE external_id = $1),
			$2, $3, $4, $5, $6, $7, $8
		)
		RETURNING id, name, token_prefix, permissions, mail_scope, drive_scope, created_at, last_used_at, expires_at
	`, externalID, name, prefix, HashSecret(secret), permJSON, mailJSON, driveJSON, expiresAt).Scan(
		&item.ID,
		&item.Name,
		&item.TokenPrefix,
		&permJSON,
		&mailJSON,
		&driveJSON,
		&item.CreatedAt,
		&item.LastUsedAt,
		&item.ExpiresAt,
	)
	if err != nil {
		return CreatedToken{}, err
	}
	if err := decodeTokenJSON(permJSON, mailJSON, driveJSON, &item); err != nil {
		return CreatedToken{}, err
	}

	return CreatedToken{Token: item, TokenSecret: secret}, nil
}

func Revoke(ctx context.Context, db *pgxpool.Pool, externalID, tokenID string) error {
	result, err := db.Exec(ctx, `
		UPDATE api_tokens
		SET revoked_at = now(), updated_at = now()
		WHERE id = $1
		  AND user_id = (SELECT id FROM users WHERE external_id = $2)
		  AND revoked_at IS NULL
	`, tokenID, externalID)
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return ErrNotFound
	}
	return nil
}

func Authenticate(ctx context.Context, db *pgxpool.Pool, secret string) (*AuthContext, error) {
	secret = strings.TrimSpace(secret)
	if !strings.HasPrefix(secret, tokenPrefix) {
		return nil, fmt.Errorf("not an api token")
	}

	hash := HashSecret(secret)
	row := db.QueryRow(ctx, `
		SELECT t.id, u.id::text, u.external_id, u.email, COALESCE(u.name, ''),
		       t.permissions, t.mail_scope, t.drive_scope, t.expires_at, t.revoked_at
		FROM api_tokens t
		JOIN users u ON u.id = t.user_id
		WHERE t.secret_hash = $1
		LIMIT 1
	`, hash)

	var auth AuthContext
	var permJSON, mailJSON, driveJSON []byte
	var expiresAt *time.Time
	var revokedAt *time.Time
	if err := row.Scan(
		&auth.TokenID,
		&auth.UserID,
		&auth.ExternalID,
		&auth.Email,
		&auth.Name,
		&permJSON,
		&mailJSON,
		&driveJSON,
		&expiresAt,
		&revokedAt,
	); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrNotFound
		}
		return nil, err
	}
	if revokedAt != nil {
		return nil, ErrRevoked
	}
	if expiresAt != nil && time.Now().After(*expiresAt) {
		return nil, ErrExpired
	}
	if err := json.Unmarshal(permJSON, &auth.Permissions); err != nil {
		return nil, err
	}
	if err := json.Unmarshal(mailJSON, &auth.MailScope); err != nil {
		return nil, err
	}
	if err := json.Unmarshal(driveJSON, &auth.DriveScope); err != nil {
		return nil, err
	}

	_, _ = db.Exec(ctx, `
		UPDATE api_tokens SET last_used_at = now(), updated_at = now() WHERE id = $1
	`, auth.TokenID)

	return &auth, nil
}

func HasPermission(auth *AuthContext, resource string, write bool) bool {
	if auth == nil {
		return false
	}
	for _, grant := range auth.Permissions {
		if grant.Resource != resource {
			continue
		}
		if write {
			return grant.Write
		}
		return grant.Read || grant.Write
	}
	return false
}

func ConstantTimeEqual(a, b string) bool {
	return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
}

type rowScanner interface {
	Scan(dest ...any) error
}

func scanToken(rows rowScanner) (Token, error) {
	var item Token
	var permJSON, mailJSON, driveJSON []byte
	if err := rows.Scan(
		&item.ID,
		&item.Name,
		&item.TokenPrefix,
		&permJSON,
		&mailJSON,
		&driveJSON,
		&item.CreatedAt,
		&item.LastUsedAt,
		&item.ExpiresAt,
	); err != nil {
		return Token{}, err
	}
	if err := decodeTokenJSON(permJSON, mailJSON, driveJSON, &item); err != nil {
		return Token{}, err
	}
	return item, nil
}

func decodeTokenJSON(permJSON, mailJSON, driveJSON []byte, item *Token) error {
	if err := json.Unmarshal(permJSON, &item.Permissions); err != nil {
		return err
	}
	if err := json.Unmarshal(mailJSON, &item.MailScope); err != nil {
		return err
	}
	if err := json.Unmarshal(driveJSON, &item.DriveScope); err != nil {
		return err
	}
	return nil
}