Fix DecodeDatabaseTo() with lock page (#51)

Author: benbjohnson

decoder.go

@@ -72,6 +72,8 @@ func (dec *Decoder) Close() error {
 		return fmt.Errorf("unmarshal trailer: %w", err)
 	}
 
+	// TODO: Ensure last read page is equal to the commit for snapshot LTX files
+
 	dec.writeToHash(b[:TrailerChecksumOffset])
 
 	// Compare checksum with checksum in trailer.
@@ -203,24 +205,44 @@ func (dec *Decoder) Verify() error {
 func (dec *Decoder) DecodeDatabaseTo(w io.Writer) error {
 	if err := dec.DecodeHeader(); err != nil {
 		return fmt.Errorf("decode header: %w", err)
-	} else if !dec.header.IsSnapshot() {
+	}
+
+	hdr := dec.Header()
+	lockPgno := hdr.LockPgno()
+	if !dec.header.IsSnapshot() {
 		return fmt.Errorf("cannot decode non-snapshot LTX file to SQLite database")
 	}
 
 	var pageHeader PageHeader
 	data := make([]byte, dec.header.PageSize)
-	for i := 0; ; i++ {
-		if err := dec.DecodePage(&pageHeader, data); err == io.EOF {
-			break
-		} else if err != nil {
-			return fmt.Errorf("decode page %d: %w", i, err)
+	for pgno := uint32(1); pgno <= hdr.Commit; pgno++ {
+		if pgno == lockPgno {
+			// Write empty page for lock page.
+			for i := range data {
+				data[i] = 0
+			}
+		} else {
+			// Otherwise read the page from the LTX decoder.
+			if err := dec.DecodePage(&pageHeader, data); err != nil {
+				return fmt.Errorf("decode page %d: %w", pgno, err)
+			} else if pageHeader.Pgno != pgno {
+				return fmt.Errorf("unexpected pgno while decoding page: read %d, expected %d", pageHeader.Pgno, pgno)
+			}
 		}
 
 		if _, err := w.Write(data); err != nil {
-			return fmt.Errorf("write page %d: %w", i, err)
+			return fmt.Errorf("write page %d: %w", pgno, err)
 		}
 	}
 
+	// Issue one more final read and expect to see an EOF. This is required so
+	// that the decoder can successfully close and validate.
+	if err := dec.DecodePage(&pageHeader, data); err == nil {
+		return fmt.Errorf("unexpected page %d after commit %d", pageHeader.Pgno, hdr.Commit)
+	} else if err != io.EOF {
+		return fmt.Errorf("unexpected error decoding after end of database: %w", err)
+	}
+
 	if err := dec.Close(); err != nil {
 		return fmt.Errorf("close decoder: %w", err)
 	}

decoder_test.go

@@ -147,6 +147,46 @@ func TestDecoder_DecodeDatabaseTo(t *testing.T) {
 		}
 	})
 
+	t.Run("WithLockPage", func(t *testing.T) {
+		lockPgno := ltx.LockPgno(4096)
+		commit := lockPgno + 10
+
+		var want bytes.Buffer
+		var buf bytes.Buffer
+		enc := ltx.NewEncoder(&buf)
+		if err := enc.EncodeHeader(ltx.Header{Version: 1, Flags: 0, PageSize: 4096, Commit: commit, MinTXID: 1, MaxTXID: 2, Timestamp: 1000}); err != nil {
+			t.Fatal(err)
+		}
+
+		pageData := bytes.Repeat([]byte("x"), 4096)
+		for pgno := uint32(1); pgno <= commit; pgno++ {
+			if pgno == lockPgno {
+				_, _ = want.Write(make([]byte, 4096))
+				continue
+			}
+
+			_, _ = want.Write(pageData)
+			if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, pageData); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		enc.SetPostApplyChecksum(0xc19b668c376662c7)
+		if err := enc.Close(); err != nil {
+			t.Fatal(err)
+		}
+
+		// Decode serialized LTX file.
+		dec := ltx.NewDecoder(&buf)
+
+		var out bytes.Buffer
+		if err := dec.DecodeDatabaseTo(&out); err != nil {
+			t.Fatal(err)
+		} else if got, want := out.Bytes(), want.Bytes(); !bytes.Equal(got, want) {
+			t.Fatal("output mismatch")
+		}
+	})
+
 	t.Run("ErrNonSnapshot", func(t *testing.T) {
 		spec := &ltx.FileSpec{
 			Header: ltx.Header{Version: 1, Flags: 0, PageSize: 512, Commit: 2, MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1},

ltx.go

@@ -223,6 +223,11 @@ func (h *Header) IsSnapshot() bool {
 	return h.MinTXID == 1
 }
 
+// LockPgno returns the lock page number based on the header's page size.
+func (h *Header) LockPgno() uint32 {
+	return LockPgno(h.PageSize)
+}
+
 // Validate returns an error if h is invalid.
 func (h *Header) Validate() error {
 	if h.Version != Version {