diff --git a/internal/tgb/BUILD.bazel b/internal/tgb/BUILD.bazel index ff05fdf0..f182e9ce 100644 --- a/internal/tgb/BUILD.bazel +++ b/internal/tgb/BUILD.bazel @@ -23,6 +23,7 @@ go_library( go_test( name = "tgb_test", srcs = [ + "expansion_test.go", "fuzz_test.go", "hashwidth_test.go", "sentinel_test.go", diff --git a/internal/tgb/expansion_test.go b/internal/tgb/expansion_test.go new file mode 100644 index 00000000..21c6469e --- /dev/null +++ b/internal/tgb/expansion_test.go @@ -0,0 +1,141 @@ +package tgb_test + +import ( + "encoding/binary" + "encoding/hex" + "fmt" + "hash/crc32" + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/tango/entity" + "github.com/uber/tango/internal/tgb" +) + +// taglessGraph builds an n-node graph in which no target has tags, like a repo +// whose rules never set them. Everything else varies so that TAG_DEG is the +// only highly compressible column. +func taglessGraph(n int) *tgb.Graph { + rng := rand.New(rand.NewSource(1)) + idMap := make(map[int32]string, n) + targets := make([]entity.OptimizedTarget, n) + for i := range targets { + id := int32(i + 1) + idMap[id] = fmt.Sprintf("//pkg/p%d:t%d", rng.Intn(n/3+1), i) + raw := make([]byte, 8) + rng.Read(raw) + attrs := map[int32]int32{} + if rng.Intn(4) == 0 { + attrs[0] = int32(rng.Intn(3)) + } + var deps []int32 + for j := rng.Intn(4); i > 0 && j > 0; j-- { + deps = append(deps, int32(rng.Intn(i)+1)) + } + targets[i] = entity.OptimizedTarget{ + ID: id, + Hash: hex.EncodeToString(raw), + DirectDependencies: deps, + RuleType: int32(rng.Intn(4)), + Attributes: attrs, + } + } + return &tgb.Graph{ + Targets: targets, + Metadata: entity.Metadata{ + TargetIDMapping: idMap, + RuleTypeMapping: map[int32]string{0: "source file", 1: "js_library", 2: "ts_project", 3: "js_test"}, + AttributeNameMapping: map[int32]string{0: "visibility"}, + AttributeStringValueMapping: map[int32]string{0: "//visibility:public", 1: "true", 2: "false"}, + }, + } +} + +// TestTaglessGraphReadable guards against rejecting graphs whose per-node +// columns are all zeros. With no tags TAG_DEG is nodeCount zero bytes, which +// zstd compresses to a near-constant ~100 bytes, so its expansion ratio grows +// with nodeCount without bound even though the blob is valid. +func TestTaglessGraphReadable(t *testing.T) { + const nodes = 700_000 + data := mustEncode(t, taglessGraph(nodes), tgb.EncodeOptions{HashBytes: 8, BlockSize: 512}) + + r, err := tgb.NewReader(data) + require.NoError(t, err, "valid tagless graph must be readable") + assert.Equal(t, nodes, r.NodeCount()) + + // Keep the test honest: it only covers the bug if TAG_DEG expands past + // the 4096x per-column ratio cap this reader used to enforce. + stats, err := tgb.ColumnStats(data) + require.NoError(t, err) + var found bool + for _, s := range stats { + if s.Name != "TAG_DEG" { + continue + } + found = true + require.NotZero(t, s.CompressedSize) + require.Greater(t, s.RawSize/s.CompressedSize, uint64(4096), + "TAG_DEG no longer exceeds 4096x; grow nodes") + } + require.True(t, found, "TAG_DEG column missing") + + g, err := tgb.Decode(data) + require.NoError(t, err) + assert.Len(t, g.Targets, nodes) +} + +// withColumnRawSize returns data with the directory's rawSize for column id +// replaced and the directory CRC recomputed, as a hostile blob would. +func withColumnRawSize(t *testing.T, data []byte, id, rawSize uint64) []byte { + t.Helper() + const footerSize = 16 + dirOff := binary.LittleEndian.Uint64(data[len(data)-footerSize:]) + dir := data[dirOff : len(data)-footerSize] + + next := func() uint64 { + v, n := binary.Uvarint(dir) + require.Positive(t, n, "truncated directory") + dir = dir[n:] + return v + } + count := next() + newDir := binary.AppendUvarint(nil, count) + var patched bool + for i := uint64(0); i < count; i++ { + colID := next() + codec := dir[0] + dir = dir[1:] + offset, comp, raw := next(), next(), next() + if colID == id { + raw, patched = rawSize, true + } + newDir = binary.AppendUvarint(newDir, colID) + newDir = append(newDir, codec) + newDir = binary.AppendUvarint(newDir, offset) + newDir = binary.AppendUvarint(newDir, comp) + newDir = binary.AppendUvarint(newDir, raw) + } + require.True(t, patched, "column %d not in directory", id) + + out := append([]byte(nil), data[:dirOff]...) + out = append(out, newDir...) + out = binary.LittleEndian.AppendUint64(out, dirOff) + out = binary.LittleEndian.AppendUint32(out, crc32.Checksum(newDir, crc32.MakeTable(crc32.Castagnoli))) + return append(out, "TGB1"...) +} + +// TestHostileColumnExpansionStillRejected checks that, without a per-column +// ratio cap, a small blob still cannot claim a huge column. +func TestHostileColumnExpansionStillRejected(t *testing.T) { + const tagDeg = 10 + data := mustEncode(t, buildTinyGraph(), tgb.EncodeOptions{HashBytes: 8, BlockSize: 4}) + + _, err := tgb.NewReader(data) + require.NoError(t, err, "unmodified blob must be readable") + + _, err = tgb.NewReader(withColumnRawSize(t, data, tagDeg, 1<<28)) + require.Error(t, err) + assert.Contains(t, err.Error(), "total raw bytes") +} diff --git a/internal/tgb/format.go b/internal/tgb/format.go index 27a72722..6e39baab 100644 --- a/internal/tgb/format.go +++ b/internal/tgb/format.go @@ -41,20 +41,13 @@ const ( // raw; anything near this cap is not a real graph. maxColumnRawSize = uint64(1) << 30 - // maxZstdExpansion bounds rawSize/compressedSize for zstd columns. The - // most compressible real column (TAG_DEG, almost all zeros) expands - // ~75x; 4096x is generous headroom while still stopping a kilobyte - // blob from claiming gigabytes. - maxZstdExpansion = uint64(4096) - // maxTotalRawFactor bounds the *sum* of claimed rawSize across all // columns relative to the blob size, with maxTotalRawFloor as an - // absolute floor so tiny legitimate blobs are never rejected. The - // per-column expansion cap alone is not enough: a small blob full of - // RLE frames can honestly expand 30,000x, and everything downstream - // (dict expansion, degree indexes, Decode) allocates in proportion to - // raw bytes. The real graph measures 1.81x raw/file; 16x is ~9x - // headroom. + // absolute floor so tiny legitimate blobs are never rejected. A small + // blob full of RLE frames can honestly expand 30,000x, and everything + // downstream (dict expansion, degree indexes, Decode) allocates in + // proportion to raw bytes. The real graph measures 1.81x raw/file; 16x + // is ~9x headroom. maxTotalRawFactor = uint64(16) maxTotalRawFloor = uint64(4) << 20 diff --git a/internal/tgb/reader.go b/internal/tgb/reader.go index c684cfed..2759571f 100644 --- a/internal/tgb/reader.go +++ b/internal/tgb/reader.go @@ -383,10 +383,6 @@ func validateColumns(hdr fileHeader, cols map[uint64]colEntry, fileSize, dirOffs if e.rawSize > 0 && e.compressedSize == 0 { return fmt.Errorf("tgb: column %d claims %d raw bytes from empty data", id, e.rawSize) } - if e.compressedSize > 0 && e.rawSize > e.compressedSize*maxZstdExpansion { - return fmt.Errorf("tgb: column %d claims %dx expansion (max %d)", - id, e.rawSize/e.compressedSize, maxZstdExpansion) - } default: return fmt.Errorf("tgb: column %d has unknown codec %d", id, e.codec) } @@ -396,11 +392,11 @@ func validateColumns(hdr fileHeader, cols map[uint64]colEntry, fileSize, dirOffs } totalRaw += e.rawSize } - // The per-column expansion cap cannot stop a blob of highly-compressible - // frames from honestly claiming raw bytes far beyond anything a real - // graph produces; every downstream allocation is proportional to raw - // bytes, so bound their sum against the file itself (real graphs - // measure ~1.8x). + // Every downstream allocation is proportional to raw bytes, so bound + // their sum against the file itself (real graphs measure ~1.8x). A + // per-column compression ratio is no bound: a valid per-node column of + // constant values (TAG_DEG in a graph with no tags) compresses to a + // near-fixed size, so its ratio grows without limit as nodeCount does. if budget := max(fileSize*maxTotalRawFactor, maxTotalRawFloor); totalRaw > budget { return fmt.Errorf("tgb: columns claim %d total raw bytes from a %d-byte file (max %d)", totalRaw, fileSize, budget)