ultisuite-backend/internal/users/avatar.go

package users

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"strings"

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

const maxAvatarBytes = 512 * 1024

var (
	ErrAvatarTooLarge   = errors.New("avatar too large")
	ErrAvatarInvalid    = errors.New("avatar invalid")
	ErrAvatarNotFound   = errors.New("avatar not found")
)

var allowedAvatarMIME = map[string]struct{}{
	"image/jpeg": {},
	"image/png":  {},
	"image/gif":  {},
	"image/webp": {},
}

// GetAvatarURL returns the stored avatar URL/data URI for external_id.
func GetAvatarURL(ctx context.Context, db *pgxpool.Pool, externalID string) (string, error) {
	if db == nil || strings.TrimSpace(externalID) == "" {
		return "", nil
	}
	var avatarURL *string
	err := db.QueryRow(ctx, `
		SELECT avatar_url FROM users WHERE external_id = $1
	`, externalID).Scan(&avatarURL)
	if errors.Is(err, pgx.ErrNoRows) {
		return "", nil
	}
	if err != nil {
		return "", err
	}
	if avatarURL == nil {
		return "", nil
	}
	return strings.TrimSpace(*avatarURL), nil
}

// SetAvatarURL validates and stores avatar_url for external_id.
func SetAvatarURL(ctx context.Context, db *pgxpool.Pool, externalID, avatarURL string) error {
	if db == nil {
		return fmt.Errorf("database not configured")
	}
	if strings.TrimSpace(externalID) == "" {
		return fmt.Errorf("missing external id")
	}
	normalized, err := normalizeAvatarURL(avatarURL)
	if err != nil {
		return err
	}
	tag, err := db.Exec(ctx, `
		UPDATE users
		SET avatar_url = $2, updated_at = NOW()
		WHERE external_id = $1
	`, externalID, normalized)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}

// ClearAvatarURL removes the stored avatar for external_id.
func ClearAvatarURL(ctx context.Context, db *pgxpool.Pool, externalID string) error {
	if db == nil {
		return fmt.Errorf("database not configured")
	}
	if strings.TrimSpace(externalID) == "" {
		return fmt.Errorf("missing external id")
	}
	tag, err := db.Exec(ctx, `
		UPDATE users
		SET avatar_url = NULL, updated_at = NOW()
		WHERE external_id = $1
	`, externalID)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}

func normalizeAvatarURL(raw string) (string, error) {
	trimmed := strings.TrimSpace(raw)
	if trimmed == "" {
		return "", ErrAvatarInvalid
	}
	if strings.HasPrefix(trimmed, "https://") || strings.HasPrefix(trimmed, "http://") {
		if len(trimmed) > 2048 {
			return "", ErrAvatarTooLarge
		}
		return trimmed, nil
	}
	if !strings.HasPrefix(trimmed, "data:") {
		return "", ErrAvatarInvalid
	}
	comma := strings.Index(trimmed, ",")
	if comma == -1 {
		return "", ErrAvatarInvalid
	}
	meta := trimmed[:comma]
	payload := strings.TrimSpace(trimmed[comma+1:])
	if !strings.Contains(meta, ";base64") {
		return "", ErrAvatarInvalid
	}
	mimePart := strings.TrimPrefix(meta, "data:")
	mimePart = strings.Split(mimePart, ";")[0]
	if _, ok := allowedAvatarMIME[strings.ToLower(mimePart)]; !ok {
		return "", ErrAvatarInvalid
	}
	decoded, err := base64.StdEncoding.DecodeString(payload)
	if err != nil {
		return "", ErrAvatarInvalid
	}
	if len(decoded) == 0 || len(decoded) > maxAvatarBytes {
		return "", ErrAvatarTooLarge
	}
	return trimmed, nil
}