| 1 | package lease |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "database/sql" |
| 6 | "errors" |
| 7 | "os" |
| 8 | "path/filepath" |
| 9 | "testing" |
| 10 | |
| 11 | _ "modernc.org/sqlite" |
| 12 | ) |
| 13 | |
| 14 | // TestCheckDB pins the three verdicts a restored file can earn. Ported from |
| 15 | // ops.TestCheckRestoredDBRejectsEmptySchema, which now runs against this |
| 16 | // implementation through the delegating wrapper. |
| 17 | func TestCheckDB(t *testing.T) { |
| 18 | dir := t.TempDir() |
| 19 | |
| 20 | good := filepath.Join(dir, "good.db") |
| 21 | db, err := sql.Open("sqlite", good) |
| 22 | if err != nil { |
| 23 | t.Fatal(err) |
| 24 | } |
| 25 | if _, err := db.Exec("CREATE TABLE t (id INTEGER PRIMARY KEY); INSERT INTO t VALUES (1)"); err != nil { |
| 26 | t.Fatal(err) |
| 27 | } |
| 28 | db.Close() |
| 29 | if err := CheckDB(good); err != nil { |
| 30 | t.Fatalf("CheckDB(good) = %v, want nil", err) |
| 31 | } |
| 32 | |
| 33 | // A schema-less database restores "successfully" but would be useless in a |
| 34 | // disaster — the check must refuse it. |
| 35 | empty := filepath.Join(dir, "empty.db") |
| 36 | db, err = sql.Open("sqlite", empty) |
| 37 | if err != nil { |
| 38 | t.Fatal(err) |
| 39 | } |
| 40 | if _, err := db.Exec("PRAGMA user_version=1"); err != nil { // force file creation |
| 41 | t.Fatal(err) |
| 42 | } |
| 43 | db.Close() |
| 44 | if err := CheckDB(empty); err == nil { |
| 45 | t.Fatal("CheckDB accepted a database with no tables") |
| 46 | } |
| 47 | |
| 48 | // Garbage bytes must fail loudly, not pass as an open-able file. |
| 49 | junk := filepath.Join(dir, "junk.db") |
| 50 | if err := os.WriteFile(junk, []byte("this is not a sqlite database at all"), 0o644); err != nil { |
| 51 | t.Fatal(err) |
| 52 | } |
| 53 | if err := CheckDB(junk); err == nil { |
| 54 | t.Fatal("CheckDB accepted a non-SQLite file") |
| 55 | } |
| 56 | |
| 57 | // A file that is not there at all is a failure, never a silent pass. |
| 58 | if err := CheckDB(filepath.Join(dir, "absent.db")); err == nil { |
| 59 | t.Fatal("CheckDB accepted a missing file") |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | // TestCheckDBDetectsPageCorruption is the case the whole verification pass |
| 64 | // exists for: a file that opens fine and has a schema, but whose pages are |
| 65 | // damaged. integrity_check is the only thing that catches it. |
| 66 | func TestCheckDBDetectsPageCorruption(t *testing.T) { |
| 67 | dir := t.TempDir() |
| 68 | path := filepath.Join(dir, "corrupt.db") |
| 69 | |
| 70 | db, err := sql.Open("sqlite", path) |
| 71 | if err != nil { |
| 72 | t.Fatal(err) |
| 73 | } |
| 74 | // No WAL: the damage below must land in the main database file, not in a |
| 75 | // sidecar the reopen would replay over. |
| 76 | if _, err := db.Exec("PRAGMA journal_mode=DELETE"); err != nil { |
| 77 | t.Fatal(err) |
| 78 | } |
| 79 | if _, err := db.Exec("CREATE TABLE t (id INTEGER PRIMARY KEY, body TEXT)"); err != nil { |
| 80 | t.Fatal(err) |
| 81 | } |
| 82 | for i := 0; i < 2000; i++ { |
| 83 | if _, err := db.Exec("INSERT INTO t (body) VALUES (?)", "row padding to force several pages"); err != nil { |
| 84 | t.Fatal(err) |
| 85 | } |
| 86 | } |
| 87 | db.Close() |
| 88 | |
| 89 | if fi, err := os.Stat(path); err != nil { |
| 90 | t.Fatal(err) |
| 91 | } else if fi.Size() < 3*4096 { |
| 92 | t.Fatalf("test fixture is only %d bytes — too few pages to corrupt one meaningfully", fi.Size()) |
| 93 | } |
| 94 | |
| 95 | if err := CheckDB(path); err != nil { |
| 96 | t.Fatalf("CheckDB on the intact file = %v, want nil", err) |
| 97 | } |
| 98 | |
| 99 | // Scribble over page 2's header — its page type and cell-pointer array — |
| 100 | // leaving the file length and page 1 (the schema) intact so the open still |
| 101 | // succeeds and only integrity_check can tell. |
| 102 | f, err := os.OpenFile(path, os.O_WRONLY, 0o644) |
| 103 | if err != nil { |
| 104 | t.Fatal(err) |
| 105 | } |
| 106 | junk := make([]byte, 512) |
| 107 | for i := range junk { |
| 108 | junk[i] = 0xA5 |
| 109 | } |
| 110 | if _, err := f.WriteAt(junk, 4096); err != nil { |
| 111 | t.Fatal(err) |
| 112 | } |
| 113 | f.Close() |
| 114 | |
| 115 | if err := CheckDB(path); err == nil { |
| 116 | t.Fatal("CheckDB passed a file with a scribbled-over page") |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | // The restored file is WAL-journaled — that is what instances write and what |
| 121 | // litestream restores — so the check must pass one. This is the shape both |
| 122 | // callers actually hand it, and the box-wide drill has been opening exactly |
| 123 | // this way in production daily. |
| 124 | // |
| 125 | // It also pins what the "mode=ro" in the DSN does NOT buy: the driver still |
| 126 | // attaches WAL sidecars. That is harmless at both call sites and neither may |
| 127 | // come to depend on the opposite — the wake path calls this holding the |
| 128 | // lease, on a file it has just restored with the sidecars already cleared |
| 129 | // (nothing else has the database open), and VerifyReplica works on a |
| 130 | // throwaway copy it deletes, sidecars included. |
| 131 | func TestCheckDBAcceptsWALDatabase(t *testing.T) { |
| 132 | dir := t.TempDir() |
| 133 | path := filepath.Join(dir, "wal.db") |
| 134 | seedDB(t, path) // WAL mode, 50 rows — the lease package's own fixture |
| 135 | |
| 136 | // litestream restores a checkpointed file with no live -wal beside it; |
| 137 | // mimic that so this tests the shape the wake path actually sees. |
| 138 | for _, p := range []string{path + "-wal", path + "-shm"} { |
| 139 | _ = os.Remove(p) |
| 140 | } |
| 141 | if err := CheckDB(path); err != nil { |
| 142 | t.Fatalf("CheckDB on a WAL-mode restore = %v, want nil", err) |
| 143 | } |
| 144 | } |
| 145 | |
| 146 | // VerifyReplica must leave nothing behind — including the WAL sidecars the |
| 147 | // check attaches. A drill that littered the box with restored copies of every |
| 148 | // tenant database would be a disk-usage bug and a data-exposure one. |
| 149 | func TestVerifyReplicaCleansUpEvenOnFailure(t *testing.T) { |
| 150 | dir := t.TempDir() |
| 151 | |
| 152 | // A restore that "succeeds" into a corrupt file: the check fails, and the |
| 153 | // wreckage must still be swept. |
| 154 | prev := restoreForVerify |
| 155 | restoreForVerify = func(_ context.Context, out, _, _, _ string) (bool, error) { |
| 156 | if err := os.WriteFile(out, []byte("not a database"), 0o600); err != nil { |
| 157 | return false, err |
| 158 | } |
| 159 | return true, nil |
| 160 | } |
| 161 | t.Cleanup(func() { restoreForVerify = prev }) |
| 162 | |
| 163 | if _, err := VerifyReplica(context.Background(), "bucket", "hib/x.example.com", "", dir); err == nil { |
| 164 | t.Fatal("VerifyReplica accepted a corrupt restore") |
| 165 | } |
| 166 | left, err := os.ReadDir(dir) |
| 167 | if err != nil { |
| 168 | t.Fatal(err) |
| 169 | } |
| 170 | if len(left) != 0 { |
| 171 | names := []string{} |
| 172 | for _, e := range left { |
| 173 | names = append(names, e.Name()) |
| 174 | } |
| 175 | t.Fatalf("VerifyReplica left %v behind", names) |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | // A prefix that has never been replicated to is ErrNoReplica — "nothing to |
| 180 | // restore yet", which callers must not report as a failure. |
| 181 | func TestVerifyReplicaReportsNoReplica(t *testing.T) { |
| 182 | prev := restoreForVerify |
| 183 | restoreForVerify = func(context.Context, string, string, string, string) (bool, error) { |
| 184 | return false, nil |
| 185 | } |
| 186 | t.Cleanup(func() { restoreForVerify = prev }) |
| 187 | |
| 188 | _, err := VerifyReplica(context.Background(), "bucket", "hib/new.example.com", "", t.TempDir()) |
| 189 | if !errors.Is(err, ErrNoReplica) { |
| 190 | t.Fatalf("VerifyReplica on a never-replicated prefix = %v, want ErrNoReplica", err) |
| 191 | } |
| 192 | } |
| 193 | |
| 194 | // VerifyReplica hands back the restored copy's change counter, which is what |
| 195 | // lets the hibernation drill ask the second question — is this replica still |
| 196 | // TRACKING its source? — after the throwaway copy has been swept. |
| 197 | func TestVerifyReplicaReturnsTheRestoredChangeCounter(t *testing.T) { |
| 198 | dir := t.TempDir() |
| 199 | src := filepath.Join(t.TempDir(), "source.db") |
| 200 | seedDB(t, src) |
| 201 | want, err := ChangeCounter(src) |
| 202 | if err != nil { |
| 203 | t.Fatal(err) |
| 204 | } |
| 205 | |
| 206 | prev := restoreForVerify |
| 207 | restoreForVerify = func(_ context.Context, out, _, _, _ string) (bool, error) { |
| 208 | body, err := os.ReadFile(src) |
| 209 | if err != nil { |
| 210 | return false, err |
| 211 | } |
| 212 | return true, os.WriteFile(out, body, 0o600) |
| 213 | } |
| 214 | t.Cleanup(func() { restoreForVerify = prev }) |
| 215 | |
| 216 | got, err := VerifyReplica(context.Background(), "bucket", "hib/x.example.com", "", dir) |
| 217 | if err != nil { |
| 218 | t.Fatalf("VerifyReplica: %v", err) |
| 219 | } |
| 220 | if got != want { |
| 221 | t.Fatalf("restored change counter = %d, want %d", got, want) |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | // A stop reaches the integrity scan: CheckDBContext returns the context's |
| 226 | // error instead of reading the whole file first. |
| 227 | func TestCheckDBContextHonoursCancellation(t *testing.T) { |
| 228 | path := filepath.Join(t.TempDir(), "c.db") |
| 229 | seedDB(t, path) |
| 230 | if err := CheckDBContext(context.Background(), path); err != nil { |
| 231 | t.Fatalf("live context: %v", err) |
| 232 | } |
| 233 | ctx, cancel := context.WithCancel(context.Background()) |
| 234 | cancel() |
| 235 | if err := CheckDBContext(ctx, path); !errors.Is(err, context.Canceled) { |
| 236 | t.Fatalf("cancelled context = %v, want context.Canceled", err) |
| 237 | } |
| 238 | } |
| 239 | |
amadan