From eb83c202bad5b920ed4612102369f7d28b0a23fb Mon Sep 17 00:00:00 2001 From: Anirudh Oppiliappan Date: Wed, 29 Apr 2026 10:25:06 +0300 Subject: [PATCH] knotserver/db: add owner UID assignment and repo isolation --- knotserver/db/db.go | 51 +++++++++++++--- knotserver/db/owner_uid.go | 102 ++++++++++++++++++++++++++++++++ knotserver/db/owner_uid_test.go | 91 ++++++++++++++++++++++++++++ 3 files changed, 237 insertions(+), 7 deletions(-) create mode 100644 knotserver/db/owner_uid.go create mode 100644 knotserver/db/owner_uid_test.go diff --git a/knotserver/db/db.go b/knotserver/db/db.go index b02edaf6..3687ea42 100644 --- a/knotserver/db/db.go +++ b/knotserver/db/db.go @@ -7,6 +7,7 @@ import ( "log/slog" "os" "strings" + "sync" securejoin "github.com/cyphar/filepath-securejoin" _ "github.com/mattn/go-sqlite3" @@ -17,6 +18,10 @@ import ( type DB struct { db *sql.DB logger *slog.Logger + + // uidAssignMu serialises GetOrAssignOwnerUID across goroutines so that + // concurrent callers don't race on the uid_counter read-modify-write. + uidAssignMu sync.Mutex } type DBTX interface { @@ -193,14 +198,46 @@ func Setup(ctx context.Context, dbPath string) (*DB, error) { if err := orm.RunMigration(conn, logger, "create-knot-members", func(tx *sql.Tx) error { _, mErr := tx.ExecContext(ctx, ` - create table if not exists knot_members ( - id integer primary key autoincrement, - did text not null, - rkey text not null, - subject text not null, - created text not null default (strftime('%Y-%m-%dT%H:%M:%SZ', 'now')), - unique (did, rkey) + create table if not exists knot_members ( + id integer primary key autoincrement, + did text not null, + rkey text not null, + subject text not null, + created text not null default (strftime('%Y-%m-%dT%H:%M:%SZ', 'now')), + unique (did, rkey) + ); + `) + return mErr + }); err != nil { + return nil, err + } + + if err := orm.RunMigration(conn, logger, "add-isolated-at-to-repo-keys", func(tx *sql.Tx) error { + _, mErr := tx.ExecContext(ctx, `ALTER TABLE repo_keys ADD COLUMN isolated_at DATETIME`) + return mErr + }); err != nil { + return nil, err + } + + if err := orm.RunMigration(conn, logger, "add-owner-uid-tables", func(tx *sql.Tx) error { + _, mErr := tx.ExecContext(ctx, ` + CREATE TABLE IF NOT EXISTS owner_uid_assignments ( + owner_did TEXT PRIMARY KEY, + uid INTEGER NOT NULL UNIQUE, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP ); + + CREATE TABLE IF NOT EXISTS uid_counter ( + next_uid INTEGER NOT NULL DEFAULT 100000 + ); + `) + if mErr != nil { + return mErr + } + // Seed the counter only if the table is empty. + _, mErr = tx.ExecContext(ctx, ` + INSERT INTO uid_counter (next_uid) + SELECT 100000 WHERE NOT EXISTS (SELECT 1 FROM uid_counter) `) return mErr }); err != nil { diff --git a/knotserver/db/owner_uid.go b/knotserver/db/owner_uid.go new file mode 100644 index 00000000..73e549f2 --- /dev/null +++ b/knotserver/db/owner_uid.go @@ -0,0 +1,102 @@ +package db + +import ( + "database/sql" +) + +// GetOrAssignOwnerUID returns the virtual UID for ownerDID, minting a new one +// from the uid_counter table if this owner has not been seen before. +// UIDs start at 100000 and increment by one per unique owner. +// +// A process-wide mutex serialises concurrent callers so two simultaneous +// requests for distinct DIDs do not race to claim the same counter value. +// The mutex is local to this DB instance and SQLite itself only permits one +// writer at a time, so this does not impede throughput in practice. +func (d *DB) GetOrAssignOwnerUID(ownerDID string) (uint32, error) { + d.uidAssignMu.Lock() + defer d.uidAssignMu.Unlock() + + tx, err := d.db.Begin() + if err != nil { + return 0, err + } + defer tx.Rollback() + + var uid uint32 + err = tx.QueryRow( + `SELECT uid FROM owner_uid_assignments WHERE owner_did = ?`, + ownerDID, + ).Scan(&uid) + if err == nil { + return uid, tx.Commit() + } + if err != sql.ErrNoRows { + return 0, err + } + + if err := tx.QueryRow(`SELECT next_uid FROM uid_counter`).Scan(&uid); err != nil { + return 0, err + } + if _, err := tx.Exec(`UPDATE uid_counter SET next_uid = next_uid + 1`); err != nil { + return 0, err + } + if _, err := tx.Exec( + `INSERT INTO owner_uid_assignments (owner_did, uid) VALUES (?, ?)`, + ownerDID, uid, + ); err != nil { + return 0, err + } + + return uid, tx.Commit() +} + +// AllReposForMigration returns all (repo_did, owner_did) pairs with a +// non-null owner. Pass force=true to include already-migrated repos. +func (d *DB) AllReposForMigration(force bool) ([]RepoMigrationRow, error) { + query := `SELECT repo_did, owner_did FROM repo_keys WHERE owner_did IS NOT NULL` + if !force { + query += ` AND isolated_at IS NULL` + } + rows, err := d.db.Query(query) + if err != nil { + return nil, err + } + defer rows.Close() + + var result []RepoMigrationRow + for rows.Next() { + var r RepoMigrationRow + if err := rows.Scan(&r.RepoDID, &r.OwnerDID); err != nil { + return nil, err + } + result = append(result, r) + } + return result, rows.Err() +} + +// CountUnmigratedRepos returns the number of repos that have not yet been +// isolation-migrated. +func (d *DB) CountUnmigratedRepos() (int, error) { + var n int + err := d.db.QueryRow(` + SELECT count(1) FROM repo_keys + WHERE owner_did IS NOT NULL + AND isolated_at IS NULL + `).Scan(&n) + return n, err +} + +// MarkRepoIsolated sets isolated_at to the current time for repoDID. +func (d *DB) MarkRepoIsolated(repoDID string) error { + _, err := d.db.Exec( + `UPDATE repo_keys SET isolated_at = strftime('%Y-%m-%dT%H:%M:%SZ','now') WHERE repo_did = ?`, + repoDID, + ) + return err +} + +// RepoMigrationRow is a row returned by AllReposForMigration. +type RepoMigrationRow struct { + RepoDID string + OwnerDID string +} diff --git a/knotserver/db/owner_uid_test.go b/knotserver/db/owner_uid_test.go new file mode 100644 index 00000000..389d49a0 --- /dev/null +++ b/knotserver/db/owner_uid_test.go @@ -0,0 +1,91 @@ +package db + +import ( + "context" + "fmt" + "path/filepath" + "sync" + "testing" +) + +// TestGetOrAssignOwnerUID_Concurrent verifies that concurrent calls for +// distinct DIDs each get a unique UID, with no duplicates from a race on +// the uid_counter table. +func TestGetOrAssignOwnerUID_Concurrent(t *testing.T) { + d := newTestDB(t) + + const n = 50 + results := make([]uint32, n) + var wg sync.WaitGroup + wg.Add(n) + for i := 0; i < n; i++ { + go func(i int) { + defer wg.Done() + uid, err := d.GetOrAssignOwnerUID(fmt.Sprintf("did:plc:test-%03d", i)) + if err != nil { + t.Errorf("GetOrAssignOwnerUID(%d): %v", i, err) + return + } + results[i] = uid + }(i) + } + wg.Wait() + + seen := map[uint32]int{} + for i, uid := range results { + if uid == 0 { + t.Errorf("result[%d] is zero (probably an error above)", i) + continue + } + if prev, ok := seen[uid]; ok { + t.Errorf("uid %d was returned for both index %d and %d", uid, prev, i) + } + seen[uid] = i + } + + if len(seen) != n { + t.Errorf("got %d unique UIDs, want %d", len(seen), n) + } +} + +// TestGetOrAssignOwnerUID_Idempotent verifies that repeated calls for the +// same DID return the same UID. +func TestGetOrAssignOwnerUID_Idempotent(t *testing.T) { + d := newTestDB(t) + const did = "did:plc:stable" + + first, err := d.GetOrAssignOwnerUID(did) + if err != nil { + t.Fatalf("first call: %v", err) + } + for i := 0; i < 5; i++ { + uid, err := d.GetOrAssignOwnerUID(did) + if err != nil { + t.Fatalf("repeat call %d: %v", i, err) + } + if uid != first { + t.Errorf("repeat call %d returned %d, want %d", i, uid, first) + } + } +} + +// TestGetOrAssignOwnerUID_StartsAt100000 verifies the counter seeds correctly. +func TestGetOrAssignOwnerUID_StartsAt100000(t *testing.T) { + d := newTestDB(t) + uid, err := d.GetOrAssignOwnerUID("did:plc:first") + if err != nil { + t.Fatalf("first call: %v", err) + } + if uid != 100000 { + t.Errorf("first UID = %d, want 100000", uid) + } +} + +func newTestDB(t *testing.T) *DB { + t.Helper() + d, err := Setup(context.Background(), filepath.Join(t.TempDir(), "test.db")) + if err != nil { + t.Fatalf("db.Setup: %v", err) + } + return d +} -- 2.51.2