🐛 fix: validate .tsf files and mmap frame data

src/frameformat.go

@@ -6,16 +6,68 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"sync"
 )
 
 // .tsf container, little-endian: "TSFR" | version uint16 | fps float64 |
 // count uint32 | count × (colorLen uint32, color JPEG).
 const (
-	tsfMagic   = "TSFR"
-	tsfVersion = 1
+	tsfMagic         = "TSFR"
+	tsfVersion       = 1
+	maxTSFFPS        = 240
+	maxTSFFrameCount = 10_000_000
 )
 
+type frameFile struct {
+	data    []byte
+	cleanup func() error
+	once    sync.Once
+	err     error
+}
+
+func (f *frameFile) Close() error {
+	if f == nil {
+		return nil
+	}
+	f.once.Do(func() {
+		if f.cleanup != nil {
+			f.err = f.cleanup()
+		}
+		f.data = nil
+	})
+	return f.err
+}
+
+var frameFileOwners sync.Map // map[*FramesContainer]*frameFile
+
+func (data *FramesContainer) Close() error {
+	if data == nil {
+		return nil
+	}
+	owner, ok := frameFileOwners.LoadAndDelete(data)
+	if !ok {
+		return nil
+	}
+	data.ColorFrames = nil
+	return owner.(*frameFile).Close()
+}
+
 func writeTSF(output string, data *FramesContainer) error {
+	if data == nil {
+		return fmt.Errorf("cannot write nil frames container")
+	}
+	if math.IsNaN(data.FPS) || math.IsInf(data.FPS, 0) || data.FPS <= 0 || data.FPS > maxTSFFPS {
+		return fmt.Errorf("cannot write .tsf: fps must be finite, positive, and at most %d", maxTSFFPS)
+	}
+	if len(data.ColorFrames) > maxTSFFrameCount || uint64(len(data.ColorFrames)) > math.MaxUint32 {
+		return fmt.Errorf("cannot write .tsf: frame count exceeds limit")
+	}
+	for i, frame := range data.ColorFrames {
+		if uint64(len(frame)) > math.MaxUint32 {
+			return fmt.Errorf("cannot write .tsf: frame %d length exceeds uint32", i)
+		}
+	}
+
 	f, err := os.Create(output)
 	if err != nil {
 		return err
@@ -48,10 +100,17 @@ func writeTSF(output string, data *FramesContainer) error {
 }
 
 func loadTSF(filename string) (*FramesContainer, error) {
-	raw, err := readFrameFile(filename)
+	file, err := readFrameFile(filename)
 	if err != nil {
 		return nil, err
 	}
+	owned := false
+	defer func() {
+		if !owned {
+			_ = file.Close()
+		}
+	}()
+	raw := file.data
 	invalid := func() error {
 		return fmt.Errorf("invalid frames file %q: corrupt .tsf container", filename)
 	}
@@ -65,27 +124,40 @@ func loadTSF(filename string) (*FramesContainer, error) {
 	}
 	fps := math.Float64frombits(binary.LittleEndian.Uint64(raw[6:]))
 	count := binary.LittleEndian.Uint32(raw[14:])
-
-	colorFrames := make([][]byte, 0, count)
-	off := 18
-	for range count {
-		if off+4 > len(raw) {
-			return nil, invalid()
-		}
-		n := int(binary.LittleEndian.Uint32(raw[off:]))
-		off += 4
-		if off+n > len(raw) {
-			return nil, invalid()
-		}
-		colorFrames = append(colorFrames, raw[off:off+n])
-		off += n
-	}
-
-	if len(colorFrames) == 0 || fps <= 0 {
+	if math.IsNaN(fps) || math.IsInf(fps, 0) || fps <= 0 || fps > maxTSFFPS {
 		return nil, fmt.Errorf(
-			"invalid frames file %q: expected non-empty frames and a positive fps",
-			filename,
+			"invalid frames file %q: fps must be finite, greater than 0, and at most %d",
+			filename, maxTSFFPS,
 		)
 	}
-	return &FramesContainer{ColorFrames: colorFrames, FPS: fps}, nil
+	if count == 0 {
+		return nil, fmt.Errorf("invalid frames file %q: expected non-empty frames", filename)
+	}
+	if count > maxTSFFrameCount || uint64(count) > uint64((len(raw)-18)/4) {
+		return nil, invalid()
+	}
+
+	colorFrames := make([][]byte, 0, int(count))
+	off := 18
+	for range count {
+		if len(raw)-off < 4 {
+			return nil, invalid()
+		}
+		n := uint64(binary.LittleEndian.Uint32(raw[off:]))
+		off += 4
+		if n > uint64(len(raw)-off) {
+			return nil, invalid()
+		}
+		nativeLen := int(n)
+		colorFrames = append(colorFrames, raw[off:off+nativeLen])
+		off += nativeLen
+	}
+
+	if off != len(raw) {
+		return nil, invalid()
+	}
+	data := &FramesContainer{ColorFrames: colorFrames, FPS: fps}
+	frameFileOwners.Store(data, file)
+	owned = true
+	return data, nil
 }

src/frameformat_hardening_test.go

@@ -0,0 +1,82 @@
+package main
+
+import (
+	"encoding/binary"
+	"math"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func tsfHeader(fps float64, count uint32) []byte {
+	raw := make([]byte, 18)
+	copy(raw, tsfMagic)
+	binary.LittleEndian.PutUint16(raw[4:], tsfVersion)
+	binary.LittleEndian.PutUint64(raw[6:], math.Float64bits(fps))
+	binary.LittleEndian.PutUint32(raw[14:], count)
+	return raw
+}
+
+func writeRawTSF(t *testing.T, raw []byte) string {
+	t.Helper()
+	path := filepath.Join(t.TempDir(), "frames.tsf")
+	if err := os.WriteFile(path, raw, 0o644); err != nil {
+		t.Fatal(err)
+	}
+	return path
+}
+
+func TestTSFRejectsInvalidFPS(t *testing.T) {
+	for _, fps := range []float64{math.NaN(), math.Inf(1), math.Inf(-1), -1, 0, 240.01} {
+		raw := append(tsfHeader(fps, 1), 0, 0, 0, 0)
+		if _, err := loadTSF(writeRawTSF(t, raw)); err == nil {
+			t.Errorf("loadTSF accepted fps %v", fps)
+		}
+	}
+}
+
+func TestTSFRejectsImpossibleCountsAndLengths(t *testing.T) {
+	if _, err := loadTSF(writeRawTSF(t, tsfHeader(30, math.MaxUint32))); err == nil {
+		t.Fatal("loadTSF accepted impossible frame count")
+	}
+
+	raw := append(tsfHeader(30, 1), 0xff, 0xff, 0xff, 0xff)
+	if _, err := loadTSF(writeRawTSF(t, raw)); err == nil {
+		t.Fatal("loadTSF accepted overflowing frame length")
+	}
+}
+
+func TestTSFCloseReleasesOwnedFrames(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "frames.tsf")
+	if err := writeTSF(path, &FramesContainer{FPS: 30, ColorFrames: [][]byte{{1, 2, 3}}}); err != nil {
+		t.Fatal(err)
+	}
+	frames, err := loadTSF(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := frames.ColorFrames[0]; len(got) != 3 || got[0] != 1 {
+		t.Fatalf("unexpected zero-copy frame data: %v", got)
+	}
+	if err := frames.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	if frames.ColorFrames != nil {
+		t.Fatal("Close retained references to released frame data")
+	}
+	if err := frames.Close(); err != nil {
+		t.Fatalf("second Close: %v", err)
+	}
+}
+
+func TestTSFWriteRejectsInvalidHeaderValuesBeforeCreate(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "frames.tsf")
+	err := writeTSF(path, &FramesContainer{FPS: math.NaN(), ColorFrames: [][]byte{{1}}})
+	if err == nil || !strings.Contains(err.Error(), "fps") {
+		t.Fatalf("writeTSF error = %v", err)
+	}
+	if _, err := os.Stat(path); !os.IsNotExist(err) {
+		t.Fatalf("invalid write created output: %v", err)
+	}
+}

src/mmap_other.go

@@ -4,6 +4,7 @@ package main
 
 import "os"
 
-func readFrameFile(filename string) ([]byte, error) {
-	return os.ReadFile(filename)
+func readFrameFile(filename string) (*frameFile, error) {
+	data, err := os.ReadFile(filename)
+	return &frameFile{data: data}, err
 }

src/mmap_unix.go

@@ -7,7 +7,7 @@ import (
 	"syscall"
 )
 
-func readFrameFile(filename string) ([]byte, error) {
+func readFrameFile(filename string) (*frameFile, error) {
 	f, err := os.Open(filename)
 	if err != nil {
 		return nil, err
@@ -20,11 +20,18 @@ func readFrameFile(filename string) ([]byte, error) {
 	}
 	size := info.Size()
 	if size <= 0 || size != int64(int(size)) {
-		return os.ReadFile(filename)
+		data, err := os.ReadFile(filename)
+		return &frameFile{data: data}, err
 	}
 	data, err := syscall.Mmap(int(f.Fd()), 0, int(size), syscall.PROT_READ, syscall.MAP_SHARED)
 	if err != nil {
-		return os.ReadFile(filename)
+		data, err := os.ReadFile(filename)
+		return &frameFile{data: data}, err
 	}
-	return data, nil
+	return &frameFile{
+		data: data,
+		cleanup: func() error {
+			return syscall.Munmap(data)
+		},
+	}, nil
 }