byteclaw/internal/store/invites.go

package store

import (
	"context"
	"crypto/rand"
	"database/sql"
	"errors"
	"math/big"
	"strings"
	"time"
)

var (
	ErrInviteInvalid = errors.New("invalid invite code")
	ErrInviteUsed    = errors.New("invite code already used")
	ErrInviteExpired = errors.New("invite code expired")
)

// Invite is a registration code. ExpiresAt is zero when it never expires.
type Invite struct {
	Code      string     `json:"code"`
	MaxUses   int        `json:"max_uses"`
	Uses      int        `json:"uses"`
	ExpiresAt *time.Time `json:"expires_at,omitempty"`
	CreatedAt time.Time  `json:"created_at"`
	UsedBy    []string   `json:"used_by"`
}

// No 0/O/1/I/L to keep codes easy to read and type.
const inviteAlphabet = "23456789ABCDEFGHJKMNPQRSTUVWXYZ"

func newInviteCode() (string, error) {
	var b strings.Builder
	b.WriteString("BC-")
	for i := 0; i < 8; i++ {
		if i == 4 {
			b.WriteByte('-')
		}
		n, err := rand.Int(rand.Reader, big.NewInt(int64(len(inviteAlphabet))))
		if err != nil {
			return "", err
		}
		b.WriteByte(inviteAlphabet[n.Int64()])
	}
	return b.String(), nil
}

// NormalizeInviteCode upper-cases and trims user input ("bc-abcd-efgh " -> "BC-ABCD-EFGH").
func NormalizeInviteCode(c string) string { return strings.ToUpper(strings.TrimSpace(c)) }

// CreateInvite stores a new random code. maxUses < 1 means 1; ttl <= 0 means no expiry.
func (s *Store) CreateInvite(ctx context.Context, maxUses int, ttl time.Duration) (Invite, error) {
	if maxUses < 1 {
		maxUses = 1
	}
	now := time.Now().UTC()
	inv := Invite{MaxUses: maxUses, CreatedAt: now, UsedBy: []string{}}
	var exp int64
	if ttl > 0 {
		t := now.Add(ttl)
		inv.ExpiresAt = &t
		exp = t.UnixMilli()
	}
	for attempt := 0; attempt < 5; attempt++ {
		code, err := newInviteCode()
		if err != nil {
			return Invite{}, err
		}
		_, err = s.db.ExecContext(ctx, `INSERT INTO invites(code,max_uses,uses,expires_at,created_at) VALUES(?,?,0,?,?)`, code, maxUses, exp, now.UnixMilli())
		if err == nil {
			inv.Code = code
			return inv, nil
		}
		if !strings.Contains(err.Error(), "UNIQUE") {
			return Invite{}, err
		}
	}
	return Invite{}, errors.New("could not generate unique invite code")
}

// ListInvites returns all codes, newest first, with the usernames that used them.
func (s *Store) ListInvites(ctx context.Context) ([]Invite, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT code,max_uses,uses,expires_at,created_at FROM invites ORDER BY created_at DESC`)
	if err != nil {
		return nil, err
	}
	var out []Invite
	idx := map[string]int{}
	for rows.Next() {
		var inv Invite
		var exp, c int64
		if err := rows.Scan(&inv.Code, &inv.MaxUses, &inv.Uses, &exp, &c); err != nil {
			rows.Close()
			return nil, err
		}
		inv.CreatedAt = time.UnixMilli(c).UTC()
		if exp > 0 {
			t := time.UnixMilli(exp).UTC()
			inv.ExpiresAt = &t
		}
		inv.UsedBy = []string{}
		idx[inv.Code] = len(out)
		out = append(out, inv)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return nil, err
	}
	ur, err := s.db.QueryContext(ctx, `SELECT code,username FROM invite_uses ORDER BY used_at`)
	if err != nil {
		return nil, err
	}
	defer ur.Close()
	for ur.Next() {
		var code, name string
		if err := ur.Scan(&code, &name); err != nil {
			return nil, err
		}
		if i, ok := idx[code]; ok {
			out[i].UsedBy = append(out[i].UsedBy, name)
		}
	}
	if out == nil {
		out = []Invite{}
	}
	return out, ur.Err()
}

// DeleteInvite revokes a code. Users who already registered with it stay.
func (s *Store) DeleteInvite(ctx context.Context, code string) error {
	res, err := s.db.ExecContext(ctx, `DELETE FROM invites WHERE code=?`, NormalizeInviteCode(code))
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrNotFound
	}
	return nil
}

func consumeInvite(ctx context.Context, tx *sql.Tx, code string, now time.Time) error {
	var maxUses, uses int
	var exp int64
	err := tx.QueryRowContext(ctx, `SELECT max_uses,uses,expires_at FROM invites WHERE code=?`, code).Scan(&maxUses, &uses, &exp)
	if errors.Is(err, sql.ErrNoRows) {
		return ErrInviteInvalid
	}
	if err != nil {
		return err
	}
	if exp > 0 && now.UnixMilli() >= exp {
		return ErrInviteExpired
	}
	// Conditional update: never exceeds max_uses even under concurrency.
	res, err := tx.ExecContext(ctx, `UPDATE invites SET uses=uses+1 WHERE code=? AND uses<max_uses`, code)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrInviteUsed
	}
	return nil
}