Warp/server/internal/storage/admin_test.go

package storage_test

import (
	"database/sql"
	"errors"
	"os"
	"testing"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// migratedAdminDB opens a migrated database with empty Phase 9 tables.
func migratedAdminDB(t *testing.T, driver string) storage.Database {
	t.Helper()
	db := openTestDB(t, driver)
	if err := db.Migrate(t.Context()); err != nil {
		t.Fatalf("Migrate: %v", err)
	}
	if driver == storage.DriverPostgres {
		pool, err := sql.Open("pgx", os.Getenv("WARP_TEST_POSTGRES_DSN"))
		if err != nil {
			t.Fatalf("open postgres to reset the admin tables: %v", err)
		}
		defer func() { _ = pool.Close() }()
		for _, stmt := range []string{
			`TRUNCATE admin_registration_tokens, admin_server_notice_rooms, admin_audit_log,
				media_quotas, appservice_state`,
			`UPDATE admin_sequences SET value = 0`,
		} {
			if _, err := pool.ExecContext(t.Context(), stmt); err != nil {
				t.Fatalf("reset the admin tables: %v", err)
			}
		}
	}
	return db
}

func TestRegistrationTokens(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			admin := migratedAdminDB(t, driver).Admin()

			two := 2
			if err := admin.CreateRegistrationToken(ctx, storage.RegistrationToken{
				Token: "abc", UsesAllowed: &two, CreatedAtMS: 1,
			}); err != nil {
				t.Fatal(err)
			}
			if err := admin.CreateRegistrationToken(ctx, storage.RegistrationToken{Token: "abc", CreatedAtMS: 1}); !errors.Is(err, storage.ErrRegistrationTokenExists) {
				t.Errorf("duplicate token: %v", err)
			}
			if err := admin.CreateRegistrationToken(ctx, storage.RegistrationToken{
				Token: "open", CreatedAtMS: 1, ExpiresAtMS: 1000,
			}); err != nil {
				t.Fatal(err)
			}

			// Two uses, then no more; giving one back makes room again.
			for i := range 2 {
				if err := admin.UseRegistrationToken(ctx, "abc", 10); err != nil {
					t.Fatalf("use %d: %v", i, err)
				}
			}
			if err := admin.UseRegistrationToken(ctx, "abc", 10); !errors.Is(err, storage.ErrNoSuchRegistrationToken) {
				t.Errorf("third use: %v", err)
			}
			if err := admin.ReleaseRegistrationToken(ctx, "abc"); err != nil {
				t.Fatal(err)
			}
			got, err := admin.RegistrationToken(ctx, "abc")
			if err != nil || got.Completed != 1 || got.UsesAllowed == nil || *got.UsesAllowed != 2 || !got.Usable(10) {
				t.Errorf("after release = %+v, %v", got, err)
			}

			// An unlimited token stops at its expiry.
			if err := admin.UseRegistrationToken(ctx, "open", 999); err != nil {
				t.Errorf("before expiry: %v", err)
			}
			if err := admin.UseRegistrationToken(ctx, "open", 1000); !errors.Is(err, storage.ErrNoSuchRegistrationToken) {
				t.Errorf("at expiry: %v", err)
			}

			// Updating replaces the limit and the expiry and keeps the count.
			if err := admin.UpdateRegistrationToken(ctx, "open", &two, 0); err != nil {
				t.Fatal(err)
			}
			got, _ = admin.RegistrationToken(ctx, "open")
			if got.Completed != 1 || got.ExpiresAtMS != 0 || *got.UsesAllowed != 2 {
				t.Errorf("after update = %+v", got)
			}
			if err := admin.UpdateRegistrationToken(ctx, "nope", nil, 0); !errors.Is(err, storage.ErrNoSuchRegistrationToken) {
				t.Errorf("update unknown: %v", err)
			}

			all, err := admin.RegistrationTokens(ctx)
			if err != nil || len(all) != 2 || all[0].Token != "abc" || all[1].Token != "open" {
				t.Errorf("RegistrationTokens = %+v, %v", all, err)
			}
			if err := admin.DeleteRegistrationToken(ctx, "abc"); err != nil {
				t.Fatal(err)
			}
			if err := admin.DeleteRegistrationToken(ctx, "abc"); !errors.Is(err, storage.ErrNoSuchRegistrationToken) {
				t.Errorf("second delete: %v", err)
			}
			if _, err := admin.RegistrationToken(ctx, "abc"); !errors.Is(err, storage.ErrNoSuchRegistrationToken) {
				t.Errorf("read deleted: %v", err)
			}
		})
	}
}

func TestAuditLog(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			admin := migratedAdminDB(t, driver).Admin()

			for i, e := range []storage.AuditEntry{
				{TSMS: 1, Actor: "@root:hs", Action: "user.create", Target: "@a:hs"},
				{TSMS: 2, Actor: "@root:hs", Action: "user.deactivate", Target: "@a:hs", Details: []byte(`{"erase":true}`)},
				{TSMS: 3, Actor: "@ops:hs", Action: "user.create", Target: "@b:hs"},
			} {
				id, err := admin.AppendAudit(ctx, e)
				if err != nil || id != int64(i+1) {
					t.Fatalf("AppendAudit %d = %d, %v", i, id, err)
				}
			}

			all, err := admin.AuditLog(ctx, storage.AuditFilter{})
			if err != nil || len(all) != 3 || all[0].ID != 3 || all[2].ID != 1 {
				t.Fatalf("AuditLog = %+v, %v", all, err)
			}
			if string(all[1].Details) != `{"erase":true}` || string(all[0].Details) != `{}` {
				t.Errorf("details = %s, %s", all[1].Details, all[0].Details)
			}

			page, _ := admin.AuditLog(ctx, storage.AuditFilter{Limit: 1, Before: 3})
			if len(page) != 1 || page[0].ID != 2 {
				t.Errorf("paged = %+v", page)
			}
			byActor, _ := admin.AuditLog(ctx, storage.AuditFilter{Actor: "@root:hs"})
			byAction, _ := admin.AuditLog(ctx, storage.AuditFilter{Action: "user.create"})
			if len(byActor) != 2 || len(byAction) != 2 || byAction[0].Target != "@b:hs" {
				t.Errorf("filtered = %+v / %+v", byActor, byAction)
			}
		})
	}
}

func TestServerNoticeRooms(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			admin := migratedAdminDB(t, driver).Admin()

			if _, err := admin.ServerNoticeRoom(ctx, "alice"); !errors.Is(err, storage.ErrNoServerNoticeRoom) {
				t.Errorf("unset: %v", err)
			}
			for _, room := range []string{"!one:hs", "!two:hs"} {
				if err := admin.SetServerNoticeRoom(ctx, "alice", room); err != nil {
					t.Fatal(err)
				}
			}
			if room, err := admin.ServerNoticeRoom(ctx, "alice"); err != nil || room != "!two:hs" {
				t.Errorf("ServerNoticeRoom = %q, %v", room, err)
			}
		})
	}
}

func TestAppServiceState(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			apps := migratedAdminDB(t, driver).AppServices()

			st, err := apps.AppServiceState(ctx, "irc", 42)
			if err != nil || st.StreamPos != 42 || st.TxnTo != 0 {
				t.Fatalf("first read = %+v, %v", st, err)
			}
			// The start position only applies the first time.
			if err := apps.BeginTransaction(ctx, "irc", 1, 50); err != nil {
				t.Fatal(err)
			}
			st, _ = apps.AppServiceState(ctx, "irc", 99)
			if st.StreamPos != 42 || st.TxnID != 1 || st.TxnTo != 50 {
				t.Errorf("in flight = %+v", st)
			}
			if err := apps.CompleteTransaction(ctx, "irc", 50); err != nil {
				t.Fatal(err)
			}
			st, _ = apps.AppServiceState(ctx, "irc", 0)
			if st.StreamPos != 50 || st.TxnID != 1 || st.TxnTo != 0 {
				t.Errorf("completed = %+v", st)
			}
		})
	}
}

func TestMediaUsageAndQuota(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			db := migratedAdminDB(t, driver)
			if driver == storage.DriverPostgres {
				pool, err := sql.Open("pgx", os.Getenv("WARP_TEST_POSTGRES_DSN"))
				if err != nil {
					t.Fatal(err)
				}
				defer func() { _ = pool.Close() }()
				if _, err := pool.ExecContext(ctx, `TRUNCATE media_repository`); err != nil {
					t.Fatal(err)
				}
			}
			media := db.Media()

			for _, m := range []storage.Media{
				{Origin: "hs", MediaID: "a", Size: 100, SHA256: "x", UserID: "@alice:hs", CreatedAtMS: 1},
				{Origin: "hs", MediaID: "b", Size: 100, SHA256: "x", UserID: "@alice:hs", CreatedAtMS: 1},
				{Origin: "hs", MediaID: "c", UserID: "@alice:hs", CreatedAtMS: 1, Pending: true, ExpiresAtMS: 10},
				{Origin: "hs", MediaID: "d", Size: 7, SHA256: "y", UserID: "@bob:hs", CreatedAtMS: 1},
				{Origin: "remote", MediaID: "e", Size: 5000, SHA256: "z", CreatedAtMS: 1},
			} {
				if err := media.CreateMedia(ctx, m); err != nil {
					t.Fatal(err)
				}
			}
			// The same bytes twice count twice, and a reservation counts
			// for nothing.
			usage, err := media.MediaUsage(ctx, "hs", "@alice:hs")
			if err != nil || usage.Count != 2 || usage.Bytes != 200 {
				t.Errorf("alice = %+v, %v", usage, err)
			}
			if usage, _ := media.MediaUsage(ctx, "hs", "@carol:hs"); usage.Count != 0 || usage.Bytes != 0 {
				t.Errorf("carol = %+v", usage)
			}

			if _, ok, err := media.MediaQuota(ctx, "@alice:hs"); ok || err != nil {
				t.Errorf("unset quota: %v, %v", ok, err)
			}
			for _, q := range []int64{1000, 500} {
				if err := media.SetMediaQuota(ctx, "@alice:hs", q); err != nil {
					t.Fatal(err)
				}
			}
			if q, ok, err := media.MediaQuota(ctx, "@alice:hs"); !ok || q != 500 || err != nil {
				t.Errorf("quota = %d, %v, %v", q, ok, err)
			}
			if err := media.DeleteMediaQuota(ctx, "@alice:hs"); err != nil {
				t.Fatal(err)
			}
			if _, ok, _ := media.MediaQuota(ctx, "@alice:hs"); ok {
				t.Error("quota survived its deletion")
			}
		})
	}
}