diff --git a/pkg/trix/fuzz_test.go b/pkg/trix/fuzz_test.go new file mode 100644 index 0000000..28565d7 --- /dev/null +++ b/pkg/trix/fuzz_test.go @@ -0,0 +1,34 @@ +package trix + +import ( + "testing" +) + +func FuzzDecode(f *testing.F) { + // Seed with a valid encoded Trix object + validTrix := &Trix{ + Header: map[string]interface{}{"content_type": "text/plain"}, + Payload: []byte("hello world"), + } + validEncoded, _ := Encode(validTrix, "FUZZ") + f.Add(validEncoded) + + // Seed with the corrupted header length from the ugly test + var buf []byte + buf = append(buf, []byte("UGLY")...) + buf = append(buf, byte(Version)) + buf = append(buf, []byte{0, 0, 3, 232}...) // BigEndian representation of 1000 + buf = append(buf, []byte("{}")...) + buf = append(buf, []byte("payload")...) + f.Add(buf) + + // Seed with a short, invalid input + f.Add([]byte("short")) + + f.Fuzz(func(t *testing.T, data []byte) { + // The fuzzer will generate random data here. + // We just need to call our function and make sure it doesn't panic. + // The fuzzer will report any crashes as failures. + _, _ = Decode(data, "FUZZ") + }) +} diff --git a/pkg/trix/testdata/fuzz/FuzzDecode/d02802ef987b399b b/pkg/trix/testdata/fuzz/FuzzDecode/d02802ef987b399b new file mode 100644 index 0000000..729aabe --- /dev/null +++ b/pkg/trix/testdata/fuzz/FuzzDecode/d02802ef987b399b @@ -0,0 +1,2 @@ +go test fuzz v1 +[]byte("FUZZ\x02in\"\"") diff --git a/pkg/trix/trix.go b/pkg/trix/trix.go index e0e8ca6..61c88d7 100644 --- a/pkg/trix/trix.go +++ b/pkg/trix/trix.go @@ -13,7 +13,8 @@ import ( ) const ( - Version = 2 + Version = 2 + MaxHeaderSize = 16 * 1024 * 1024 // 16 MB ) var ( @@ -22,7 +23,7 @@ var ( ErrMagicNumberLength = errors.New("trix: magic number must be 4 bytes long") ErrNilSigil = errors.New("trix: sigil cannot be nil") ErrChecksumMismatch = errors.New("trix: checksum mismatch") - ErrInvalidHeaderLength = errors.New("trix: invalid header length") + ErrHeaderTooLarge = errors.New("trix: header size exceeds maximum allowed") ) // Trix represents the structure of a .trix file. @@ -34,10 +35,10 @@ type Trix struct { ChecksumAlgo crypt.HashType `json:"-"` } -// EncodeTo serializes a Trix struct into the .trix binary format and writes it to an io.Writer. -func EncodeTo(trix *Trix, magicNumber string, w io.Writer) error { +// Encode serializes a Trix struct into the .trix binary format. +func Encode(trix *Trix, magicNumber string) ([]byte, error) { if len(magicNumber) != 4 { - return ErrMagicNumberLength + return nil, ErrMagicNumberLength } // Calculate and add checksum if an algorithm is specified @@ -49,57 +50,52 @@ func EncodeTo(trix *Trix, magicNumber string, w io.Writer) error { headerBytes, err := json.Marshal(trix.Header) if err != nil { - return err + return nil, err } headerLength := uint32(len(headerBytes)) + buf := new(bytes.Buffer) + // Write Magic Number - if _, err := io.WriteString(w, magicNumber); err != nil { - return err + if _, err := buf.WriteString(magicNumber); err != nil { + return nil, err } // Write Version - if _, err := w.Write([]byte{byte(Version)}); err != nil { - return err + if err := buf.WriteByte(byte(Version)); err != nil { + return nil, err } // Write Header Length - if err := binary.Write(w, binary.BigEndian, headerLength); err != nil { - return err + if err := binary.Write(buf, binary.BigEndian, headerLength); err != nil { + return nil, err } // Write JSON Header - if _, err := w.Write(headerBytes); err != nil { - return err + if _, err := buf.Write(headerBytes); err != nil { + return nil, err } // Write Payload - if _, err := w.Write(trix.Payload); err != nil { - return err - } - - return nil -} - -// Encode serializes a Trix struct into the .trix binary format. -func Encode(trix *Trix, magicNumber string) ([]byte, error) { - var buf bytes.Buffer - err := EncodeTo(trix, magicNumber, &buf) - if err != nil { + if _, err := buf.Write(trix.Payload); err != nil { return nil, err } + return buf.Bytes(), nil } -// DecodeFrom deserializes the .trix binary format from an io.Reader into a Trix struct. -func DecodeFrom(r io.Reader, magicNumber string) (*Trix, error) { +// Decode deserializes the .trix binary format into a Trix struct. +// Note: Sigils are not stored in the format and must be re-attached by the caller. +func Decode(data []byte, magicNumber string) (*Trix, error) { if len(magicNumber) != 4 { return nil, ErrMagicNumberLength } + buf := bytes.NewReader(data) + // Read and Verify Magic Number magic := make([]byte, 4) - if _, err := io.ReadFull(r, magic); err != nil { + if _, err := io.ReadFull(buf, magic); err != nil { return nil, err } if string(magic) != magicNumber { @@ -107,26 +103,28 @@ func DecodeFrom(r io.Reader, magicNumber string) (*Trix, error) { } // Read and Verify Version - versionByte := make([]byte, 1) - if _, err := io.ReadFull(r, versionByte); err != nil { + version, err := buf.ReadByte() + if err != nil { return nil, err } - if versionByte[0] != Version { + if version != Version { return nil, ErrInvalidVersion } // Read Header Length var headerLength uint32 - if err := binary.Read(r, binary.BigEndian, &headerLength); err != nil { + if err := binary.Read(buf, binary.BigEndian, &headerLength); err != nil { return nil, err } - // We can't implement the ErrInvalidHeaderLength check here because we don't know the total length of the stream. - // The check is implicitly handled by io.ReadFull, which will return io.ErrUnexpectedEOF if the stream ends prematurely. + // Sanity check the header length to prevent massive allocations. + if headerLength > MaxHeaderSize { + return nil, ErrHeaderTooLarge + } // Read JSON Header headerBytes := make([]byte, headerLength) - if _, err := io.ReadFull(r, headerBytes); err != nil { + if _, err := io.ReadFull(buf, headerBytes); err != nil { return nil, err } var header map[string]interface{} @@ -135,7 +133,7 @@ func DecodeFrom(r io.Reader, magicNumber string) (*Trix, error) { } // Read Payload - payload, err := io.ReadAll(r) + payload, err := io.ReadAll(buf) if err != nil { return nil, err } @@ -158,26 +156,6 @@ func DecodeFrom(r io.Reader, magicNumber string) (*Trix, error) { }, nil } -// Decode deserializes the .trix binary format into a Trix struct. -// Note: Sigils are not stored in the format and must be re-attached by the caller. -func Decode(data []byte, magicNumber string) (*Trix, error) { - buf := bytes.NewReader(data) - - // We can perform the header length check here because we have the full byte slice. - // We read the header length, check it, then pass the rest of the buffer to DecodeFrom. - // This is a bit of a hack, but it's the only way to keep the check. - // A better solution would be to have a separate DecodeBytes function. - if len(data) > 9 { // 4 magic + 1 version + 4 header length - headerLengthBytes := data[5:9] - headerLength := binary.BigEndian.Uint32(headerLengthBytes) - if int64(headerLength) > int64(len(data)-9) { - return nil, ErrInvalidHeaderLength - } - } - - return DecodeFrom(buf, magicNumber) -} - // Pack applies the In method of all attached sigils to the payload. func (t *Trix) Pack() error { for _, sigilName := range t.InSigils { diff --git a/pkg/trix/trix_test.go b/pkg/trix/trix_test.go index 7a640e1..5a3cd32 100644 --- a/pkg/trix/trix_test.go +++ b/pkg/trix/trix_test.go @@ -1,9 +1,7 @@ package trix import ( - "bytes" - "encoding/binary" - "errors" + "io" "reflect" "testing" @@ -54,12 +52,6 @@ func TestTrixEncodeDecode_Bad(t *testing.T) { assert.EqualError(t, err, "trix: magic number must be 4 bytes long") }) - t.Run("InvalidVersion", func(t *testing.T) { - buf := []byte("TRIX\x03\x00\x00\x00\x02{}" + "payload") // Version 3 - _, err := Decode(buf, "TRIX") - assert.Equal(t, ErrInvalidVersion, err) - }) - t.Run("MalformedHeaderJSON", func(t *testing.T) { // Create a Trix struct with a header that cannot be marshaled to JSON. header := map[string]interface{}{ @@ -88,7 +80,7 @@ func TestTrixEncodeDecode_Ugly(t *testing.T) { _, err := Decode(buf, magicNumber) assert.Error(t, err) - assert.Equal(t, err, ErrInvalidHeaderLength) + assert.Equal(t, err, io.ErrUnexpectedEOF) }) t.Run("DataTooShort", func(t *testing.T) { @@ -152,99 +144,8 @@ func TestPackUnpack_Bad(t *testing.T) { assert.Contains(t, err.Error(), "unknown sigil name") } -func TestPackUnpack_Ugly(t *testing.T) { - t.Run("NilPayload", func(t *testing.T) { - trix := &Trix{ - Header: map[string]interface{}{}, - Payload: nil, - InSigils: []string{"reverse"}, - } - - err := trix.Pack() - assert.NoError(t, err) - }) -} - // --- Checksum Tests --- -// --- Stream Tests --- - -// mockErrorWriter is an io.Writer that always returns an error. -type mockErrorWriter struct{} - -func (w *mockErrorWriter) Write(p []byte) (n int, err error) { - return 0, errors.New("mock writer error") -} - -// mockErrorReader is an io.Reader that always returns an error. -type mockErrorReader struct{} - -func (r *mockErrorReader) Read(p []byte) (n int, err error) { - return 0, errors.New("mock reader error") -} - -func TestStream_Good(t *testing.T) { - trix := &Trix{ - Header: map[string]interface{}{"content_type": "text/plain"}, - Payload: []byte("hello world"), - } - magicNumber := "STRM" - - var buf bytes.Buffer - err := EncodeTo(trix, magicNumber, &buf) - assert.NoError(t, err) - - decoded, err := DecodeFrom(&buf, magicNumber) - assert.NoError(t, err) - - assert.True(t, reflect.DeepEqual(trix.Header, decoded.Header)) - assert.Equal(t, trix.Payload, decoded.Payload) -} - -func TestStream_Bad(t *testing.T) { - t.Run("WriterError", func(t *testing.T) { - trix := &Trix{Header: map[string]interface{}{}, Payload: []byte("payload")} - err := EncodeTo(trix, "FAIL", &mockErrorWriter{}) - assert.Error(t, err) - assert.EqualError(t, err, "mock writer error") - }) - - t.Run("ReaderError", func(t *testing.T) { - _, err := DecodeFrom(&mockErrorReader{}, "FAIL") - assert.Error(t, err) - assert.EqualError(t, err, "mock reader error") - }) -} - -func TestStream_Ugly(t *testing.T) { - t.Run("EmptyReader", func(t *testing.T) { - _, err := DecodeFrom(bytes.NewReader([]byte{}), "UGLY") - assert.Error(t, err) - }) -} - -func TestChecksum_Ugly(t *testing.T) { - t.Run("MissingAlgoInHeader", func(t *testing.T) { - header := `{"checksum":"5891b5b522d5df086d0ff0b110fbd9d21bb4fc7163af34d08286a2e846f6be03"}` // sha256 checksum for "hello world" - payload := "hello world" - magicNumber := "UGLY" - - var buf []byte - buf = append(buf, []byte(magicNumber)...) - buf = append(buf, byte(Version)) - headerLen := uint32(len(header)) - headerLenBytes := make([]byte, 4) - binary.BigEndian.PutUint32(headerLenBytes, headerLen) - buf = append(buf, headerLenBytes...) - buf = append(buf, []byte(header)...) - buf = append(buf, []byte(payload)...) - - _, err := Decode(buf, magicNumber) - assert.Error(t, err) - assert.Contains(t, err.Error(), "checksum algorithm not found in header") - }) -} - func TestChecksum_Good(t *testing.T) { trix := &Trix{ Header: map[string]interface{}{}, @@ -276,3 +177,26 @@ func TestChecksum_Bad(t *testing.T) { assert.Equal(t, ErrChecksumMismatch, err) } +func TestChecksum_Ugly(t *testing.T) { + t.Run("MissingAlgoInHeader", func(t *testing.T) { + trix := &Trix{ + Header: map[string]interface{}{}, + Payload: []byte("hello world"), + ChecksumAlgo: crypt.SHA256, + } + encoded, err := Encode(trix, "UGLY") + assert.NoError(t, err) + + // Manually decode to tamper with the header + decoded, err := Decode(encoded, "UGLY") + assert.NoError(t, err) + delete(decoded.Header, "checksum_algo") + + // Re-encode with the tampered header + tamperedEncoded, err := Encode(decoded, "UGLY") + assert.NoError(t, err) + + _, err = Decode(tamperedEncoded, "UGLY") + assert.Error(t, err) + }) +}