didplcbft / abciapp / snapshots.go

package abciapp

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	abcitypes "github.com/cometbft/cometbft/abci/types"
	"github.com/cosmos/iavl"
	"github.com/klauspost/compress/zstd"
	"github.com/palantir/stacktrace"
)

const snapshotChunkSize = 10 * 1024 * 1024 // 10 MB
const snapshotChunkHashSize = 32

// ListSnapshots implements [types.Application].
func (d *DIDPLCApplication) ListSnapshots(context.Context, *abcitypes.RequestListSnapshots) (*abcitypes.ResponseListSnapshots, error) {
	files, err := filepath.Glob(filepath.Join(d.snapshotDirectory, "*.snapshot"))
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}

	snapshots := make([]*abcitypes.Snapshot, 0, len(files))
	for _, filename := range files {
		s, err := readSnapshotMetadata(filename)
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}

		snapshots = append(snapshots, s)
	}

	return &abcitypes.ResponseListSnapshots{
		Snapshots: snapshots,
	}, nil
}

func readSnapshotMetadata(filename string) (*abcitypes.Snapshot, error) {
	// Extract height from filename pattern: %020d.snapshot
	base := filepath.Base(filename)
	if !strings.HasSuffix(base, ".snapshot") {
		return nil, stacktrace.NewError("invalid snapshot filename format: %s", filename)
	}
	heightStr := strings.TrimSuffix(base, ".snapshot")
	height, err := strconv.ParseInt(heightStr, 10, 64)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to parse height from filename: %s", filename)
	}

	// Open and read snapshot file header
	f, err := os.Open(filename)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to open snapshot file: %s", filename)
	}
	defer f.Close()

	// Read file magic (18 bytes)
	magic := make([]byte, 18)
	_, err = io.ReadFull(f, magic)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to read file magic")
	}
	if string(magic) != "didplcbft-snapshot" {
		return nil, stacktrace.NewError("invalid file magic: expected 'didplcbft-snapshot', got '%s'", string(magic))
	}

	// Read version bytes (6 bytes)
	versionBytes := make([]byte, 6)
	_, err = io.ReadFull(f, versionBytes)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to read version bytes")
	}
	format := binary.BigEndian.Uint32(versionBytes[2:])

	// Read height (8 bytes, big-endian)
	heightBytes := make([]byte, 8)
	_, err = io.ReadFull(f, heightBytes)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to read height")
	}
	fileHeight := int64(binary.BigEndian.Uint64(heightBytes))
	if fileHeight != height {
		return nil, stacktrace.NewError("height mismatch: filename indicates %d, file header contains %d", height, fileHeight)
	}

	// Read tree hash (32 bytes)
	hash := make([]byte, 32)
	_, err = io.ReadFull(f, hash)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to read tree hash")
	}

	// Read corresponding chunksums file
	chunksumsFilename := strings.TrimSuffix(filename, ".snapshot") + ".chunksums"
	chunksumsData, err := os.ReadFile(chunksumsFilename)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to read chunksums file: %s", chunksumsFilename)
	}

	// Calculate number of chunks (each chunk hash is 32 bytes)
	chunks := int64(len(chunksumsData)) / snapshotChunkHashSize

	return &abcitypes.Snapshot{
		Height:   uint64(height),
		Format:   format,
		Chunks:   uint32(chunks),
		Hash:     hash,
		Metadata: chunksumsData,
	}, nil
}

// LoadSnapshotChunk implements [types.Application].
func (d *DIDPLCApplication) LoadSnapshotChunk(_ context.Context, req *abcitypes.RequestLoadSnapshotChunk) (*abcitypes.ResponseLoadSnapshotChunk, error) {
	if req.Format != 1 {
		// just in case CometBFT asks us to load a chunk of a format we didn't declare to support in ListSnapshots...
		return nil, stacktrace.NewError("unsupported snapshot format")
	}

	// Construct filename from height using the same pattern as createSnapshot
	snapshotFilename := filepath.Join(d.snapshotDirectory, fmt.Sprintf("%020d.snapshot", req.Height))

	// Open the snapshot file
	f, err := os.Open(snapshotFilename)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to open snapshot file: %s", snapshotFilename)
	}
	defer f.Close()

	// Calculate the offset for the requested chunk (start from beginning of file, including header)
	offset := int64(req.Chunk) * snapshotChunkSize
	_, err = f.Seek(offset, io.SeekStart)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to seek to chunk offset")
	}

	// Read up to snapshotChunkSize bytes
	chunkData := make([]byte, snapshotChunkSize)
	n, err := f.Read(chunkData)
	if err != nil && err != io.EOF {
		return nil, stacktrace.Propagate(err, "failed to read chunk data")
	}

	// If we read less than snapshotChunkSize, trim the slice
	if n < snapshotChunkSize {
		chunkData = chunkData[:n]
	}

	return &abcitypes.ResponseLoadSnapshotChunk{
		Chunk: chunkData,
	}, nil
}

// ApplySnapshotChunk implements [types.Application].
func (d *DIDPLCApplication) ApplySnapshotChunk(_ context.Context, req *abcitypes.RequestApplySnapshotChunk) (*abcitypes.ResponseApplySnapshotChunk, error) {
	if d.snapshotApplier == nil {
		return nil, stacktrace.NewError("snapshot not offered yet, can't apply chunk")
	}

	err := d.snapshotApplier.Apply(int(req.Index), req.Chunk)
	if err != nil {
		if errors.Is(err, errMalformedChunk) {
			return &abcitypes.ResponseApplySnapshotChunk{
				Result:        abcitypes.ResponseApplySnapshotChunk_RETRY,
				RefetchChunks: []uint32{req.Index},
				RejectSenders: []string{req.Sender},
			}, nil
		} else if errors.Is(err, errTreeHashMismatch) {
			return &abcitypes.ResponseApplySnapshotChunk{
				Result:        abcitypes.ResponseApplySnapshotChunk_REJECT_SNAPSHOT,
				RejectSenders: []string{req.Sender},
			}, nil
		}
		return nil, stacktrace.NewError("failed to apply")
	}

	if d.snapshotApplier.Done() {
		d.snapshotApplier = nil
		err := d.plc.RefreshTreeData()
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}
	}

	return &abcitypes.ResponseApplySnapshotChunk{
		Result: abcitypes.ResponseApplySnapshotChunk_ACCEPT,
	}, nil
}

// OfferSnapshot implements [types.Application].
func (d *DIDPLCApplication) OfferSnapshot(_ context.Context, req *abcitypes.RequestOfferSnapshot) (*abcitypes.ResponseOfferSnapshot, error) {
	if d.snapshotApplier != nil {
		err := d.snapshotApplier.Abort()
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}
	}

	if req.Snapshot.Format != 1 {
		return &abcitypes.ResponseOfferSnapshot{
			Result: abcitypes.ResponseOfferSnapshot_REJECT_FORMAT,
		}, nil
	}

	var err error
	d.snapshotApplier, err = d.beginApplyingSnapshot(int64(req.Snapshot.Height), req.AppHash, int(req.Snapshot.Chunks), req.Snapshot.Metadata)
	if err != nil {
		d.snapshotApplier = nil
		if errors.Is(err, errInvalidMetadata) {
			return &abcitypes.ResponseOfferSnapshot{
				Result: abcitypes.ResponseOfferSnapshot_REJECT_SENDER,
			}, nil
		}
		return nil, stacktrace.Propagate(err, "")
	}

	return &abcitypes.ResponseOfferSnapshot{
		Result: abcitypes.ResponseOfferSnapshot_ACCEPT,
	}, nil
}

func (d *DIDPLCApplication) createSnapshot(treeVersion int64, tempFilename string) error {
	it, err := d.tree.GetImmutable(treeVersion)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	// Delete tempFilename if it exists to ensure a fresh file is created
	_ = os.Remove(tempFilename)

	f, err := os.Create(tempFilename)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	defer f.Close()

	st := time.Now()

	err = writeSnapshot(f, it)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = f.Sync()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	hf, err := os.Create(filepath.Join(d.snapshotDirectory, fmt.Sprintf("%020d.chunksums", treeVersion)))
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	defer hf.Close()

	err = writeChunkHashes(f, hf)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = hf.Sync()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = f.Close()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	os.Rename(tempFilename, filepath.Join(d.snapshotDirectory, fmt.Sprintf("%020d.snapshot", treeVersion)))

	fmt.Println("Took", time.Since(st), "to export")

	return nil
}

func writeSnapshot(writerSeeker io.WriteSeeker, it *iavl.ImmutableTree) error {
	writtenUntilReservedFields := 0

	bw := bufio.NewWriter(writerSeeker)

	// file magic and version
	c, err := bw.Write([]byte("didplcbft-snapshot"))
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	writtenUntilReservedFields += c

	c, err = bw.Write([]byte{0, 0, 0, 0, 0, 1})
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	writtenUntilReservedFields += c

	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(it.Version()))
	c, err = bw.Write(b)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	writtenUntilReservedFields += c

	c, err = bw.Write(it.Hash())
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	writtenUntilReservedFields += c

	// reserve space for writing number of bytes, number of nodes
	// 8 bytes for node list size in bytes
	// 8 bytes for number of nodes
	sizeOfReservedFields := 8 + 8
	b = make([]byte, sizeOfReservedFields)
	_, err = bw.Write(b)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	zstdw, err := zstd.NewWriter(bw, zstd.WithEncoderLevel(zstd.SpeedBetterCompression))
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	numNodes, err := exportNodes(it, zstdw)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = zstdw.Close()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = bw.Flush()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	// find total compressed node list file size
	offset, err := writerSeeker.Seek(0, io.SeekCurrent)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	compressedNodeListSize := offset - int64(writtenUntilReservedFields) - int64(sizeOfReservedFields)

	// seek back and write empty header fields

	offset, err = writerSeeker.Seek(int64(writtenUntilReservedFields), io.SeekStart)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}
	if offset != int64(writtenUntilReservedFields) {
		return stacktrace.NewError("unexpected seek result")
	}

	b = make([]byte, sizeOfReservedFields)
	binary.BigEndian.PutUint64(b, uint64(compressedNodeListSize))
	binary.BigEndian.PutUint64(b[8:], uint64(numNodes))
	_, err = writerSeeker.Write(b)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	return nil
}

func exportNodes(it *iavl.ImmutableTree, w io.Writer) (int64, error) {
	exporter, err := it.Export()
	if err != nil {
		return 0, stacktrace.Propagate(err, "")
	}
	defer exporter.Close()
	cexporter := iavl.NewCompressExporter(exporter)

	numNodes := int64(0)
	for {
		node, err := cexporter.Next()
		if errors.Is(err, iavl.ErrorExportDone) {
			break
		}
		if err != nil {
			return 0, stacktrace.Propagate(err, "")
		}

		b := make([]byte, 9)
		b[0] = byte(node.Height)

		binary.BigEndian.PutUint64(b[1:], uint64(node.Version))
		_, err = w.Write(b)
		if err != nil {
			return 0, stacktrace.Propagate(err, "")
		}

		// nil node values are different from 0-byte values
		b = []byte{0xff, 0xff, 0xff, 0xff}
		if node.Key != nil {
			binary.BigEndian.PutUint32(b, uint32(len(node.Key)))
		}
		_, err = w.Write(b)
		if err != nil {
			return 0, stacktrace.Propagate(err, "")
		}

		b = []byte{0xff, 0xff, 0xff, 0xff}
		if node.Value != nil {
			binary.BigEndian.PutUint32(b, uint32(len(node.Value)))
		}
		_, err = w.Write(b)
		if err != nil {
			return 0, stacktrace.Propagate(err, "")
		}

		_, err = w.Write(node.Key)
		if err != nil {
			return 0, stacktrace.Propagate(err, "")
		}

		_, err = w.Write(node.Value)
		if err != nil {
			return 0, stacktrace.Propagate(err, "")
		}
		numNodes++
	}

	return numNodes, nil
}

func writeChunkHashes(snapshotFile io.ReadSeeker, w io.Writer) error {
	bw := bufio.NewWriter(w)
	defer bw.Flush()

	_, err := snapshotFile.Seek(0, io.SeekStart)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	buf := make([]byte, snapshotChunkSize)
	for {
		n, err := io.ReadFull(snapshotFile, buf)
		if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
			return stacktrace.Propagate(err, "")
		}
		if n == 0 {
			break
		}

		hash := sha256.Sum256(buf[:n])
		c, err := w.Write(hash[:])
		if err != nil {
			return stacktrace.Propagate(err, "")
		}

		if c != snapshotChunkHashSize {
			return stacktrace.NewError("unexpected chunk hash size")
		}

		if n < snapshotChunkSize {
			break
		}
	}

	return nil
}

type snapshotApplier struct {
	tree                *iavl.MutableTree
	treeVersion         int64
	expectedFinalHash   []byte
	expectedChunkHashes [][]byte

	pipeWriter *io.PipeWriter
	pipeReader *io.PipeReader
	zstdReader io.ReadCloser

	importer         *iavl.Importer
	compressImporter iavl.NodeImporter
	importerWg       sync.WaitGroup

	numImportedNodes int
	claimedNodeCount int
	done             bool
}

var errMalformedChunk = errors.New("malformed chunk")
var errInvalidMetadata = errors.New("invalid metadata")
var errTreeHashMismatch = errors.New("tree hash mismatch")

func (d *DIDPLCApplication) beginApplyingSnapshot(treeVersion int64, expectedFinalHash []byte, expectedNumChunks int, chunksums []byte) (*snapshotApplier, error) {
	if len(chunksums)%snapshotChunkHashSize != 0 || len(chunksums)/snapshotChunkHashSize != expectedNumChunks {
		return nil, stacktrace.Propagate(errInvalidMetadata, "")
	}

	if !d.tree.IsEmpty() {
		err := d.fullyClearTree()
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}
	}

	importer, err := d.tree.Import(treeVersion)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}

	pipeReader, pipeWriter := io.Pipe()

	zstdReader, err := zstd.NewReader(pipeReader)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}

	chunkHashes := make([][]byte, 0, expectedNumChunks)
	for hash := range slices.Chunk(chunksums, snapshotChunkHashSize) {
		chunkHashes = append(chunkHashes, hash)
	}

	return &snapshotApplier{
		tree:                d.tree,
		treeVersion:         treeVersion,
		expectedFinalHash:   expectedFinalHash,
		expectedChunkHashes: chunkHashes,

		pipeWriter: pipeWriter,
		pipeReader: pipeReader,
		zstdReader: zstdReader.IOReadCloser(),

		importer:         importer,
		compressImporter: iavl.NewCompressImporter(importer),
	}, nil
}

func (a *snapshotApplier) Apply(chunkIndex int, chunkBytes []byte) error {
	if len(chunkBytes) > snapshotChunkSize {
		return stacktrace.Propagate(errMalformedChunk, "chunk too large")
	}
	hash := sha256.Sum256(chunkBytes)
	if !bytes.Equal(a.expectedChunkHashes[chunkIndex], hash[:]) {
		return stacktrace.Propagate(errMalformedChunk, "hash mismatch")
	}

	if chunkIndex == 0 {
		if len(chunkBytes) < 80 {
			return stacktrace.Propagate(errMalformedChunk, "chunk too small")
		}

		if string(chunkBytes[0:18]) != "didplcbft-snapshot" {
			return stacktrace.Propagate(errMalformedChunk, "invalid file magic")
		}

		if binary.BigEndian.Uint32(chunkBytes[20:]) != 1 {
			return stacktrace.Propagate(errMalformedChunk, "invalid snapshot format")
		}

		if binary.BigEndian.Uint64(chunkBytes[24:]) != uint64(a.treeVersion) {
			return stacktrace.Propagate(errMalformedChunk, "mismatched tree version")
		}

		if !bytes.Equal(chunkBytes[32:64], a.expectedFinalHash) {
			return stacktrace.Propagate(errMalformedChunk, "mismatched declared tree hash")
		}

		declaredFileSize := 80 + binary.BigEndian.Uint64(chunkBytes[64:])
		minExpectedSize := uint64((len(a.expectedChunkHashes) - 1) * snapshotChunkSize)
		maxExpectedSize := uint64(len(a.expectedChunkHashes) * snapshotChunkSize)
		if declaredFileSize < minExpectedSize ||
			declaredFileSize > maxExpectedSize {
			return stacktrace.Propagate(errMalformedChunk, "unexpected compressed node list length")
		}

		a.claimedNodeCount = int(binary.BigEndian.Uint64(chunkBytes[72:]))

		// move to the start of the compressed portion
		chunkBytes = chunkBytes[80:]

		a.importerWg.Go(a.streamingImporter)
	}

	isLastChunk := chunkIndex == len(a.expectedChunkHashes)-1
	go func(b []byte) {
		// From the docs:
		// It is safe to call Read and Write in parallel with each other or with Close.
		// Parallel calls to Read and parallel calls to Write are also safe:
		// the individual calls will be gated sequentially.

		// so even if not everything gets written from this chunk (e.g. because the zstd decoder decided not to advance)
		// it'll eventually be written, in the correct order
		_, _ = a.pipeWriter.Write(b)
		if isLastChunk {
			_ = a.pipeWriter.Close()
		}
	}(chunkBytes)

	if isLastChunk {
		// wait for importer to finish reading and importing everything
		a.importerWg.Wait()

		if a.numImportedNodes != a.claimedNodeCount {
			return stacktrace.Propagate(errTreeHashMismatch, "imported node count mismatch")
		}

		err := a.importer.Commit()
		if err != nil {
			if strings.Contains(err.Error(), "invalid node structure") {
				return stacktrace.Propagate(errors.Join(errMalformedChunk, err), "")
			}
			return stacktrace.Propagate(err, "")
		}

		a.closeCommons()
		a.done = true

		if !bytes.Equal(a.tree.Hash(), a.expectedFinalHash) {
			return stacktrace.Propagate(errTreeHashMismatch, "")
		}
	}

	return nil
}

func (a *snapshotApplier) streamingImporter() {
	for {
		nodeHeader := make([]byte, 9+4+4)
		n, err := io.ReadFull(a.zstdReader, nodeHeader)
		if err != nil || n != 9+4+4 {
			// err may be EOF here, which is expected
			return
		}

		// validate lengths against sensible limits to prevent OOM DoS by malicious third parties
		keyLength := binary.BigEndian.Uint32(nodeHeader[9:13])
		var key []byte
		if keyLength != 0xffffffff {
			if keyLength > 1024*1024 {
				return
			}
			key = make([]byte, keyLength)

			n, err = io.ReadFull(a.zstdReader, key)
			if err != nil || n != len(key) {
				// this shouldn't happen unless the data is corrupt
				// we can return silently here because since we didn't import all nodes, the tree hash won't match anyway
				return
			}
		}

		valueLength := binary.BigEndian.Uint32(nodeHeader[13:17])
		var value []byte
		if valueLength != 0xffffffff {
			if valueLength > 1024*1024 {
				return
			}
			value = make([]byte, valueLength)
			n, err = io.ReadFull(a.zstdReader, value)
			if err != nil || n != len(value) {
				return
			}
		}

		err = a.compressImporter.Add(&iavl.ExportNode{
			Height:  int8(nodeHeader[0]),
			Version: int64(binary.BigEndian.Uint64(nodeHeader[1:9])),
			Key:     key,
			Value:   value,
		})
		if err != nil {
			// this shouldn't happen unless the data is corrupt
			// we can return silently here because since we didn't import all nodes, the tree hash won't match anyway
			return
		}
		a.numImportedNodes++
	}
}

func (a *snapshotApplier) Abort() error {
	err := a.closeCommons()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = a.tree.DeleteVersionsFrom(0)
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	return nil
}

func (a *snapshotApplier) closeCommons() error {
	err := a.zstdReader.Close()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = a.pipeReader.Close()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	err = a.pipeWriter.Close()
	if err != nil {
		return stacktrace.Propagate(err, "")
	}

	a.importerWg.Wait()
	a.importer.Close()

	return nil
}

func (a *snapshotApplier) Done() bool {
	return a.done
}