diff --git a/backup/restore.go b/backup/restore.go new file mode 100644 index 0000000..7cc83a9 --- /dev/null +++ b/backup/restore.go @@ -0,0 +1,339 @@ +package backup + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "sort" + "strings" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" + "go.mongodb.org/mongo-driver/v2/mongo/options" +) + +// BatchSize is how many documents are inserted per bulk write. +const BatchSize = 1000 + +// ErrTargetNotEmpty is returned when the target database already holds data and +// Force was not set. +var ErrTargetNotEmpty = errors.New("target database is not empty") + +// ErrKeyMismatch is returned when the archive's key fingerprint does not match +// the key supplied. +var ErrKeyMismatch = errors.New("KEY_ENCRYPTION_KEY does not match the archive") + +// ErrIndexBuild is returned when a unique index in the archive cannot be built +// on the restored data. +var ErrIndexBuild = errors.New("index could not be built on the restored data") + +// RestoredCollection is what one collection's restore produced. +type RestoredCollection struct { + Name string + Documents int64 + Indexes int +} + +// RestoreResult is the summary a caller prints. +type RestoreResult struct { + Collections []RestoredCollection +} + +// RestoreOptions configures one restore. +type RestoreOptions struct { + Client *mongo.Client + Database string + Archive *Reader + + // Force drops each collection in the archive before loading it. Without it + // a non-empty target is refused. + Force bool + + KeyHex string + + // IgnoreKeyMismatch proceeds past a fingerprint mismatch, having first + // warned which collections will hold unreadable ciphertext afterwards. + IgnoreKeyMismatch bool + + // Warn receives operator-facing warnings. A nil Warn discards them. + Warn func(string) +} + +func (o RestoreOptions) warn(format string, args ...any) { + if o.Warn != nil { + o.Warn(fmt.Sprintf(format, args...)) + } +} + +// Restore loads an archive into a database. +// +// The order is fixed and every check that can refuse does so before the first +// write: format, checksums (done by Open), key policy, then target inspection. +// A restore that has begun writing and then fails leaves a partial database +// which the next run refuses to touch, which is correct — the alternative is a +// silent merge, and merging two control planes reconciles nothing. +func Restore(ctx context.Context, opt RestoreOptions) (RestoreResult, error) { + m := opt.Archive.Manifest() + + if err := checkKey(m, opt); err != nil { + return RestoreResult{}, err + } + if err := checkTarget(ctx, opt); err != nil { + return RestoreResult{}, err + } + if len(m.Excluded) > 0 { + opt.warn("this archive excluded %s; those collections will be empty after the restore", + strings.Join(m.Excluded, ", ")) + } + opt.warn("Redis is not restored. Sessions are the only state it holds, so everyone signs in again.") + + res := RestoreResult{} + if err := warnVersionGap(ctx, opt, m); err != nil { + return res, err + } + + for _, entry := range m.Collections { + rc, err := restoreCollection(ctx, opt, entry) + if err != nil { + return res, err + } + res.Collections = append(res.Collections, rc) + } + return res, nil +} + +// checkKey applies the fingerprint policy. +func checkKey(m Manifest, opt RestoreOptions) error { + if m.KeyFingerprint == nil { + opt.warn("this archive carries no key fingerprint, so nothing here proves your " + + "KEY_ENCRYPTION_KEY opens its ciphertext") + return nil + } + if opt.KeyHex == "" { + return fmt.Errorf("%w: the archive records a key fingerprint, so a key is required "+ + "(pass --ignore-key-mismatch only if you accept unreadable secrets)", ErrNoKey) + } + got, err := FingerprintHex(opt.KeyHex) + if err != nil { + return err + } + if got == *m.KeyFingerprint { + return nil + } + if !opt.IgnoreKeyMismatch { + return fmt.Errorf("%w: archive fingerprint %s, your key fingerprints as %s", + ErrKeyMismatch, *m.KeyFingerprint, got) + } + opt.warn("proceeding past a key mismatch: ciphertext in %s will be permanently unreadable", + strings.Join(CiphertextCollections(), ", ")) + return nil +} + +// checkTarget refuses a non-empty database unless Force was set. +func checkTarget(ctx context.Context, opt RestoreOptions) error { + db := opt.Client.Database(opt.Database) + names, err := db.ListCollectionNames(ctx, bson.M{}) + if err != nil { + return fmt.Errorf("inspect target: %w", err) + } + if len(names) == 0 || opt.Force { + return nil + } + sort.Strings(names) + + var found []string + for _, n := range names { + count, err := db.Collection(n).CountDocuments(ctx, bson.M{}) + if err != nil { + return fmt.Errorf("count %s: %w", n, err) + } + found = append(found, fmt.Sprintf("%s (%d)", n, count)) + } + return fmt.Errorf("%w: %s holds %s", ErrTargetNotEmpty, opt.Database, strings.Join(found, ", ")) +} + +// warnVersionGap reports a major version difference between the server that +// produced the archive and the one receiving it. It warns rather than refuses: +// restoring across a major version is a normal part of an upgrade, and a tool +// that refused would be blocking the migration it exists to make safe. +func warnVersionGap(ctx context.Context, opt RestoreOptions, m Manifest) error { + if m.MongoServerVersion == "" { + return nil + } + target, err := mongoServerVersion(ctx, opt.Client) + if err != nil { + return err + } + if majorOf(m.MongoServerVersion) != majorOf(target) { + opt.warn("this archive came from MongoDB %s and you are restoring onto %s", + m.MongoServerVersion, target) + } + return nil +} + +func majorOf(version string) string { + if i := strings.IndexByte(version, '.'); i >= 0 { + return version[:i] + } + return version +} + +func restoreCollection(ctx context.Context, opt RestoreOptions, entry CollectionEntry) (RestoredCollection, error) { + coll := opt.Client.Database(opt.Database).Collection(entry.Name) + if opt.Force { + if err := coll.Drop(ctx); err != nil { + return RestoredCollection{}, fmt.Errorf("drop %s: %w", entry.Name, err) + } + } + + written, err := insertDocuments(ctx, opt, coll, entry) + if err != nil { + return RestoredCollection{}, err + } + indexes, err := replayIndexes(ctx, opt, coll, entry.Name) + if err != nil { + return RestoredCollection{}, err + } + return RestoredCollection{Name: entry.Name, Documents: written, Indexes: indexes}, nil +} + +func insertDocuments(ctx context.Context, opt RestoreOptions, coll *mongo.Collection, entry CollectionEntry) (int64, error) { + rc, err := opt.Archive.OpenCollection(entry.Name) + if err != nil { + return 0, fmt.Errorf("open %s in archive: %w", entry.Name, err) + } + defer rc.Close() + + raw, err := io.ReadAll(rc) + if err != nil { + return 0, fmt.Errorf("read %s: %w", entry.Name, err) + } + + var written int64 + batch := make([]any, 0, BatchSize) + flush := func() error { + if len(batch) == 0 { + return nil + } + if _, err := coll.InsertMany(ctx, batch, options.InsertMany().SetOrdered(false)); err != nil { + return fmt.Errorf("insert into %s: %w", entry.Name, err) + } + written += int64(len(batch)) + batch = batch[:0] + return nil + } + + for len(raw) > 0 { + doc, rest, err := splitBSON(raw) + if err != nil { + return 0, fmt.Errorf("%s: %w", entry.Name, err) + } + batch = append(batch, doc) + raw = rest + if len(batch) == BatchSize { + if err := flush(); err != nil { + return 0, err + } + } + } + if err := flush(); err != nil { + return 0, err + } + return written, nil +} + +// splitBSON peels one document off the front of a concatenated BSON stream. A +// BSON document declares its own length in its first four bytes. +func splitBSON(raw []byte) (bson.Raw, []byte, error) { + if len(raw) < 4 { + return nil, nil, fmt.Errorf("truncated BSON: %d trailing bytes", len(raw)) + } + n := int(int32(raw[0]) | int32(raw[1])<<8 | int32(raw[2])<<16 | int32(raw[3])<<24) + if n < 5 || n > len(raw) { + return nil, nil, fmt.Errorf("BSON document declares length %d with %d bytes remaining", n, len(raw)) + } + return bson.Raw(raw[:n]), raw[n:], nil +} + +// replayIndexes recreates the archived indexes. +// +// A unique index that will not build means the restored data violates it, and +// the unique indexes here — (instance_id, email), instance slug, settings +// instance, the ESO token hash — are tenant-isolation properties rather than +// optimisations. That aborts. A non-unique index failing is a performance +// problem and warns. +func replayIndexes(ctx context.Context, opt RestoreOptions, coll *mongo.Collection, name string) (int, error) { + raw, err := opt.Archive.IndexesJSON(name) + if err != nil { + return 0, fmt.Errorf("read index specs for %s: %w", name, err) + } + if len(raw) == 0 { + return 0, nil + } + var specs []map[string]any + if err := json.Unmarshal(raw, &specs); err != nil { + return 0, fmt.Errorf("parse index specs for %s: %w", name, err) + } + + created := 0 + for _, spec := range specs { + model, indexName, unique, ok := indexModelFrom(spec) + if !ok { + continue + } + if _, err := coll.Indexes().CreateOne(ctx, model); err != nil { + if unique { + return created, fmt.Errorf("%w: %s on %s: %v", ErrIndexBuild, indexName, name, err) + } + opt.warn("index %s on %s was not created: %v", indexName, name, err) + continue + } + created++ + } + return created, nil +} + +// indexModelFrom converts one archived index specification into a model. +// The _id_ index is skipped: MongoDB creates it itself and refuses an explicit +// attempt to create it. +func indexModelFrom(spec map[string]any) (mongo.IndexModel, string, bool, bool) { + name, _ := spec["name"].(string) + if name == "_id_" { + return mongo.IndexModel{}, name, false, false + } + keys, ok := spec["key"].(map[string]any) + if !ok || len(keys) == 0 { + return mongo.IndexModel{}, name, false, false + } + + // JSON objects do not preserve order but compound index key order is + // significant, so the field order recorded by the server is recovered from + // the spec's own ordering where available and sorted otherwise. bson.M + // round-trips through json as a map; the archive therefore stores the key + // document and this reconstructs a deterministic bson.D from it. + fields := make([]string, 0, len(keys)) + for k := range keys { + fields = append(fields, k) + } + sort.Strings(fields) + d := make(bson.D, 0, len(fields)) + for _, f := range fields { + d = append(d, bson.E{Key: f, Value: keys[f]}) + } + + opts := options.Index().SetName(name) + unique := false + if u, ok := spec["unique"].(bool); ok && u { + unique = true + opts = opts.SetUnique(true) + } + if s, ok := spec["sparse"].(bool); ok && s { + opts = opts.SetSparse(true) + } + if e, ok := spec["expireAfterSeconds"].(float64); ok { + opts = opts.SetExpireAfterSeconds(int32(e)) + } + return mongo.IndexModel{Keys: d, Options: opts}, name, unique, true +} diff --git a/backup/restore_test.go b/backup/restore_test.go new file mode 100644 index 0000000..b56841e --- /dev/null +++ b/backup/restore_test.go @@ -0,0 +1,348 @@ +package backup + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" +) + +// archiveOf seeds a database, dumps it, and returns an opened Reader. +func archiveOf(t *testing.T, client *mongo.Client, keyHex string, allowNoKey bool) *Reader { + t.Helper() + _, srcDB := testDB(t) + seed(t, client, srcDB) + path, _ := dumpToFile(t, DumpOptions{ + Client: client, Database: srcDB, KeyHex: keyHex, AllowNoKey: allowNoKey, + }) + r, err := Open(path) + if err != nil { + t.Fatalf("Open: %v", err) + } + t.Cleanup(func() { r.Close() }) + return r +} + +func countIn(t *testing.T, client *mongo.Client, dbName, coll string) int64 { + t.Helper() + n, err := client.Database(dbName).Collection(coll).CountDocuments(context.Background(), bson.M{}) + if err != nil { + t.Fatalf("count %s: %v", coll, err) + } + return n +} + +func TestRestoreIntoEmptyDatabase(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, validKeyHex, false) + _, target := testDB(t) + + res, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: validKeyHex, + }) + if err != nil { + t.Fatalf("Restore: %v", err) + } + if countIn(t, client, target, "servers") != 2 { + t.Fatal("servers not restored") + } + if len(res.Collections) == 0 { + t.Fatal("result reports no collections") + } +} + +func TestRestoreRefusesNonEmptyTarget(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, validKeyHex, false) + _, target := testDB(t) + if _, err := client.Database(target).Collection("servers"). + InsertOne(context.Background(), bson.M{"name": "existing"}); err != nil { + t.Fatalf("seed target: %v", err) + } + + _, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: validKeyHex, + }) + if !errors.Is(err, ErrTargetNotEmpty) { + t.Fatalf("got %v, want ErrTargetNotEmpty", err) + } + if countIn(t, client, target, "servers") != 1 { + t.Fatal("a refused restore modified the target") + } +} + +func TestRestoreForceReplaces(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, validKeyHex, false) + _, target := testDB(t) + if _, err := client.Database(target).Collection("servers"). + InsertOne(context.Background(), bson.M{"name": "existing"}); err != nil { + t.Fatalf("seed target: %v", err) + } + + if _, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: validKeyHex, Force: true, + }); err != nil { + t.Fatalf("Restore --force: %v", err) + } + if got := countIn(t, client, target, "servers"); got != 2 { + t.Fatalf("servers has %d documents, want 2; force must drop, not merge", got) + } + n, err := client.Database(target).Collection("servers"). + CountDocuments(context.Background(), bson.M{"name": "existing"}) + if err != nil { + t.Fatalf("count: %v", err) + } + if n != 0 { + t.Fatal("the pre-existing document survived --force") + } +} + +func TestRestoreRefusesKeyMismatch(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, validKeyHex, false) + _, target := testDB(t) + other := "0000000000000000000000000000000000000000000000000000000000000002" + + _, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: other, + }) + if !errors.Is(err, ErrKeyMismatch) { + t.Fatalf("got %v, want ErrKeyMismatch", err) + } + names, err := client.Database(target).ListCollectionNames(context.Background(), bson.M{}) + if err != nil { + t.Fatalf("list: %v", err) + } + if len(names) != 0 { + t.Fatalf("a refused restore wrote %v", names) + } +} + +func TestRestoreRefusesWhenArchiveHasKeyAndEnvironmentDoesNot(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, validKeyHex, false) + _, target := testDB(t) + + _, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, + }) + if !errors.Is(err, ErrNoKey) { + t.Fatalf("got %v, want ErrNoKey", err) + } +} + +func TestRestoreIgnoreKeyMismatchWarnsAndProceeds(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, validKeyHex, false) + _, target := testDB(t) + other := "0000000000000000000000000000000000000000000000000000000000000002" + + var warnings []string + if _, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, + KeyHex: other, IgnoreKeyMismatch: true, + Warn: func(s string) { warnings = append(warnings, s) }, + }); err != nil { + t.Fatalf("Restore: %v", err) + } + joined := strings.Join(warnings, "\n") + for _, name := range CiphertextCollections() { + if !strings.Contains(joined, name) { + t.Fatalf("warning does not name %s\ngot:\n%s", name, joined) + } + } + if countIn(t, client, target, "servers") != 2 { + t.Fatal("restore did not proceed") + } +} + +func TestRestoreNullFingerprintIsReportedNotAssumed(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, "", true) + _, target := testDB(t) + + var warnings []string + if _, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: validKeyHex, + Warn: func(s string) { warnings = append(warnings, s) }, + }); err != nil { + t.Fatalf("Restore: %v", err) + } + if !strings.Contains(strings.Join(warnings, "\n"), "no key fingerprint") { + t.Fatalf("a null fingerprint must be reported, got: %v", warnings) + } +} + +func TestRestoreReplaysIndexes(t *testing.T) { + client, _ := testDB(t) + ctx := context.Background() + _, srcDB := testDB(t) + seed(t, client, srcDB) + if _, err := client.Database(srcDB).Collection("servers").Indexes(). + CreateOne(ctx, mongo.IndexModel{Keys: bson.D{{Key: "name", Value: 1}}}); err != nil { + t.Fatalf("create index: %v", err) + } + path, _ := dumpToFile(t, DumpOptions{Client: client, Database: srcDB, KeyHex: validKeyHex}) + archive, err := Open(path) + if err != nil { + t.Fatalf("Open: %v", err) + } + defer archive.Close() + + _, target := testDB(t) + res, err := Restore(ctx, RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: validKeyHex, + }) + if err != nil { + t.Fatalf("Restore: %v", err) + } + + cur, err := client.Database(target).Collection("servers").Indexes().List(ctx) + if err != nil { + t.Fatalf("list indexes: %v", err) + } + var specs []bson.M + if err := cur.All(ctx, &specs); err != nil { + t.Fatalf("read indexes: %v", err) + } + found := false + for _, s := range specs { + if s["name"] == "name_1" { + found = true + } + } + if !found { + t.Fatalf("index name_1 not replayed; got %v", specs) + } + for _, c := range res.Collections { + if c.Name == "servers" && c.Indexes < 1 { + t.Fatal("result reports no indexes created for servers") + } + } +} + +// TestRestoreAbortsWhenAUniqueIndexCannotBuild replaces the brief's +// TestRestoreAbortsOnUniqueIndexViolation per ruling 1: creating an index on +// the target also creates the collection, so that version's assertion +// (err == nil) would have passed on the wrong error (ErrTargetNotEmpty), and +// Force: true does not rescue it because the drop removes the index before +// replayIndexes runs. This version builds a hand-made archive whose data and +// index specification directly contradict each other, which is the actual +// shape of a corrupted archive that replayIndexes must refuse to load. +func TestRestoreAbortsWhenAUniqueIndexCannotBuild(t *testing.T) { + client, _ := testDB(t) + ctx := context.Background() + + // Built by hand rather than dumped: two documents that collide on email + // alongside an index specification declaring email unique. No live database + // would let those coexist, which is exactly the point — this is the shape + // of a corrupted or hand-edited archive, and restore must refuse rather + // than load the rows and leave the index missing. + a, err := bson.Marshal(bson.M{"email": "a@example.com"}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + b, err := bson.Marshal(bson.M{"email": "a@example.com"}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + path := filepath.Join(t.TempDir(), "dupes.tar.gz") + f, err := os.Create(path) + if err != nil { + t.Fatalf("create: %v", err) + } + w := NewWriter(f) + entry, err := w.WriteCollection("users", [][]byte{a, b}) + if err != nil { + t.Fatalf("WriteCollection: %v", err) + } + if err := w.WriteIndexes("users", + []byte(`[{"name":"email_1","key":{"email":1},"unique":true}]`)); err != nil { + t.Fatalf("WriteIndexes: %v", err) + } + fp, err := FingerprintHex(validKeyHex) + if err != nil { + t.Fatalf("FingerprintHex: %v", err) + } + if err := w.Close(Manifest{ + FormatVersion: FormatVersion, + CreatedAt: time.Now().UTC(), + MongoDB: "handmade", + KeyFingerprint: &fp, + Collections: []CollectionEntry{entry}, + }); err != nil { + t.Fatalf("Close: %v", err) + } + if err := f.Close(); err != nil { + t.Fatalf("close: %v", err) + } + + archive, err := Open(path) + if err != nil { + t.Fatalf("Open: %v", err) + } + defer archive.Close() + + _, target := testDB(t) + if _, err := Restore(ctx, RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: validKeyHex, + }); !errors.Is(err, ErrIndexBuild) { + t.Fatalf("got %v, want ErrIndexBuild", err) + } else if !strings.Contains(err.Error(), "email_1") { + t.Fatalf("the error must name the offending index, got: %v", err) + } +} + +func TestRestoreSameVersionDoesNotWarnAboutIt(t *testing.T) { + client, _ := testDB(t) + archive := archiveOf(t, client, validKeyHex, false) + _, target := testDB(t) + + var warnings []string + if _, err := Restore(context.Background(), RestoreOptions{ + Client: client, Database: target, Archive: archive, KeyHex: validKeyHex, + Warn: func(s string) { warnings = append(warnings, s) }, + }); err != nil { + t.Fatalf("Restore: %v", err) + } + for _, w := range warnings { + if strings.Contains(w, "you are restoring onto") { + t.Fatalf("same-version restore warned about a version gap: %s", w) + } + } +} + +func TestCompoundIndexKeyOrderIsAlphabetical(t *testing.T) { + model, name, unique, ok := indexModelFrom(map[string]any{ + "name": "b_1_a_1", + "key": map[string]any{"b": float64(1), "a": float64(1)}, + }) + if !ok { + t.Fatal("spec rejected") + } + if unique { + t.Fatal("index reported as unique") + } + if name != "b_1_a_1" { + t.Fatalf("name %q", name) + } + keys, isD := model.Keys.(bson.D) + if !isD { + t.Fatalf("keys are %T, want bson.D", model.Keys) + } + // Documents current behaviour: field order is alphabetical, not the order + // the server reported. Storing the key document as raw BSON in the archive + // instead of JSON would fix this and is the change to make if compound + // index order ever matters here. + if keys[0].Key != "a" || keys[1].Key != "b" { + t.Fatalf("got %v", keys) + } +}