tensor/io/hostile_test.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package io

import (
	"encoding/binary"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// Hostile inputs: a data file is untrusted input, so every size field
// the header carries must be checked against the bytes actually
// present before it is used. These cases pin the contract that a
// malformed file is an error, never a panic and never an allocation
// sized by the header alone.

// hostileNCName writes a 4-byte-padded NetCDF name.
func hostileNCName(b []byte, name string) []byte {
	b = binary.BigEndian.AppendUint32(b, uint32(len(name)))
	b = append(b, name...)
	if pad := (4 - len(name)%4) % 4; pad != 0 {
		b = append(b, make([]byte, pad)...)
	}
	return b
}

// ncHostileHeader builds a minimal CDF-1 file: magic, numrecs, a
// dimension list with the given lengths, an absent attribute list, and
// one NC_DOUBLE variable over those dimensions. Nothing follows the
// header, so any declared data is truncated by construction.
func ncHostileHeader(names []string, lengths []uint32) []byte {
	b := []byte{'C', 'D', 'F', 1}
	b = binary.BigEndian.AppendUint32(b, 0) // numrecs
	b = binary.BigEndian.AppendUint32(b, ncTagDimension)
	b = binary.BigEndian.AppendUint32(b, uint32(len(lengths)))
	for i, l := range lengths {
		b = hostileNCName(b, names[i])
		b = binary.BigEndian.AppendUint32(b, l)
	}
	b = binary.BigEndian.AppendUint32(b, 0) // absent attribute list
	b = binary.BigEndian.AppendUint32(b, 0)
	b = binary.BigEndian.AppendUint32(b, ncTagVariable)
	b = binary.BigEndian.AppendUint32(b, 1)
	b = hostileNCName(b, "v")
	b = binary.BigEndian.AppendUint32(b, uint32(len(lengths)))
	for i := range lengths {
		b = binary.BigEndian.AppendUint32(b, uint32(i))
	}
	b = binary.BigEndian.AppendUint32(b, 0) // absent variable attributes
	b = binary.BigEndian.AppendUint32(b, 0)
	b = binary.BigEndian.AppendUint32(b, ncTypeDouble)
	b = binary.BigEndian.AppendUint32(b, 0) // vsize
	b = binary.BigEndian.AppendUint32(b, 0) // begin
	return b
}

// writeHostile drops the bytes into a temp file and returns the path.
func writeHostile(t *testing.T, name string, data []byte) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), name)
	if err := os.WriteFile(path, data, 0o644); err != nil {
		t.Fatal(err)
	}
	return path
}

// TestLoadNetCDFHostileCounts covers the three ways a declared count
// can be used to demand memory the file cannot back: a product that
// wraps to a small positive number, a product that wraps negative, and
// record counts that are simply larger than the file.
func TestLoadNetCDFHostileCounts(t *testing.T) {
	cases := []struct {
		name    string
		names   []string
		lengths []uint32
	}{
		{
			// 2^21 * 2^20 * 2^20 = 2^61, which times the 8-byte
			// element width wraps to zero in a 64-bit multiply.
			name:    "product wraps to zero",
			names:   []string{"a", "b", "c"},
			lengths: []uint32{1 << 21, 1 << 20, 1 << 20},
		},
		{
			// 4294967295 squared wraps to a negative product.
			name:    "product wraps negative",
			names:   []string{"a", "b"},
			lengths: []uint32{4294967295, 4294967295},
		},
		{
			name:    "one dimension longer than the file",
			names:   []string{"a"},
			lengths: []uint32{1 << 30},
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			path := writeHostile(t, "hostile.nc", ncHostileHeader(tc.names, tc.lengths))
			if _, _, _, err := LoadNetCDF(path); err == nil {
				t.Fatal("expected an error for a header whose data cannot fit the file")
			}
		})
	}
}

// TestLoadNetCDFHostileRecordCounts pins the list caps: a header that
// declares millions of dimensions, attributes or variables in a file
// of a few bytes must fail before the list is allocated.
func TestLoadNetCDFHostileRecordCounts(t *testing.T) {
	const declared = 2000000

	dimCount := []byte{'C', 'D', 'F', 1}
	dimCount = binary.BigEndian.AppendUint32(dimCount, 0)
	dimCount = binary.BigEndian.AppendUint32(dimCount, ncTagDimension)
	dimCount = binary.BigEndian.AppendUint32(dimCount, declared)

	attrCount := []byte{'C', 'D', 'F', 1}
	attrCount = binary.BigEndian.AppendUint32(attrCount, 0)
	attrCount = binary.BigEndian.AppendUint32(attrCount, 0) // absent dimensions
	attrCount = binary.BigEndian.AppendUint32(attrCount, 0)
	attrCount = binary.BigEndian.AppendUint32(attrCount, ncTagAttribute)
	attrCount = binary.BigEndian.AppendUint32(attrCount, declared)

	varCount := []byte{'C', 'D', 'F', 1}
	varCount = binary.BigEndian.AppendUint32(varCount, 0)
	varCount = binary.BigEndian.AppendUint32(varCount, 0) // absent dimensions
	varCount = binary.BigEndian.AppendUint32(varCount, 0)
	varCount = binary.BigEndian.AppendUint32(varCount, 0) // absent attributes
	varCount = binary.BigEndian.AppendUint32(varCount, 0)
	varCount = binary.BigEndian.AppendUint32(varCount, ncTagVariable)
	varCount = binary.BigEndian.AppendUint32(varCount, declared)

	cases := []struct {
		name string
		data []byte
	}{
		{"dimensions", dimCount},
		{"attributes", attrCount},
		{"variables", varCount},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			path := writeHostile(t, "counts.nc", tc.data)
			_, _, _, err := LoadNetCDF(path)
			if err == nil {
				t.Fatalf("a %d-byte file declaring %d %s must be refused", len(tc.data), declared, tc.name)
			}
			if !strings.Contains(err.Error(), "remaining bytes") {
				t.Fatalf("error = %v, want the declared-count bound", err)
			}
		})
	}
}

// TestLoadNetCDFTruncatedVariable pins the data-block check: a header
// that points its variable past the end of the file fails without
// reading beyond it.
func TestLoadNetCDFTruncatedVariable(t *testing.T) {
	data := ncHostileHeader([]string{"a"}, []uint32{4})
	binary.BigEndian.PutUint32(data[len(data)-4:], 1<<30) // begin past EOF
	path := writeHostile(t, "short.nc", data)
	if _, _, _, err := LoadNetCDF(path); err == nil {
		t.Fatal("expected an error for a variable whose data lies past the end of the file")
	}
}

// TestLoadFITSTableTruncatedHostile pins the table data guard: a file
// that ends inside the header block has no data at all, whether or not
// the column forms give it a row size to divide by.
func TestLoadFITSTableTruncatedHostile(t *testing.T) {
	// No data block follows the header: the cards stop at 640 bytes.
	header := func(form string) []byte {
		var b []byte
		for _, body := range []string{
			"XTENSION= 'BINTABLE'", "BITPIX  =                    8", "NAXIS   =                    2",
			"NAXIS1  =                    4", "NAXIS2  =                    2", "TFIELDS =                    1",
			"TFORM1  = '" + form + strings.Repeat(" ", 8-len(form)) + "'", "END",
		} {
			b = append(b, card(body)...)
		}
		return b
	}
	for _, form := range []string{"X", "D"} {
		t.Run(form, func(t *testing.T) {
			data := header(form)
			if len(data) != 640 {
				t.Fatalf("the test header is %d bytes, want 640", len(data))
			}
			path := writeHostile(t, "trunc.fits", data)
			if _, err := LoadFITSTable(path); err == nil {
				t.Fatal("expected an error for a table whose data block is missing")
			}
		})
	}
}

// TestLoadNetCDFUnusedLongDimension pins the other side of the bound: a
// header may declare a dimension longer than the data any variable
// uses, because the format allows it, so the reader must accept the
// file rather than refuse a legal one.
func TestLoadNetCDFUnusedLongDimension(t *testing.T) {
	// One dimension of a million elements, one variable using none of
	// it: the count product is 1, so nothing is read.
	var b []byte
	b = append(b, 'C', 'D', 'F', 1)
	b = binary.BigEndian.AppendUint32(b, 0)  // numrecs
	b = binary.BigEndian.AppendUint32(b, 10) // NC_DIMENSION
	b = binary.BigEndian.AppendUint32(b, 1)
	b = hostileNCName(b, "big")
	b = binary.BigEndian.AppendUint32(b, 1<<20)
	b = binary.BigEndian.AppendUint32(b, 0) // absent attributes
	b = binary.BigEndian.AppendUint32(b, 0)
	b = binary.BigEndian.AppendUint32(b, 11) // NC_VARIABLE
	b = binary.BigEndian.AppendUint32(b, 1)
	b = hostileNCName(b, "scalar")
	b = binary.BigEndian.AppendUint32(b, 0) // rank 0
	b = binary.BigEndian.AppendUint32(b, 0) // absent attributes
	b = binary.BigEndian.AppendUint32(b, 0)
	b = binary.BigEndian.AppendUint32(b, 6) // NC_DOUBLE
	b = binary.BigEndian.AppendUint32(b, 8) // vsize
	b = binary.BigEndian.AppendUint32(b, 0) // begin
	b = append(b, 0, 0, 0, 0, 0, 0, 0, 0)   // one double at offset 0
	path := writeHostile(t, "unused.nc", b)
	dims, vars, _, err := LoadNetCDF(path)
	if err != nil {
		t.Fatalf("LoadNetCDF refused a legal file: %v", err)
	}
	if len(dims) != 1 || dims[0].Length != 1<<20 {
		t.Fatalf("dims = %+v, want one dimension of 2^20", dims)
	}
	if len(vars) != 1 || vars[0].Values.Len() != 1 {
		t.Fatalf("vars = %+v, want one scalar", vars)
	}
}

// TestLoadHDF5ChunkTreeCycle pins the visited set on the chunk B-tree
// walk: an inner node that lists itself among its children must be
// refused. Without the set the walk multiplies into entries^depth node
// visits before the depth guard can fire, so a few hundred crafted
// bytes never terminate.
func TestLoadHDF5ChunkTreeCycle(t *testing.T) {
	const entries = 512
	rank := 1
	keySize := 8 + 8*(rank+1)
	entrySize := keySize + 8
	buf := make([]byte, 24+entries*entrySize)
	copy(buf, []byte("TREE"))
	buf[4] = 1 // a chunk node
	buf[5] = 1 // inner level: every child is recursed, not placed
	binary.LittleEndian.PutUint16(buf[6:], entries)
	for i := range entries {
		p := 24 + i*entrySize
		binary.LittleEndian.PutUint32(buf[p:], 16)  // chunk size
		binary.LittleEndian.PutUint32(buf[p+4:], 0) // filter mask
		// The key offsets stay zero; the child address points at
		// this very node.
		binary.LittleEndian.PutUint64(buf[p+keySize:], 0)
	}
	f := &hdf5File{data: buf, offSize: 8, lenSize: 8}
	place := func(uint64, []uint64, uint32, int) error { return nil }
	if err := f.chunkTree(0, hdf5Layout{}, rank, map[uint64]bool{}, place, 0); err == nil || !strings.Contains(err.Error(), "revisits") {
		t.Fatalf("chunkTree on a self-referencing node: %v", err)
	}
}