Bound download expansion and escape control chars on the terminal (closes #164) #197

Merged
clawbot merged 1 commits from issue-164-bound-download-expansion into next 2026-09-22 17:00:36 +02:00
14 changed files with 431 additions and 87 deletions
+49
View File
@@ -0,0 +1,49 @@
package blobgen
import (
"errors"
"io"
)
// ErrOutputTooLarge is returned by a reader from LimitReader once it has
// been asked for more than its limit. It bounds how far an untrusted
// compressed stream may expand, so a small, highly compressible object
// from the store cannot decompress without limit.
var ErrOutputTooLarge = errors.New("output exceeds size limit")
// LimitReader returns a reader that yields at most limit bytes from r and
// then fails with ErrOutputTooLarge. Unlike io.LimitReader, which reports
// a silent io.EOF at the limit (indistinguishable from a stream that
// simply ended), this fails, so a caller decoding or copying the stream
// sees an error rather than a truncated value. A stream of exactly limit
// bytes reads back cleanly to EOF; the first byte beyond it is the error.
func LimitReader(r io.Reader, limit int64) io.Reader {
// remaining counts down from limit+1: the extra byte is the one that,
// if it ever arrives, proves the stream is longer than the limit.
return &limitReader{r: r, remaining: limit + 1}
}
type limitReader struct {
r io.Reader
remaining int64
}
func (l *limitReader) Read(p []byte) (int, error) {
if l.remaining <= 0 {
return 0, ErrOutputTooLarge
}
if int64(len(p)) > l.remaining {
p = p[:l.remaining]
}
n, err := l.r.Read(p)
l.remaining -= int64(n)
if l.remaining <= 0 {
// The (limit+1)th byte was just read: the stream is too long.
return n, ErrOutputTooLarge
}
return n, err
}
+43
View File
@@ -0,0 +1,43 @@
package blobgen_test
import (
"bytes"
"io"
"testing"
"github.com/stretchr/testify/require"
"sneak.berlin/go/vaultik/internal/blobgen"
)
// TestLimitReaderPassesExactSize checks that a stream of exactly the limit
// reads back cleanly to EOF: the bound must not reject a legitimate blob
// whose plaintext equals its recorded size.
func TestLimitReaderPassesExactSize(t *testing.T) {
t.Parallel()
const n = 1000
r := blobgen.LimitReader(bytes.NewReader(bytes.Repeat([]byte("a"), n)), n)
got, err := io.ReadAll(r)
require.NoError(t, err)
require.Len(t, got, n)
}
// TestLimitReaderFailsPastLimit feeds a large, highly compressible run of
// zeros — the decompressed output a zip bomb would produce — through a
// small limit and checks it fails within the bound rather than passing
// the whole stream through.
func TestLimitReaderFailsPastLimit(t *testing.T) {
t.Parallel()
const limit = 1000
r := blobgen.LimitReader(
bytes.NewReader(bytes.Repeat([]byte{0}, limit*1000)), limit)
n, err := io.Copy(io.Discard, r)
require.ErrorIs(t, err, blobgen.ErrOutputTooLarge)
require.LessOrEqual(t, n, int64(limit)+1,
"reader must stop within one byte of the limit")
}
+39
View File
@@ -0,0 +1,39 @@
package log_test
import (
"bytes"
"log/slog"
"strings"
"testing"
"github.com/stretchr/testify/require"
"sneak.berlin/go/vaultik/internal/log"
)
// TestTTYHandlerEscapesControlCharacters logs a message and an attribute
// value that each carry an ESC and a newline — the shape a crafted path or
// storage error from the destination would take — and checks neither raw
// byte reaches the output. The handler's own colour codes (ESC ... m) are
// stripped first; any ESC left after that came from the untrusted value.
func TestTTYHandlerEscapesControlCharacters(t *testing.T) {
t.Parallel()
var buf bytes.Buffer
logger := slog.New(log.NewTTYHandler(&buf, debugHandlerOptions()))
logger.Info("start\x1b[31mZAP\nend", "target", "a\x1b[31mZAP\nb")
out := buf.String()
// The only newline is the line terminator; the injected ones were escaped.
require.Equal(t, 1, strings.Count(out, "\n"),
"a newline in the message or a value must be escaped, not emitted raw")
// After the handler's own colour codes are removed, no ESC survives.
stripped := ansiEscape.ReplaceAllString(out, "")
require.NotContains(t, stripped, "\x1b",
"a raw ESC from the message or a value must not reach the terminal")
// The escaped form is what appears instead.
require.Contains(t, out, `\x1b`)
}
+29 -4
View File
@@ -5,9 +5,11 @@ import (
"fmt" "fmt"
"io" "io"
"log/slog" "log/slog"
"strconv"
"strings" "strings"
"sync" "sync"
"time" "time"
"unicode"
) )
// groupSeparator joins an open group path to an attribute key. This // groupSeparator joins an open group path to an attribute key. This
@@ -116,11 +118,14 @@ func (h *TTYHandler) Handle(_ context.Context, r slog.Record) error {
levelColor = colorReset levelColor = colorReset
} }
// Print main message // Print main message. The message is escaped before the colour codes
// are written around it: it can carry text from an untrusted source
// (a storage error, for one), and a raw control character would
// otherwise reach the terminal.
_, _ = fmt.Fprintf(h.out, "%s%s%s %s%s%s %s%s%s", _, _ = fmt.Fprintf(h.out, "%s%s%s %s%s%s %s%s%s",
colorGray, timestamp, colorReset, colorGray, timestamp, colorReset,
levelColor, level, colorReset, levelColor, level, colorReset,
colorBold, r.Message, colorReset) colorBold, sanitize(r.Message), colorReset)
// Attributes carried by the handler come first, then the record's // Attributes carried by the handler come first, then the record's
// own. Handler attributes were qualified when they were added; the // own. Handler attributes were qualified when they were added; the
@@ -260,9 +265,29 @@ func (h *TTYHandler) writeAttr(a slog.Attr) {
// Future kinds also use the plain string form. // Future kinds also use the plain string form.
} }
// Escape the key and value before the colour codes are written around
// them. Both can carry text from an untrusted source — a manifest
// timestamp, a storage error, a path or symlink target read back from
// the snapshot database — so a control character in one of them must
// be rendered as an escape sequence rather than reaching the terminal,
// where it could move the cursor or inject its own colours.
_, _ = fmt.Fprintf(h.out, " %s%s%s=%s%s%s", _, _ = fmt.Fprintf(h.out, " %s%s%s=%s%s%s",
colorCyan, a.Key, colorReset, colorCyan, sanitize(a.Key), colorReset,
colorBlue, value, colorReset) colorBlue, sanitize(value), colorReset)
}
// sanitize returns s unchanged when every rune in it is printable, and a
// double-quoted, backslash-escaped form (\n, \x1b, …) otherwise. It is
// applied to untrusted text before any colour code is written, so a
// control character can never reach the terminal raw.
func sanitize(s string) string {
for _, r := range s {
if !unicode.IsPrint(r) {
return strconv.Quote(s)
}
}
return s
} }
// formatDuration formats a duration in a human-readable way // formatDuration formats a duration in a human-readable way
+30 -5
View File
@@ -7,6 +7,19 @@ import (
"io" "io"
"github.com/klauspost/compress/zstd" "github.com/klauspost/compress/zstd"
"sneak.berlin/go/vaultik/internal/blobgen"
)
// Manifest size bounds. A manifest lists one small entry per blob, and
// blobs are large (the default target is 10 GB), so even a manifest for a
// petabyte-scale backup is a few megabytes. These caps are far above any
// manifest the writer can emit, yet stop a crafted, highly compressible
// manifest from expanding without limit when decoded: the manifest is
// fetched from the store, which is not trusted, and json.Decode buffers
// the whole value in memory.
const (
manifestMaxCompressed = 256 * 1024 * 1024 // 256 MiB
manifestMaxDecompressed = 1024 * 1024 * 1024 // 1 GiB
) )
// Manifest represents the structure of a snapshot's blob manifest // Manifest represents the structure of a snapshot's blob manifest
@@ -28,19 +41,31 @@ type BlobInfo struct {
CompressedSize int64 `json:"compressed_size"` CompressedSize int64 `json:"compressed_size"`
} }
// DecodeManifest decodes a manifest from a reader containing compressed JSON // DecodeManifest decodes a manifest from a reader containing compressed
// JSON, reading through byte limits on both the compressed input and the
// decompressed output so an untrusted manifest cannot exhaust memory.
func DecodeManifest(r io.Reader) (*Manifest, error) { func DecodeManifest(r io.Reader) (*Manifest, error) {
// Decompress using zstd return decodeManifest(r, manifestMaxCompressed, manifestMaxDecompressed)
zr, err := zstd.NewReader(r) }
// decodeManifest is DecodeManifest with explicit limits, so tests can drive
// the bounds with small inputs instead of gigabyte-scale ones.
func decodeManifest(
r io.Reader, maxCompressed, maxDecompressed int64,
) (*Manifest, error) {
// Decompress using zstd, bounding how many compressed bytes are read.
zr, err := zstd.NewReader(blobgen.LimitReader(r, maxCompressed))
if err != nil { if err != nil {
return nil, fmt.Errorf("creating zstd reader: %w", err) return nil, fmt.Errorf("creating zstd reader: %w", err)
} }
defer zr.Close() defer zr.Close()
// Decode JSON manifest // Decode JSON manifest, bounding how far the compressed input may
// expand: json.Decode buffers the whole value, so without this a
// small, highly compressible manifest could expand to gigabytes.
var manifest Manifest var manifest Manifest
err = json.NewDecoder(zr).Decode(&manifest) err = json.NewDecoder(blobgen.LimitReader(zr, maxDecompressed)).Decode(&manifest)
if err != nil { if err != nil {
return nil, fmt.Errorf("decoding manifest: %w", err) return nil, fmt.Errorf("decoding manifest: %w", err)
} }
+79
View File
@@ -0,0 +1,79 @@
//nolint:testpackage // exercises the unexported decodeManifest bounds
package snapshot
import (
"bytes"
"strings"
"testing"
"github.com/stretchr/testify/require"
"sneak.berlin/go/vaultik/internal/blobgen"
)
// testSnapshotID is a stand-in snapshot ID reused across the bound cases.
const testSnapshotID = "host_home_2026-01-01T00:00:00Z"
// TestDecodeManifestRoundTrip is the baseline: with generous bounds a
// manifest the writer produced decodes back unchanged.
func TestDecodeManifestRoundTrip(t *testing.T) {
t.Parallel()
want := &Manifest{
SnapshotID: testSnapshotID,
Timestamp: "2026-01-01T00:00:00Z",
BlobCount: 2,
TotalCompressedSize: 42,
Blobs: []BlobInfo{
{Hash: "aa", CompressedSize: 21},
{Hash: "bb", CompressedSize: 21},
},
}
compressed, err := EncodeManifest(want, 3)
require.NoError(t, err)
got, err := decodeManifest(
bytes.NewReader(compressed), manifestMaxCompressed, manifestMaxDecompressed)
require.NoError(t, err)
require.Equal(t, want, got)
}
// TestDecodeManifestBoundsDecompressedOutput feeds a valid but highly
// compressible manifest — one whose timestamp is a megabyte of the same
// character — through a small decompressed bound. The compressed form is
// tiny, so only the decompressed bound stops it; decoding must fail within
// that bound rather than expanding the value in memory.
func TestDecodeManifestBoundsDecompressedOutput(t *testing.T) {
t.Parallel()
bomb := &Manifest{
SnapshotID: testSnapshotID,
Timestamp: strings.Repeat("a", 1<<20),
}
compressed, err := EncodeManifest(bomb, 3)
require.NoError(t, err)
require.Less(t, len(compressed), 4096,
"the compressible manifest must be small compressed")
_, err = decodeManifest(bytes.NewReader(compressed), 1<<20, 4096)
require.ErrorIs(t, err, blobgen.ErrOutputTooLarge)
}
// TestDecodeManifestBoundsCompressedInput checks the compressed-input
// bound fires independently: a valid manifest with a generous decompressed
// bound but a tiny compressed bound still fails.
func TestDecodeManifestBoundsCompressedInput(t *testing.T) {
t.Parallel()
manifest := &Manifest{
SnapshotID: testSnapshotID,
Timestamp: strings.Repeat("a", 4096),
}
compressed, err := EncodeManifest(manifest, 3)
require.NoError(t, err)
_, err = decodeManifest(bytes.NewReader(compressed), 8, manifestMaxDecompressed)
require.Error(t, err)
}
+22 -3
View File
@@ -23,7 +23,9 @@ import (
"fmt" "fmt"
"io" "io"
"os" "os"
"strconv"
"time" "time"
"unicode"
"github.com/dustin/go-humanize" "github.com/dustin/go-humanize"
"golang.org/x/term" "golang.org/x/term"
@@ -225,17 +227,17 @@ func (w *Writer) Hex(s string) string {
short = s[:hexAbbrevLen] + "..." short = s[:hexAbbrevLen] + "..."
} }
return w.paint(ansiCyan, short) return w.paint(ansiCyan, sanitize(short))
} }
// Snapshot colorizes a snapshot ID (full, no abbreviation). // Snapshot colorizes a snapshot ID (full, no abbreviation).
func (w *Writer) Snapshot(id string) string { func (w *Writer) Snapshot(id string) string {
return w.paint(ansiCyan+ansiBold, id) return w.paint(ansiCyan+ansiBold, sanitize(id))
} }
// Path colorizes a filesystem path. // Path colorizes a filesystem path.
func (w *Writer) Path(p string) string { func (w *Writer) Path(p string) string {
return w.paint(ansiBlue, p) return w.paint(ansiBlue, sanitize(p))
} }
// Size colorizes a byte count using humanize.Bytes. // Size colorizes a byte count using humanize.Bytes.
@@ -310,6 +312,23 @@ func (w *Writer) paint(color, s string) string {
return color + s + ansiReset return color + s + ansiReset
} }
// sanitize returns s unchanged when every rune in it is printable, and a
// double-quoted, backslash-escaped form (\n, \x1b, …) otherwise. The
// string value formatters escape their argument through this before
// painting: identifiers, paths and symlink targets they render come from
// the snapshot database, which is not trusted, and escaping must happen
// before colour is applied — the painted result already contains the
// escape codes the raw text would otherwise be indistinguishable from.
func sanitize(s string) string {
for _, r := range s {
if !unicode.IsPrint(r) {
return strconv.Quote(s)
}
}
return s
}
// emit writes "<prefix> <body>\n" with the prefix painted in prefixColor // emit writes "<prefix> <body>\n" with the prefix painted in prefixColor
// and the body optionally painted in bodyColor (empty = no body color). // and the body optionally painted in bodyColor (empty = no body color).
func (w *Writer) emit(prefixColor, prefix, bodyColor, format string, args []any) { func (w *Writer) emit(prefixColor, prefix, bodyColor, format string, args []any) {
+32
View File
@@ -0,0 +1,32 @@
package ui_test
import (
"strings"
"testing"
)
// TestValueFormattersEscapeControlCharacters checks that a path carrying an
// ESC and a newline — the shape a symlink target read back from the
// snapshot database could take — is escaped before it reaches the output.
// Colour is off here, so the only way a control byte could appear is from
// the value itself.
func TestValueFormattersEscapeControlCharacters(t *testing.T) {
t.Parallel()
w, buf := newTestWriter(false)
w.Infof("restoring %s", w.Path("a\x1b[31mZAP\nb"))
out := buf.String()
if strings.ContainsRune(out, '\x1b') {
t.Fatalf("raw ESC from a value survived in output: %q", out)
}
if strings.Count(out, "\n") != 1 {
t.Fatalf("a newline in a value must be escaped, not emitted raw: %q", out)
}
if !strings.Contains(out, `\x1b`) {
t.Fatalf("expected the escaped form of ESC in output: %q", out)
}
}
+16 -34
View File
@@ -6,11 +6,9 @@ import (
"errors" "errors"
"fmt" "fmt"
"io" "io"
"time"
"filippo.io/age" "filippo.io/age"
"sneak.berlin/go/vaultik/internal/blobgen" "sneak.berlin/go/vaultik/internal/blobgen"
"sneak.berlin/go/vaultik/internal/log"
) )
// errBlobHashMismatch is returned when a fetched blob's content hash does // errBlobHashMismatch is returned when a fetched blob's content hash does
@@ -30,13 +28,14 @@ var errBlobNotFullyRead = errors.New(
// redundant SHA-256 computation. // redundant SHA-256 computation.
type hashVerifyReader struct { type hashVerifyReader struct {
reader *blobgen.Reader // underlying decrypted blob reader (has internal hasher) reader *blobgen.Reader // underlying decrypted blob reader (has internal hasher)
limited io.Reader // reader bounded to the blob's recorded plaintext size
fetcher io.ReadCloser // raw fetched stream (closed on Close) fetcher io.ReadCloser // raw fetched stream (closed on Close)
blobHash string // expected double-SHA-256 hex blobHash string // expected double-SHA-256 hex
done bool // EOF reached done bool // EOF reached
} }
func (h *hashVerifyReader) Read(p []byte) (int, error) { func (h *hashVerifyReader) Read(p []byte) (int, error) {
n, err := h.reader.Read(p) n, err := h.limited.Read(p)
if errors.Is(err, io.EOF) { if errors.Is(err, io.EOF) {
h.done = true h.done = true
} }
@@ -73,11 +72,17 @@ func (h *hashVerifyReader) Close() error {
// returns a streaming reader that computes the double-SHA-256 hash on the fly. // returns a streaming reader that computes the double-SHA-256 hash on the fly.
// The hash is verified when the returned reader is closed (after fully reading). // The hash is verified when the returned reader is closed (after fully reading).
// This avoids buffering the entire blob in memory. // This avoids buffering the entire blob in memory.
//
// maxPlaintextSize is the blob's uncompressed_size as recorded in the
// snapshot database. Decompression stops with blobgen.ErrOutputTooLarge
// once the plaintext exceeds it, so a tampered blob cannot expand without
// limit — using the recorded size, not the restoring host's
// blob_size_limit, since that config may differ from the backup host's.
func (v *Vaultik) FetchAndDecryptBlob( func (v *Vaultik) FetchAndDecryptBlob(
ctx context.Context, blobHash string, expectedSize int64, ctx context.Context, blobHash string, maxPlaintextSize int64,
identities ...age.Identity, identities ...age.Identity,
) (io.ReadCloser, error) { ) (io.ReadCloser, error) {
rc, _, err := v.FetchBlob(ctx, blobHash, expectedSize) rc, err := v.FetchBlob(ctx, blobHash)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -91,52 +96,29 @@ func (v *Vaultik) FetchAndDecryptBlob(
return &hashVerifyReader{ return &hashVerifyReader{
reader: reader, reader: reader,
limited: blobgen.LimitReader(reader, maxPlaintextSize),
fetcher: rc, fetcher: rc,
blobHash: blobHash, blobHash: blobHash,
}, nil }, nil
} }
// FetchBlob downloads a blob and returns a reader for the encrypted data. // FetchBlob downloads a blob and returns a reader for the encrypted data.
// Times the Storage.Get and Storage.Stat round-trips separately at
// debug level so we can see whether the size-only Stat (which is an
// extra request on every fetch) is hurting throughput.
func (v *Vaultik) FetchBlob( func (v *Vaultik) FetchBlob(
ctx context.Context, blobHash string, expectedSize int64, ctx context.Context, blobHash string,
) (io.ReadCloser, int64, error) { ) (io.ReadCloser, error) {
// blobHash reaches here from the snapshot database, which is not // blobHash reaches here from the snapshot database, which is not
// trusted. Reject a malformed hash before it is spliced into a storage // trusted. Reject a malformed hash before it is spliced into a storage
// path (blobHash[:2]/blobHash[2:4]) or a fetch is attempted. // path (blobHash[:2]/blobHash[2:4]) or a fetch is attempted.
if !isBlobHash(blobHash) { if !isBlobHash(blobHash) {
return nil, 0, fmt.Errorf("%w: %s", errInvalidBlobHash, shortHash(blobHash)) return nil, fmt.Errorf("%w: %s", errInvalidBlobHash, shortHash(blobHash))
} }
blobPath := fmt.Sprintf("blobs/%s/%s/%s", blobHash[:2], blobHash[2:4], blobHash) blobPath := fmt.Sprintf("blobs/%s/%s/%s", blobHash[:2], blobHash[2:4], blobHash)
t0 := time.Now()
rc, err := v.Storage.Get(ctx, blobPath) rc, err := v.Storage.Get(ctx, blobPath)
getDur := time.Since(t0)
if err != nil { if err != nil {
return nil, 0, fmt.Errorf("downloading blob %s: %w", shortHash(blobHash), err) return nil, fmt.Errorf("downloading blob %s: %w", shortHash(blobHash), err)
} }
t0 = time.Now() return rc, nil
info, err := v.Storage.Stat(ctx, blobPath)
statDur := time.Since(t0)
if err != nil {
_ = rc.Close()
return nil, 0, fmt.Errorf("stat blob %s: %w", shortHash(blobHash), err)
}
log.Debug("FetchBlob round-trips",
"hash", shortHash(blobHash),
"ms_storage_get", getDur.Milliseconds(),
"ms_storage_stat", statDur.Milliseconds(),
"expected_size", expectedSize,
"stat_size", info.Size,
)
return rc, info.Size, nil
} }
+50
View File
@@ -0,0 +1,50 @@
package vaultik_test
import (
"context"
"io"
"testing"
"filippo.io/age"
"github.com/stretchr/testify/require"
"sneak.berlin/go/vaultik/internal/blobgen"
"sneak.berlin/go/vaultik/internal/vaultik"
)
// TestFetchAndDecryptBlobBoundsPlaintext feeds a small, highly
// compressible blob (256 KiB of zeros) whose decompressed size far exceeds
// the plaintext bound passed to FetchAndDecryptBlob. Decompression must
// stop with blobgen.ErrOutputTooLarge within the bound rather than
// expanding the whole blob into the restore cache.
func TestFetchAndDecryptBlobBoundsPlaintext(t *testing.T) {
t.Parallel()
identity, err := age.GenerateX25519Identity()
require.NoError(t, err)
plaintext := make([]byte, 256*1024)
encryptedData, correctHash := buildHashTestBlob(t, identity, plaintext)
mockStorage := NewMockStorer()
blobPath := "blobs/" + correctHash[:2] + "/" +
correctHash[2:4] + "/" + correctHash
mockStorage.mu.Lock()
mockStorage.data[blobPath] = encryptedData
mockStorage.mu.Unlock()
tv := vaultik.NewForTesting(mockStorage)
const maxPlaintext = 1024
rc, err := tv.FetchAndDecryptBlob(
context.Background(), correctHash, maxPlaintext, identity)
require.NoError(t, err)
n, copyErr := io.Copy(io.Discard, rc)
_ = rc.Close()
require.ErrorIs(t, copyErr, blobgen.ErrOutputTooLarge)
require.LessOrEqual(t, n, int64(maxPlaintext)+1,
"decompression must stop within the recorded plaintext bound")
}
+4 -4
View File
@@ -73,7 +73,7 @@ func TestFetchBlobRejectsMalformedHash(t *testing.T) {
strings.Repeat("A", 64), // uppercase hex is not accepted strings.Repeat("A", 64), // uppercase hex is not accepted
strings.Repeat("g", 64), // not hex strings.Repeat("g", 64), // not hex
} { } {
_, _, err := tv.FetchBlob(ctx, bad, 0) _, err := tv.FetchBlob(ctx, bad)
if err == nil { if err == nil {
t.Fatalf("expected error for malformed hash %q, got nil", bad) t.Fatalf("expected error for malformed hash %q, got nil", bad)
} }
@@ -113,7 +113,7 @@ func TestFetchAndDecryptBlobVerifiesHash(t *testing.T) {
t.Parallel() t.Parallel()
rc, err := tv.FetchAndDecryptBlob( rc, err := tv.FetchAndDecryptBlob(
ctx, correctHash, int64(len(encryptedData)), identity) ctx, correctHash, int64(len(plaintext)), identity)
if err != nil { if err != nil {
t.Fatalf("expected success, got error: %v", err) t.Fatalf("expected success, got error: %v", err)
} }
@@ -145,7 +145,7 @@ func TestFetchAndDecryptBlobVerifiesHash(t *testing.T) {
mockStorage.mu.Unlock() mockStorage.mu.Unlock()
rc, err := tv.FetchAndDecryptBlob( rc, err := tv.FetchAndDecryptBlob(
ctx, fakeHash, int64(len(encryptedData)), identity) ctx, fakeHash, int64(len(plaintext)), identity)
if err != nil { if err != nil {
t.Fatalf("unexpected error opening stream: %v", err) t.Fatalf("unexpected error opening stream: %v", err)
} }
@@ -188,7 +188,7 @@ func TestFetchAndDecryptBlobCloseBeforeEOFFails(t *testing.T) {
tv := vaultik.NewForTesting(mockStorage) tv := vaultik.NewForTesting(mockStorage)
rc, err := tv.FetchAndDecryptBlob( rc, err := tv.FetchAndDecryptBlob(
context.Background(), correctHash, int64(len(encryptedData)), identity) context.Background(), correctHash, int64(len(plaintext)), identity)
if err != nil { if err != nil {
t.Fatalf("unexpected error opening stream: %v", err) t.Fatalf("unexpected error opening stream: %v", err)
} }
+33 -34
View File
@@ -1,7 +1,6 @@
package vaultik package vaultik
import ( import (
"bytes"
"context" "context"
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
@@ -428,7 +427,7 @@ func (s *restoreSession) downloadNextBlobSet(plan *restorePlan) (bool, error) {
return false, fmt.Errorf("%w: %s", errBlobMissingFromIndex, shortHash(hash)) return false, fmt.Errorf("%w: %s", errBlobMissingFromIndex, shortHash(hash))
} }
err := s.downloadBlobToCache(hash, blob.CompressedSize) err := s.downloadBlobToCache(hash, blob.CompressedSize, blob.UncompressedSize)
if err != nil { if err != nil {
return false, fmt.Errorf("downloading blob %s: %w", shortHash(hash), err) return false, fmt.Errorf("downloading blob %s: %w", shortHash(hash), err)
} }
@@ -655,32 +654,18 @@ func (v *Vaultik) downloadSnapshotDB(
defer func() { _ = reader.Close() }() defer func() { _ = reader.Close() }()
// Read all data // Decrypt and decompress straight from the storage stream, then stream
encryptedData, err := io.ReadAll(reader) // the plaintext to a temp file. Neither the encrypted bytes nor the
if err != nil { // decrypted database is ever held whole in memory; a snapshot database
return nil, "", fmt.Errorf("reading encrypted data: %w", err) // can be large.
} blobReader, err := blobgen.NewReader(reader, identities...)
log.Debug("Downloaded encrypted database",
"size", ubytes(int64(len(encryptedData))))
// Decrypt and decompress using blobgen.Reader
blobReader, err := blobgen.NewReader(bytes.NewReader(encryptedData), identities...)
if err != nil { if err != nil {
return nil, "", fmt.Errorf("creating decryption reader: %w", err) return nil, "", fmt.Errorf("creating decryption reader: %w", err)
} }
defer func() { _ = blobReader.Close() }() defer func() { _ = blobReader.Close() }()
// Read the binary SQLite database db, tempDir, err := v.materializeSnapshotDB(blobReader)
dbData, err := io.ReadAll(blobReader)
if err != nil {
return nil, "", fmt.Errorf("decrypting and decompressing: %w", err)
}
log.Debug("Decrypted database", "size", ubytes(int64(len(dbData))))
db, tempDir, err := v.materializeSnapshotDB(dbData)
if err != nil { if err != nil {
return nil, "", err return nil, "", err
} }
@@ -729,13 +714,15 @@ func (v *Vaultik) verifySnapshotDBIdentity(
return nil return nil
} }
// materializeSnapshotDB writes the decrypted snapshot database bytes into // materializeSnapshotDB streams the decrypted snapshot database into a
// a fresh private (0700) temp directory and opens the file read-only. On // fresh private (0700) temp directory and opens the file read-only. The
// any failure it removes the directory before returning, so no decrypted // database is copied through an io.Copy buffer rather than read whole into
// metadata is left on disk when the open is interrupted or the payload is // memory. On any failure it removes the directory before returning, so no
// damaged. On success the returned directory is the caller's to remove. // decrypted metadata is left on disk when the copy is interrupted or the
// payload is damaged. On success the returned directory is the caller's to
// remove.
func (v *Vaultik) materializeSnapshotDB( func (v *Vaultik) materializeSnapshotDB(
dbData []byte, dbReader io.Reader,
) (*database.DB, string, error) { ) (*database.DB, string, error) {
tempDir, err := afero.TempDir(v.Fs, "", "vaultik-restore-") tempDir, err := afero.TempDir(v.Fs, "", "vaultik-restore-")
if err != nil { if err != nil {
@@ -752,12 +739,24 @@ func (v *Vaultik) materializeSnapshotDB(
dbPath := filepath.Join(tempDir, snapshotDBFilename) dbPath := filepath.Join(tempDir, snapshotDBFilename)
err = afero.WriteFile(v.Fs, dbPath, dbData, restoreFileMode) dbFile, err := v.Fs.OpenFile(
dbPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, restoreFileMode)
if err != nil { if err != nil {
return nil, "", fmt.Errorf("writing database file: %w", err) return nil, "", fmt.Errorf("creating database file: %w", err)
} }
log.Debug("Created restore database", "path", dbPath) written, copyErr := io.Copy(dbFile, dbReader)
closeErr := dbFile.Close()
if copyErr != nil {
return nil, "", fmt.Errorf("writing database file: %w", copyErr)
}
if closeErr != nil {
return nil, "", fmt.Errorf("closing database file: %w", closeErr)
}
log.Debug("Created restore database", "path", dbPath, "size", ubytes(written))
db, err := database.OpenReadOnly(v.ctx, dbPath) db, err := database.OpenReadOnly(v.ctx, dbPath)
if err != nil { if err != nil {
@@ -1212,12 +1211,12 @@ func (s *restoreSession) writeFileChunks(
// size, which is what makes multi-GB blobs tractable on machines with // size, which is what makes multi-GB blobs tractable on machines with
// less RAM than the blob. // less RAM than the blob.
func (s *restoreSession) downloadBlobToCache( func (s *restoreSession) downloadBlobToCache(
blobHash string, expectedSize int64, blobHash string, compressedSize, uncompressedSize int64,
) error { ) error {
start := time.Now() start := time.Now()
t0 := time.Now() t0 := time.Now()
rc, err := s.v.FetchAndDecryptBlob(s.ctx, blobHash, expectedSize, s.identities...) rc, err := s.v.FetchAndDecryptBlob(s.ctx, blobHash, uncompressedSize, s.identities...)
fetchSetupDur := time.Since(t0) fetchSetupDur := time.Since(t0)
if err != nil { if err != nil {
@@ -1247,7 +1246,7 @@ func (s *restoreSession) downloadBlobToCache(
log.Debug("Streamed blob into disk cache", log.Debug("Streamed blob into disk cache",
"hash", blobHash[:16], "hash", blobHash[:16],
"compressed_bytes", expectedSize, "compressed_bytes", compressedSize,
"plaintext_bytes", written, "plaintext_bytes", written,
"ms_total", time.Since(start).Milliseconds(), "ms_total", time.Since(start).Milliseconds(),
"ms_fetch_setup", fetchSetupDur.Milliseconds(), "ms_fetch_setup", fetchSetupDur.Milliseconds(),
+4 -2
View File
@@ -1,6 +1,7 @@
package vaultik //nolint:testpackage // inspects unexported snapshot-db materialization package vaultik //nolint:testpackage // inspects unexported snapshot-db materialization
import ( import (
"bytes"
"context" "context"
"os" "os"
"path/filepath" "path/filepath"
@@ -37,7 +38,7 @@ func TestMaterializeSnapshotDBPrivateDir(t *testing.T) {
v := &Vaultik{ctx: context.Background(), Fs: afero.NewOsFs()} v := &Vaultik{ctx: context.Background(), Fs: afero.NewOsFs()}
db, dir, err := v.materializeSnapshotDB(dbData) db, dir, err := v.materializeSnapshotDB(bytes.NewReader(dbData))
require.NoError(t, err) require.NoError(t, err)
t.Cleanup(func() { t.Cleanup(func() {
@@ -64,7 +65,8 @@ func TestMaterializeSnapshotDBRemovesDirOnOpenFailure(t *testing.T) {
v := &Vaultik{ctx: context.Background(), Fs: afero.NewOsFs()} v := &Vaultik{ctx: context.Background(), Fs: afero.NewOsFs()}
_, _, err := v.materializeSnapshotDB([]byte("this is not a sqlite database")) _, _, err := v.materializeSnapshotDB(
bytes.NewReader([]byte("this is not a sqlite database")))
require.Error(t, err) require.Error(t, err)
entries, rerr := os.ReadDir(base) entries, rerr := os.ReadDir(base)
+1 -1
View File
@@ -393,7 +393,7 @@ func (v *Vaultik) verifyBlob(
blobInfo snapshot.BlobInfo, db *sql.DB, identities []age.Identity, blobInfo snapshot.BlobInfo, db *sql.DB, identities []age.Identity,
) error { ) error {
// Download blob using shared fetch method // Download blob using shared fetch method
reader, _, err := v.FetchBlob(v.ctx, blobInfo.Hash, blobInfo.CompressedSize) reader, err := v.FetchBlob(v.ctx, blobInfo.Hash)
if err != nil { if err != nil {
return fmt.Errorf("failed to download: %w", err) return fmt.Errorf("failed to download: %w", err)
} }