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") } }) } }