Andy Dixon
71556b76de
Decompression and checksumming overlap with parsing via a pipe, and a concrete-typed varint reader removes per-byte interface calls (0.39s -> 0.29s for 458K nodes).
474 lines
13 KiB
Go
474 lines
13 KiB
Go
// Package snapshot implements the versioned, checksummed binary inventory
|
||
// format and snapshot comparison.
|
||
//
|
||
// Layout (all integers little-endian / varint):
|
||
//
|
||
// magic 8 bytes brand.SnapshotMagic
|
||
// version uint16 currently 1
|
||
// flags uint16 bit 0: body is gzip-compressed (always set by Save)
|
||
// body gzip stream of:
|
||
// meta uvarint length + JSON (root, timestamps, counters, options)
|
||
// exts uvarint count, then count × (uvarint length + bytes)
|
||
// nodes uvarint count, then per node:
|
||
// uvarint parentDelta (id-parent; 0 only for the root)
|
||
// byte kind, byte category, uvarint flags
|
||
// uvarint nameLen + name
|
||
// varint size, varint alloc, varint mtime
|
||
// uvarint mode, uid, gid, nlink, ext
|
||
// errors uvarint count, then per record: uvarint node, byte kind,
|
||
// uvarint len + name, uvarint len + message
|
||
// sha256 32 bytes over the uncompressed body
|
||
//
|
||
// Nodes are written in ID order; a parent always precedes its children, so a
|
||
// reader can rebuild the tree in one pass and malformed input cannot create
|
||
// cycles. Aggregates are recomputed on load, never trusted.
|
||
package snapshot
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"compress/gzip"
|
||
"crypto/sha256"
|
||
"encoding/binary"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/andydixon/mapsize/internal/brand"
|
||
"github.com/andydixon/mapsize/internal/inventory"
|
||
)
|
||
|
||
// Version is the current format version.
|
||
const Version = 1
|
||
|
||
// Limits applied when reading untrusted snapshots.
|
||
const (
|
||
MaxNodes = 1<<31 - 1
|
||
MaxNameLen = 1 << 16
|
||
MaxMetaLen = 16 << 20
|
||
MaxExtLen = 64
|
||
MaxMsgLen = 4096
|
||
maxExts = 1<<16 - 1
|
||
flagGzipped = 1
|
||
)
|
||
|
||
type meta struct {
|
||
Root string `json:"root"`
|
||
Created time.Time `json:"created"`
|
||
Start time.Time `json:"start"`
|
||
End time.Time `json:"end"`
|
||
Complete bool `json:"complete"`
|
||
Cancelled bool `json:"cancelled"`
|
||
Files int64 `json:"files"`
|
||
Dirs int64 `json:"dirs"`
|
||
Symlinks int64 `json:"symlinks"`
|
||
Others int64 `json:"others"`
|
||
ErrCounts []int64 `json:"err_counts"`
|
||
Excluded int64 `json:"excluded"`
|
||
SkippedMounts int64 `json:"skipped_mounts"`
|
||
VirtualSkipped int64 `json:"virtual_skipped"`
|
||
LoopsSkipped int64 `json:"loops_skipped"`
|
||
BrokenLinks int64 `json:"broken_links"`
|
||
HardlinkDups int64 `json:"hardlink_dups"`
|
||
Unscanned int64 `json:"unscanned"`
|
||
Excludes []string `json:"excludes"`
|
||
OneFileSystem bool `json:"one_file_system"`
|
||
Follow string `json:"follow"`
|
||
Workers int `json:"workers"`
|
||
Host string `json:"host"`
|
||
Generator string `json:"generator"`
|
||
}
|
||
|
||
// IsSnapshot reports whether path starts with the snapshot magic.
|
||
func IsSnapshot(path string) bool {
|
||
f, err := os.Open(path)
|
||
if err != nil {
|
||
return false
|
||
}
|
||
defer f.Close()
|
||
b := make([]byte, len(brand.SnapshotMagic))
|
||
_, err = io.ReadFull(f, b)
|
||
return err == nil && string(b) == brand.SnapshotMagic
|
||
}
|
||
|
||
type writer struct {
|
||
w io.Writer
|
||
buf [binary.MaxVarintLen64]byte
|
||
err error
|
||
}
|
||
|
||
func (w *writer) bytes(b []byte) {
|
||
if w.err == nil {
|
||
_, w.err = w.w.Write(b)
|
||
}
|
||
}
|
||
func (w *writer) uvarint(v uint64) { w.bytes(w.buf[:binary.PutUvarint(w.buf[:], v)]) }
|
||
func (w *writer) varint(v int64) { w.bytes(w.buf[:binary.PutVarint(w.buf[:], v)]) }
|
||
func (w *writer) str(s string) { w.uvarint(uint64(len(s))); w.bytes([]byte(s)) }
|
||
|
||
// Save writes t to w. The caller must hold t's read lock.
|
||
func Save(out io.Writer, t *inventory.Tree) error {
|
||
bw := bufio.NewWriterSize(out, 1<<20)
|
||
hdr := make([]byte, 0, 12)
|
||
hdr = append(hdr, brand.SnapshotMagic...)
|
||
hdr = binary.LittleEndian.AppendUint16(hdr, Version)
|
||
hdr = binary.LittleEndian.AppendUint16(hdr, flagGzipped)
|
||
bw.Write(hdr)
|
||
|
||
gz, _ := gzip.NewWriterLevel(bw, gzip.BestSpeed)
|
||
h := sha256.New()
|
||
w := &writer{w: io.MultiWriter(gz, h)}
|
||
|
||
st := &t.Stats
|
||
host, _ := os.Hostname()
|
||
m := meta{Root: st.Root, Created: time.Now().UTC(), Start: st.Start, End: st.End, Complete: st.Complete,
|
||
Cancelled: st.Cancelled, Files: st.Files, Dirs: st.Dirs, Symlinks: st.Symlinks, Others: st.Others,
|
||
ErrCounts: st.ErrCounts[:], Excluded: st.Excluded, SkippedMounts: st.SkippedMounts,
|
||
VirtualSkipped: st.VirtualSkipped, LoopsSkipped: st.LoopsSkipped, BrokenLinks: st.BrokenLinks,
|
||
HardlinkDups: st.HardlinkDups, Unscanned: st.Unscanned, Excludes: st.Excludes,
|
||
OneFileSystem: st.OneFileSystem, Follow: st.Follow, Workers: st.Workers, Host: host,
|
||
Generator: brand.Name + " " + brand.Version}
|
||
mb, err := json.Marshal(m)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
w.uvarint(uint64(len(mb)))
|
||
w.bytes(mb)
|
||
|
||
w.uvarint(uint64(t.ExtCount()))
|
||
for i := range t.ExtCount() {
|
||
w.str(t.ExtByIndex(i))
|
||
}
|
||
w.uvarint(uint64(t.Len()))
|
||
for i := range t.Len() {
|
||
id := inventory.NodeID(i)
|
||
n := t.Node(id)
|
||
if i == 0 {
|
||
w.uvarint(0)
|
||
} else {
|
||
w.uvarint(uint64(id - n.Parent))
|
||
}
|
||
w.bytes([]byte{byte(n.Kind), byte(n.Cat)})
|
||
w.uvarint(uint64(n.Flags))
|
||
w.str(n.Name)
|
||
w.varint(n.Size)
|
||
w.varint(n.Alloc)
|
||
w.varint(n.MTime)
|
||
w.uvarint(uint64(n.Mode))
|
||
w.uvarint(uint64(n.UID))
|
||
w.uvarint(uint64(n.GID))
|
||
w.uvarint(uint64(n.Nlink))
|
||
w.uvarint(uint64(n.Ext))
|
||
}
|
||
w.uvarint(uint64(len(st.Errors)))
|
||
for _, e := range st.Errors {
|
||
w.uvarint(uint64(e.Node))
|
||
w.bytes([]byte{byte(e.Kind)})
|
||
w.str(e.Name)
|
||
w.str(e.Msg)
|
||
}
|
||
if w.err != nil {
|
||
return w.err
|
||
}
|
||
if err := gz.Close(); err != nil {
|
||
return err
|
||
}
|
||
bw.Write(h.Sum(nil))
|
||
return bw.Flush()
|
||
}
|
||
|
||
// SaveFile writes a snapshot atomically (temp file + rename).
|
||
func SaveFile(path string, t *inventory.Tree) error {
|
||
tmp, err := os.CreateTemp(dirOf(path), ".mapsize-*.tmp")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer os.Remove(tmp.Name())
|
||
if err := Save(tmp, t); err != nil {
|
||
tmp.Close()
|
||
return err
|
||
}
|
||
if err := tmp.Close(); err != nil {
|
||
return err
|
||
}
|
||
return os.Rename(tmp.Name(), path)
|
||
}
|
||
|
||
func dirOf(p string) string {
|
||
i := strings.LastIndexAny(p, `/\`)
|
||
if i < 0 {
|
||
return "."
|
||
}
|
||
if i == 0 {
|
||
return p[:1]
|
||
}
|
||
return p[:i]
|
||
}
|
||
|
||
type reader struct {
|
||
r *bufio.Reader
|
||
err error
|
||
}
|
||
|
||
func (r *reader) fail(format string, a ...any) {
|
||
if r.err == nil {
|
||
r.err = fmt.Errorf("snapshot: "+format, a...)
|
||
}
|
||
}
|
||
|
||
func (r *reader) byte1() byte {
|
||
if r.err != nil {
|
||
return 0
|
||
}
|
||
b, err := r.r.ReadByte()
|
||
if err != nil {
|
||
r.fail("truncated: %v", err)
|
||
}
|
||
return b
|
||
}
|
||
|
||
func (r *reader) uvarint(max uint64, what string) uint64 {
|
||
if r.err != nil {
|
||
return 0
|
||
}
|
||
v, err := readUvarint(r.r)
|
||
if err != nil {
|
||
r.fail("bad %s: %v", what, err)
|
||
return 0
|
||
}
|
||
if v > max {
|
||
r.fail("%s %d exceeds limit %d", what, v, max)
|
||
return 0
|
||
}
|
||
return v
|
||
}
|
||
|
||
func (r *reader) varint(what string) int64 {
|
||
if r.err != nil {
|
||
return 0
|
||
}
|
||
ux, err := readUvarint(r.r)
|
||
if err != nil {
|
||
r.fail("bad %s: %v", what, err)
|
||
}
|
||
v := int64(ux >> 1) // zig-zag, as encoding/binary.PutVarint
|
||
if ux&1 != 0 {
|
||
v = ^v
|
||
}
|
||
return v
|
||
}
|
||
|
||
var errOverflow = errors.New("varint overflows 64 bits")
|
||
|
||
// readUvarint is encoding/binary.ReadUvarint specialised to *bufio.Reader
|
||
// (no interface dispatch per byte; this is the loader's hot loop).
|
||
func readUvarint(br *bufio.Reader) (uint64, error) {
|
||
var x uint64
|
||
var s uint
|
||
for i := 0; i < binary.MaxVarintLen64; i++ {
|
||
b, err := br.ReadByte()
|
||
if err != nil {
|
||
if i > 0 && err == io.EOF {
|
||
err = io.ErrUnexpectedEOF
|
||
}
|
||
return x, err
|
||
}
|
||
if b < 0x80 {
|
||
if i == binary.MaxVarintLen64-1 && b > 1 {
|
||
return x, errOverflow
|
||
}
|
||
return x | uint64(b)<<s, nil
|
||
}
|
||
x |= uint64(b&0x7f) << s
|
||
s += 7
|
||
}
|
||
return x, errOverflow
|
||
}
|
||
|
||
func (r *reader) bytesN(n uint64) []byte {
|
||
if r.err != nil {
|
||
return nil
|
||
}
|
||
b := make([]byte, n)
|
||
if _, err := io.ReadFull(r.r, b); err != nil {
|
||
r.fail("truncated: %v", err)
|
||
return nil
|
||
}
|
||
return b
|
||
}
|
||
|
||
func (r *reader) str(max uint64, what string) string {
|
||
return string(r.bytesN(r.uvarint(max, what+" length")))
|
||
}
|
||
|
||
// Load reads and validates a snapshot.
|
||
func Load(in io.Reader) (*inventory.Tree, error) {
|
||
br := bufio.NewReaderSize(in, 1<<20)
|
||
hdr := make([]byte, len(brand.SnapshotMagic)+4)
|
||
if _, err := io.ReadFull(br, hdr); err != nil {
|
||
return nil, errors.New("snapshot: file too short")
|
||
}
|
||
if string(hdr[:8]) != brand.SnapshotMagic {
|
||
return nil, errors.New("snapshot: not a " + brand.Name + " snapshot")
|
||
}
|
||
if v := binary.LittleEndian.Uint16(hdr[8:]); v != Version {
|
||
return nil, fmt.Errorf("snapshot: unsupported version %d (this build reads %d)", v, Version)
|
||
}
|
||
if binary.LittleEndian.Uint16(hdr[10:])&flagGzipped == 0 {
|
||
return nil, errors.New("snapshot: uncompressed bodies are not supported")
|
||
}
|
||
gz, err := gzip.NewReader(br)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("snapshot: %w", err)
|
||
}
|
||
gz.Multistream(false)
|
||
// Decompression and hashing run on their own goroutine, overlapping
|
||
// with parsing. The pipe's close orders the hash and br accesses: the
|
||
// parser only sees EOF after the copier has finished with both.
|
||
h := sha256.New()
|
||
pr, pw := io.Pipe()
|
||
defer pr.Close() // unblocks the copier if we bail out early
|
||
go func() {
|
||
_, err := io.CopyBuffer(pw, io.TeeReader(gz, h), make([]byte, 256<<10))
|
||
pw.CloseWithError(err)
|
||
}()
|
||
r := &reader{r: bufio.NewReaderSize(pr, 1<<16)}
|
||
|
||
var m meta
|
||
if mb := r.bytesN(r.uvarint(MaxMetaLen, "metadata length")); r.err == nil {
|
||
if err := json.Unmarshal(mb, &m); err != nil {
|
||
return nil, fmt.Errorf("snapshot: metadata: %w", err)
|
||
}
|
||
}
|
||
if r.err != nil {
|
||
return nil, r.err
|
||
}
|
||
t := inventory.New(m.Root, inventory.KindDir)
|
||
nExt := r.uvarint(maxExts, "extension count")
|
||
exts := make([]string, 0, min(nExt, 1024))
|
||
for range nExt {
|
||
exts = append(exts, r.str(MaxExtLen, "extension"))
|
||
}
|
||
if r.err != nil {
|
||
return nil, r.err
|
||
}
|
||
extIndex := t.ExtIndexer()
|
||
nNodes := r.uvarint(MaxNodes, "node count")
|
||
if nNodes == 0 {
|
||
return nil, errors.New("snapshot: no root node")
|
||
}
|
||
for i := uint64(0); i < nNodes && r.err == nil; i++ {
|
||
pd := r.uvarint(MaxNodes, "parent")
|
||
kind := inventory.Kind(r.byte1())
|
||
cat := inventory.Category(r.byte1())
|
||
flags := inventory.Flags(r.uvarint(1<<16-1, "flags"))
|
||
name := r.str(MaxNameLen, "name")
|
||
size, alloc, mtime := r.varint("size"), r.varint("alloc"), r.varint("mtime")
|
||
mode := r.uvarint(1<<32-1, "mode")
|
||
uid := r.uvarint(1<<32-1, "uid")
|
||
gid := r.uvarint(1<<32-1, "gid")
|
||
nlink := r.uvarint(1<<32-1, "nlink")
|
||
ext := r.uvarint(nExt, "ext")
|
||
if r.err != nil {
|
||
break
|
||
}
|
||
if kind > inventory.KindOther || cat >= inventory.NumCategories {
|
||
r.fail("node %d: invalid kind/category", i)
|
||
break
|
||
}
|
||
if size < 0 || alloc < 0 {
|
||
r.fail("node %d: negative size", i)
|
||
break
|
||
}
|
||
var id inventory.NodeID
|
||
if i == 0 {
|
||
if pd != 0 || kind != inventory.KindDir && nNodes > 1 {
|
||
r.fail("invalid root")
|
||
break
|
||
}
|
||
id = 0
|
||
rn := t.Node(0)
|
||
rn.Kind = kind
|
||
} else {
|
||
if pd == 0 || pd > i {
|
||
r.fail("node %d: invalid parent reference", i)
|
||
break
|
||
}
|
||
parent := inventory.NodeID(i - pd)
|
||
if t.Node(parent).Kind != inventory.KindDir {
|
||
r.fail("node %d: parent is not a directory", i)
|
||
break
|
||
}
|
||
if !validName(name) {
|
||
r.fail("node %d: invalid name %q", i, name)
|
||
break
|
||
}
|
||
id = t.Add(parent, name, kind)
|
||
}
|
||
n := t.Node(id)
|
||
n.Cat, n.Flags = cat, flags
|
||
n.Size, n.Alloc, n.MTime = size, alloc, mtime
|
||
n.Mode, n.UID, n.GID, n.Nlink = uint32(mode), uint32(uid), uint32(gid), uint32(nlink)
|
||
if kind != inventory.KindDir && ext < uint64(len(exts)) {
|
||
n.Ext = extIndex(exts[ext])
|
||
}
|
||
}
|
||
nErr := r.uvarint(inventory.MaxErrorRecords, "error count")
|
||
for range nErr {
|
||
node := r.uvarint(nNodes-1, "error node")
|
||
kind := inventory.ErrKind(r.byte1())
|
||
name := r.str(MaxNameLen, "error name")
|
||
msg := r.str(MaxMsgLen, "error message")
|
||
if r.err != nil {
|
||
break
|
||
}
|
||
if kind >= inventory.NumErrKinds {
|
||
kind = inventory.ErrOther
|
||
}
|
||
t.Stats.Errors = append(t.Stats.Errors, inventory.ErrorRecord{Node: inventory.NodeID(node), Kind: kind, Name: name, Msg: msg})
|
||
}
|
||
if r.err != nil {
|
||
return nil, r.err
|
||
}
|
||
// The body must end here, then the checksum follows the gzip stream.
|
||
if _, err := r.r.ReadByte(); err != io.EOF {
|
||
return nil, errors.New("snapshot: trailing data in body")
|
||
}
|
||
sum := make([]byte, sha256.Size)
|
||
if _, err := io.ReadFull(br, sum); err != nil {
|
||
return nil, errors.New("snapshot: missing checksum")
|
||
}
|
||
if !bytes.Equal(sum, h.Sum(nil)) {
|
||
return nil, errors.New("snapshot: checksum mismatch (file corrupt)")
|
||
}
|
||
|
||
t.Recompute()
|
||
st := &t.Stats
|
||
st.Start, st.End, st.Complete, st.Cancelled = m.Start, m.End, true, m.Cancelled
|
||
st.Files, st.Dirs, st.Symlinks, st.Others = m.Files, m.Dirs, m.Symlinks, m.Others
|
||
copy(st.ErrCounts[:], m.ErrCounts)
|
||
st.Excluded, st.SkippedMounts, st.VirtualSkipped = m.Excluded, m.SkippedMounts, m.VirtualSkipped
|
||
st.LoopsSkipped, st.BrokenLinks, st.HardlinkDups, st.Unscanned = m.LoopsSkipped, m.BrokenLinks, m.HardlinkDups, m.Unscanned
|
||
st.Excludes, st.OneFileSystem, st.Follow, st.Workers = m.Excludes, m.OneFileSystem, m.Follow, m.Workers
|
||
st.FromSnapshot = fmt.Sprintf("%s (%s, %s)", m.Created.Local().Format("2006-01-02 15:04"), m.Host, m.Generator)
|
||
return t, nil
|
||
}
|
||
|
||
func validName(s string) bool {
|
||
return s != "" && s != "." && s != ".." && !strings.ContainsAny(s, "/\x00")
|
||
}
|
||
|
||
// LoadFile opens and loads a snapshot file.
|
||
func LoadFile(path string) (*inventory.Tree, error) {
|
||
f, err := os.Open(path)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer f.Close()
|
||
return Load(f)
|
||
}
|