mesh-controller/internal/builder/xz_framing_test.go

package builder

import (
	"archive/tar"
	"bufio"
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"hash/crc32"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
)

// xzStream is one xz stream (no integrity check) of data cut into blocks of blockSize bytes, each block's
// LZMA2 data stored uncompressed, each block a delta filter before LZMA2 when delta is set: the shape of
// `xz --block-size=N --filters=delta,lzma2`, built without any xz tool. listed is how many blocks its index
// lists (-1: as many as it holds), so a stream whose index lies can be built too.
func xzStream(data []byte, blockSize int, delta bool, listed int) []byte {
	var out bytes.Buffer
	flags := []byte{0, 0}
	out.Write([]byte{0xfd, '7', 'z', 'X', 'Z', 0x00})
	out.Write(flags)
	out.Write(binary.LittleEndian.AppendUint32(nil, crc32.ChecksumIEEE(flags)))
	type record struct{ unpadded, uncompressed uint64 }
	var records []record
	for at := 0; at < len(data); at += blockSize {
		block := data[at:min(at+blockSize, len(data))]
		if delta {
			// Delta-encoded with a distance of one; the decoder's history starts empty at each block.
			encoded := make([]byte, len(block))
			for i := range block {
				encoded[i] = block[i]
				if i > 0 {
					encoded[i] -= block[i-1]
				}
			}
			block = encoded
		}
		header := []byte{2, 0} // 12 bytes
		if delta {
			header[1] = 1 // two filters
			header = append(header, 0x03, 0x01, 0x00)
		}
		header = append(header, 0x21, 0x01, 0x00) // LZMA2, a 4 KiB dictionary
		for len(header) < 8 {
			header = append(header, 0)
		}
		header = binary.LittleEndian.AppendUint32(header, crc32.ChecksumIEEE(header))
		out.Write(header)
		size := 0
		for i, chunk := 0, block; len(chunk) > 0; i++ {
			n := min(len(chunk), 1<<16)
			control := byte(2)
			if i == 0 {
				control = 1 // the first chunk of a block resets the dictionary
			}
			out.WriteByte(control)
			out.Write(binary.BigEndian.AppendUint16(nil, uint16(n-1)))
			out.Write(chunk[:n])
			size += 3 + n
			chunk = chunk[n:]
		}
		out.WriteByte(0)
		size++
		out.Write(make([]byte, (4-size%4)%4))
		records = append(records, record{uint64(len(header) + size), uint64(len(block))})
	}
	if listed >= 0 {
		for len(records) < listed {
			records = append(records, records[0])
		}
		records = records[:listed]
	}
	vli := func(b []byte, v uint64) []byte {
		for v >= 0x80 {
			b = append(b, byte(v)|0x80)
			v >>= 7
		}
		return append(b, byte(v))
	}
	index := vli([]byte{0}, uint64(len(records)))
	for _, r := range records {
		index = vli(vli(index, r.unpadded), r.uncompressed)
	}
	for len(index)%4 != 0 {
		index = append(index, 0)
	}
	index = binary.LittleEndian.AppendUint32(index, crc32.ChecksumIEEE(index))
	out.Write(index)
	footer := binary.LittleEndian.AppendUint32(nil, uint32(len(index)/4-1))
	footer = append(footer, flags...)
	out.Write(binary.LittleEndian.AppendUint32(nil, crc32.ChecksumIEEE(footer)))
	out.Write(footer)
	out.WriteString("YZ")
	return out.Bytes()
}

// tarHolding is a tar holding betterbird/data with body.
func tarHolding(t *testing.T, body []byte) []byte {
	t.Helper()
	var out bytes.Buffer
	w := tar.NewWriter(&out)
	if err := w.WriteHeader(&tar.Header{Name: "betterbird/data", Mode: 0o644, Size: int64(len(body)),
		Typeflag: tar.TypeReg}); err != nil {
		t.Fatal(err)
	}
	w.Write(body)
	if err := w.Close(); err != nil {
		t.Fatal(err)
	}
	return out.Bytes()
}

func writeTemp(t *testing.T, body []byte) string {
	t.Helper()
	file := filepath.Join(t.TempDir(), "r.tar.xz")
	if err := os.WriteFile(file, body, 0o644); err != nil {
		t.Fatal(err)
	}
	return file
}

func TestAnXzFileOfManyDeltaBlocksIsRefusedBeforeItIsDecoded(t *testing.T) {
	// Review 68 on pull request 264: `xz --block-size=1 --delta` of a 200 kB tar kept about 67 MiB live in
	// the reader, one delta decoder per block. Its index here lists one block, as a file crafted to pass a
	// check of the index alone would; the walk counts the blocks the reader would decode.
	body := make([]byte, 200_000)
	for i := range body {
		body[i] = byte(i * 7)
	}
	crafted := xzStream(tarHolding(t, body), 1, true, 1)
	path := writeTemp(t, crafted)
	var before, after runtime.MemStats
	runtime.GC()
	runtime.ReadMemStats(&before)
	err := unpackUpstream(context.Background(), path, t.TempDir(), "betterbird")
	runtime.ReadMemStats(&after)
	if err == nil || !strings.Contains(err.Error(), "more than 1024 xz blocks") {
		t.Fatalf("a file of %d bytes in blocks of one byte was not refused for its blocks: %v", len(crafted), err)
	}
	if grew := after.TotalAlloc - before.TotalAlloc; grew > 4<<20 {
		t.Fatalf("%d bytes were allocated for a refused file", grew)
	}
}

func TestTheXzBlockLimitIsExact(t *testing.T) {
	data := make([]byte, xzBlocksMax+1)
	streams, blocks, err := xzFraming(bufio.NewReader(bytes.NewReader(xzStream(data[:xzBlocksMax], 1, true, -1))))
	if err != nil || streams != 1 || blocks != xzBlocksMax {
		t.Fatalf("%d blocks: %d streams, %d blocks, %v", xzBlocksMax, streams, blocks, err)
	}
	if _, _, err := xzFraming(bufio.NewReader(bytes.NewReader(xzStream(data, 1, true, -1)))); !errors.Is(err, errXzTooManyBlocks) {
		t.Fatalf("%d blocks: %v", xzBlocksMax+1, err)
	}
}

func TestAnXzFileOfTooManyStreamsIsRefused(t *testing.T) {
	var file []byte
	for i := 0; i <= xzStreamsMax; i++ {
		file = append(file, xzStream([]byte("x"), 1, false, -1)...)
	}
	err := unpackUpstream(context.Background(), writeTemp(t, file), t.TempDir(), "")
	if err == nil || !strings.Contains(err.Error(), "more than 16 concatenated xz streams") {
		t.Fatalf("%v", err)
	}
}

func TestAnXzIndexThatDoesNotListItsBlocksIsRefused(t *testing.T) {
	file := xzStream(tarHolding(t, []byte("x")), 512, false, 1)
	err := unpackUpstream(context.Background(), writeTemp(t, file), t.TempDir(), "betterbird")
	if err == nil || !strings.Contains(err.Error(), "its index lists 1 blocks, but it holds 4") {
		t.Fatalf("%v", err)
	}
}