byteclaw/internal/server/invites_test.go

package server

import (
	"context"
	"encoding/json"
	"io"
	"regexp"
	"strings"
	"sync"
	"testing"
	"time"
)

func regBody(u, p, code string) string {
	b, _ := json.Marshal(map[string]string{"username": u, "password": p, "invite_code": code})
	return string(b)
}


func TestInviteRegistration(t *testing.T) {
	ts, st, s := authSetup(t, true)
	s.RegistrationMode = "invite"
	adm := bearerHdr(tok)

	var cfg map[string]string
	json.NewDecoder(do(t, "GET", ts.URL+"/api/auth/config", "", nil).Body).Decode(&cfg)
	if cfg["registration_mode"] != "invite" {
		t.Fatalf("%v", cfg)
	}
	if r := do(t, "POST", ts.URL+"/api/auth/register", creds("alice", "password-1"), nil); r.StatusCode != 400 {
		t.Fatal("no code must be 400", r.StatusCode)
	}
	if r := do(t, "POST", ts.URL+"/api/auth/register", regBody("alice", "password-1", "BC-NOPE-NOPE"), nil); r.StatusCode != 403 {
		t.Fatal("unknown code must be 403", r.StatusCode)
	}
	// only admin may manage codes
	ut := ""
	{
		s.RegistrationMode = "open"
		ut = registerUser(t, ts.URL, "bob", "password-1")
		s.RegistrationMode = "invite"
	}
	for _, m := range []string{"GET", "POST"} {
		if r := do(t, m, ts.URL+"/api/invites", "", bearerHdr(ut)); r.StatusCode != 403 {
			t.Fatal(m, "invites must be admin only", r.StatusCode)
		}
	}
	r := do(t, "POST", ts.URL+"/api/invites", "", adm)
	if r.StatusCode != 201 {
		b, _ := io.ReadAll(r.Body)
		t.Fatal(r.StatusCode, string(b))
	}
	var inv struct {
		Code    string `json:"code"`
		MaxUses int    `json:"max_uses"`
	}
	json.NewDecoder(r.Body).Decode(&inv)
	if !regexp.MustCompile(`^BC-[2-9A-HJKMNP-Z]{4}-[2-9A-HJKMNP-Z]{4}$`).MatchString(inv.Code) || inv.MaxUses != 1 {
		t.Fatalf("%+v", inv)
	}
	// lower-case input is accepted
	if r := do(t, "POST", ts.URL+"/api/auth/register", regBody("alice", "password-1", " "+strings.ToLower(inv.Code)), nil); r.StatusCode != 201 {
		b, _ := io.ReadAll(r.Body)
		t.Fatal("register with code", r.StatusCode, string(b))
	}
	r = do(t, "POST", ts.URL+"/api/auth/register", regBody("carol", "password-1", inv.Code), nil)
	b, _ := io.ReadAll(r.Body)
	if r.StatusCode != 403 || !strings.Contains(string(b), "already used") {
		t.Fatal("reuse", r.StatusCode, string(b))
	}
	// list shows usage
	var list []struct {
		Code   string   `json:"code"`
		Uses   int      `json:"uses"`
		UsedBy []string `json:"used_by"`
	}
	json.NewDecoder(do(t, "GET", ts.URL+"/api/invites", "", adm).Body).Decode(&list)
	if len(list) != 1 || list[0].Uses != 1 || len(list[0].UsedBy) != 1 || list[0].UsedBy[0] != "alice" {
		t.Fatalf("%+v", list)
	}
	// multi-use + revoke
	r = do(t, "POST", ts.URL+"/api/invites", `{"max_uses":5,"expires_in_hours":24}`, adm)
	json.NewDecoder(r.Body).Decode(&inv)
	if inv.MaxUses != 5 {
		t.Fatalf("%+v", inv)
	}
	if r := do(t, "DELETE", ts.URL+"/api/invites/"+inv.Code, "", adm); r.StatusCode != 204 {
		t.Fatal("revoke", r.StatusCode)
	}
	if r := do(t, "POST", ts.URL+"/api/auth/register", regBody("dave", "password-1", inv.Code), nil); r.StatusCode != 403 {
		t.Fatal("revoked code", r.StatusCode)
	}
	if r := do(t, "DELETE", ts.URL+"/api/invites/"+inv.Code, "", adm); r.StatusCode != 404 {
		t.Fatal("revoke twice", r.StatusCode)
	}
	// expired
	ex, _ := st.CreateInvite(context.Background(), 1, time.Millisecond)
	time.Sleep(5 * time.Millisecond)
	r = do(t, "POST", ts.URL+"/api/auth/register", regBody("erin", "password-1", ex.Code), nil)
	b, _ = io.ReadAll(r.Body)
	if r.StatusCode != 403 || !strings.Contains(string(b), "expired") {
		t.Fatal("expired", r.StatusCode, string(b))
	}
	// closed mode
	s.RegistrationMode = "closed"
	if r := do(t, "POST", ts.URL+"/api/auth/register", regBody("frank", "password-1", ""), nil); r.StatusCode != 403 {
		t.Fatal("closed", r.StatusCode)
	}
}

func TestInviteConsumedAtomically(t *testing.T) {
	_, st, _ := authSetup(t, true)
	ctx := context.Background()
	inv, err := st.CreateInvite(ctx, 3, 0)
	if err != nil {
		t.Fatal(err)
	}
	var wg sync.WaitGroup
	var mu sync.Mutex
	ok := 0
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			if _, err := st.CreateUserWithInvite(ctx, "u"+string(rune('a'+i))+"xx", "h", inv.Code); err == nil {
				mu.Lock()
				ok++
				mu.Unlock()
			}
		}(i)
	}
	wg.Wait()
	if ok != 3 {
		t.Fatalf("expected exactly 3 registrations, got %d", ok)
	}
	// failed user insert must not burn a use
	inv2, _ := st.CreateInvite(ctx, 1, 0)
	st.CreateUser(ctx, "dupe", "h")
	if _, err := st.CreateUserWithInvite(ctx, "dupe", "h", inv2.Code); err == nil {
		t.Fatal("duplicate username must fail")
	}
	if _, err := st.CreateUserWithInvite(ctx, "fresh", "h", inv2.Code); err != nil {
		t.Fatal("code must still be usable after failed insert:", err)
	}
}