Re: [PR] fix(puffin): read LZ4-compressed footer payloads [iceberg-go]

["laskoviymishka (via GitHub)" <gi...@apache.org>]

laskoviymishka commented on code in PR #1678:
URL: https://github.com/apache/iceberg-go/pull/1678#discussion_r3740296270


##########
puffin/puffin_reader.go:
##########
@@ -224,11 +238,32 @@ func (r *Reader) readFooter() error {
 	}
 
 	payloadReader := io.NewSectionReader(r.r, footerStart+MagicSize, payloadSize)
-	decoder := json.NewDecoder(payloadReader)
+	var footerReader io.Reader = payloadReader
+	if flags&FooterFlagCompressed != 0 {
+		lz4Reader := lz4.NewReader(payloadReader)
+		var firstByte [1]byte
+		n, err := lz4Reader.Read(firstByte[:])
+		if err != nil && !errors.Is(err, io.EOF) {
+			return fmt.Errorf("puffin: initialize compressed footer: %w", err)
+		}
+		if lz4Reader.Size() == 0 {
+			return errors.New("puffin: compressed footer LZ4 frame is missing content size")
+		}
+		if int64(lz4Reader.Size()) > r.maxFooterSize {

Review Comment:
   `lz4Reader.Size()` is a `uint64`; casting to `int64` first means a frame advertising a content size above `math.MaxInt64` wraps negative and sails straight past this guard. The `LimitedReader` still catches it downstream, but this early check is exactly the fail-fast-before-allocation gate, and a one-bit flip in the frame header disables it.
   
   Comparing in uint64 space closes it, since `r.maxFooterSize` is already guaranteed positive above:
   
   ```go
   if lz4Reader.Size() > uint64(r.maxFooterSize) {
   ```



##########
puffin/puffin_reader.go:
##########
@@ -224,11 +238,32 @@ func (r *Reader) readFooter() error {
 	}
 
 	payloadReader := io.NewSectionReader(r.r, footerStart+MagicSize, payloadSize)
-	decoder := json.NewDecoder(payloadReader)
+	var footerReader io.Reader = payloadReader
+	if flags&FooterFlagCompressed != 0 {
+		lz4Reader := lz4.NewReader(payloadReader)
+		var firstByte [1]byte
+		n, err := lz4Reader.Read(firstByte[:])
+		if err != nil && !errors.Is(err, io.EOF) {
+			return fmt.Errorf("puffin: initialize compressed footer: %w", err)
+		}
+		if lz4Reader.Size() == 0 {

Review Comment:
   This leans on the 1-byte `Read()` having populated `Size()` as a side effect. pierrec parses the frame descriptor on the first Read today, but that ordering isn't part of the documented contract. A lazy-parsing version, or a different lz4 library, would return `Size() == 0` for a perfectly valid frame and we'd reject every Trino-written file.
   
   `Size() == 0` also can't tell "no content-size field" (the case the spec actually prohibits) from "field present, declared as 0", so the "missing content size" error can misfire on the degenerate case.
   
   I'd read the frame descriptor explicitly via the lz4 header API instead of the peek-and-reconstruct dance, or just decompress into a bounded buffer and let the `LimitedReader` be the only size gate. wdyt?



##########
puffin/puffin_test.go:
##########
@@ -619,6 +638,60 @@ func TestReaderRejectsTrailingFooterData(t *testing.T) {
 	}
 }
 
+func TestReaderReadsLZ4CompressedFooter(t *testing.T) {
+	data := fileWithCompressedFooterPayload(t, []byte(`{"blobs":[],"properties":{"source":"lz4"}}`))
+
+	r, err := puffin.NewReader(bytes.NewReader(data))
+	require.NoError(t, err)
+	assert.Equal(t, "lz4", r.Properties()["source"])
+	assert.Empty(t, r.Blobs())
+}
+
+func TestReaderRejectsLZ4CompressedFooterWithoutContentSize(t *testing.T) {
+	payload := []byte(`{"blobs":[]}`)

Review Comment:
   This inlines the same frame-building as `fileWithCompressedFooterPayload` minus the `SizeOption` call, so if the helper ever regressed in a way that dropped `SizeOption` this test would still pass green. I'd factor out a `fileWithCompressedFooterPayloadNoSize(t, payload)` so the two stay in sync.



##########
puffin/puffin_reader.go:
##########
@@ -224,11 +238,32 @@ func (r *Reader) readFooter() error {
 	}
 
 	payloadReader := io.NewSectionReader(r.r, footerStart+MagicSize, payloadSize)
-	decoder := json.NewDecoder(payloadReader)
+	var footerReader io.Reader = payloadReader
+	if flags&FooterFlagCompressed != 0 {
+		lz4Reader := lz4.NewReader(payloadReader)
+		var firstByte [1]byte
+		n, err := lz4Reader.Read(firstByte[:])
+		if err != nil && !errors.Is(err, io.EOF) {
+			return fmt.Errorf("puffin: initialize compressed footer: %w", err)
+		}
+		if lz4Reader.Size() == 0 {
+			return errors.New("puffin: compressed footer LZ4 frame is missing content size")
+		}
+		if int64(lz4Reader.Size()) > r.maxFooterSize {
+			return fmt.Errorf("puffin: footer exceeds maximum size %d", r.maxFooterSize)
+		}
+		footerReader = io.MultiReader(bytes.NewReader(firstByte[:n]), lz4Reader)
+	}
+
+	limitedFooter := &io.LimitedReader{R: footerReader, N: r.maxFooterSize + 1}

Review Comment:
   Small edge: if a caller passes `WithMaxFooterSize(math.MaxInt64)`, `N: r.maxFooterSize + 1` overflows to `MinInt64`, `LimitedReader` treats `N <= 0` as an empty stream, and the decoder fails with "unexpected end of JSON input", which points nowhere near the actual misconfig. Worth capping `N` at the sentinel when `maxFooterSize` is near `MaxInt64`.



##########
puffin/puffin_reader.go:
##########
@@ -224,11 +238,32 @@ func (r *Reader) readFooter() error {
 	}
 
 	payloadReader := io.NewSectionReader(r.r, footerStart+MagicSize, payloadSize)
-	decoder := json.NewDecoder(payloadReader)
+	var footerReader io.Reader = payloadReader
+	if flags&FooterFlagCompressed != 0 {
+		lz4Reader := lz4.NewReader(payloadReader)
+		var firstByte [1]byte
+		n, err := lz4Reader.Read(firstByte[:])
+		if err != nil && !errors.Is(err, io.EOF) {
+			return fmt.Errorf("puffin: initialize compressed footer: %w", err)
+		}
+		if lz4Reader.Size() == 0 {
+			return errors.New("puffin: compressed footer LZ4 frame is missing content size")
+		}
+		if int64(lz4Reader.Size()) > r.maxFooterSize {
+			return fmt.Errorf("puffin: footer exceeds maximum size %d", r.maxFooterSize)
+		}
+		footerReader = io.MultiReader(bytes.NewReader(firstByte[:n]), lz4Reader)
+	}
+
+	limitedFooter := &io.LimitedReader{R: footerReader, N: r.maxFooterSize + 1}
+	decoder := json.NewDecoder(limitedFooter)
 	var footer Footer
 	if err := decoder.Decode(&footer); err != nil {
 		return fmt.Errorf("puffin: decode footer JSON: %w", err)
 	}
+	if limitedFooter.N == 0 {

Review Comment:
   `json.Decoder` reads ahead into its own buffer, so by the time `Decode` returns `N` may already be 0 because the decoder pulled in trailing bytes past the JSON value, not because the footer itself overflowed. In that case this reports "footer exceeds maximum size" when the real problem is trailing content, and the deliberate trailing-content check just below never gets a chance to fire.
   
   I'd move the `N == 0` guard to wrap the trailing-content read rather than sit between `Decode` and it, or check `decoder.Buffered()` to tell the two apart.



##########
puffin/puffin_test.go:
##########
@@ -619,6 +638,60 @@ func TestReaderRejectsTrailingFooterData(t *testing.T) {
 	}
 }
 
+func TestReaderReadsLZ4CompressedFooter(t *testing.T) {
+	data := fileWithCompressedFooterPayload(t, []byte(`{"blobs":[],"properties":{"source":"lz4"}}`))
+
+	r, err := puffin.NewReader(bytes.NewReader(data))
+	require.NoError(t, err)
+	assert.Equal(t, "lz4", r.Properties()["source"])
+	assert.Empty(t, r.Blobs())
+}
+
+func TestReaderRejectsLZ4CompressedFooterWithoutContentSize(t *testing.T) {
+	payload := []byte(`{"blobs":[]}`)
+	var compressed bytes.Buffer
+	writer := lz4.NewWriter(&compressed)
+	_, err := writer.Write(payload)
+	require.NoError(t, err)
+	require.NoError(t, writer.Close())
+
+	data := append([]byte("PFA1PFA1"), compressed.Bytes()...)
+	trailer := make([]byte, 12)
+	binary.LittleEndian.PutUint32(trailer[:4], uint32(compressed.Len()))
+	binary.LittleEndian.PutUint32(trailer[4:8], puffin.FooterFlagCompressed)
+	copy(trailer[8:], "PFA1")
+	data = append(data, trailer...)
+
+	_, err = puffin.NewReader(bytes.NewReader(data))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "missing content size")
+}
+
+func TestReaderEnforcesFooterSizeLimit(t *testing.T) {
+	t.Run("json payload", func(t *testing.T) {
+		data := fileWithFooterPayload([]byte(`{"blobs":[],"properties":{"large":"0123456789"}}`))
+		_, err := puffin.NewReader(bytes.NewReader(data), puffin.WithMaxFooterSize(16))
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "footer")

Review Comment:
   `assert.Contains(t, err.Error(), "footer")` would also pass on unrelated errors like "invalid footer start magic". The two sibling subtests assert the specific "footer exceeds maximum size", so I'd match that here too.



##########
puffin/puffin_test.go:
##########
@@ -619,6 +638,60 @@ func TestReaderRejectsTrailingFooterData(t *testing.T) {
 	}
 }
 
+func TestReaderReadsLZ4CompressedFooter(t *testing.T) {
+	data := fileWithCompressedFooterPayload(t, []byte(`{"blobs":[],"properties":{"source":"lz4"}}`))
+
+	r, err := puffin.NewReader(bytes.NewReader(data))
+	require.NoError(t, err)
+	assert.Equal(t, "lz4", r.Properties()["source"])
+	assert.Empty(t, r.Blobs())
+}
+
+func TestReaderRejectsLZ4CompressedFooterWithoutContentSize(t *testing.T) {
+	payload := []byte(`{"blobs":[]}`)
+	var compressed bytes.Buffer
+	writer := lz4.NewWriter(&compressed)
+	_, err := writer.Write(payload)
+	require.NoError(t, err)
+	require.NoError(t, writer.Close())
+
+	data := append([]byte("PFA1PFA1"), compressed.Bytes()...)
+	trailer := make([]byte, 12)
+	binary.LittleEndian.PutUint32(trailer[:4], uint32(compressed.Len()))
+	binary.LittleEndian.PutUint32(trailer[4:8], puffin.FooterFlagCompressed)
+	copy(trailer[8:], "PFA1")
+	data = append(data, trailer...)
+
+	_, err = puffin.NewReader(bytes.NewReader(data))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "missing content size")
+}
+
+func TestReaderEnforcesFooterSizeLimit(t *testing.T) {
+	t.Run("json payload", func(t *testing.T) {
+		data := fileWithFooterPayload([]byte(`{"blobs":[],"properties":{"large":"0123456789"}}`))
+		_, err := puffin.NewReader(bytes.NewReader(data), puffin.WithMaxFooterSize(16))
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "footer")
+	})
+
+	t.Run("trailing content beyond limit", func(t *testing.T) {
+		payload := append([]byte(`{"blobs":[]}`), bytes.Repeat([]byte{' '}, 32)...)
+		data := fileWithFooterPayload(payload)
+		_, err := puffin.NewReader(bytes.NewReader(data), puffin.WithMaxFooterSize(16))
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "footer exceeds maximum size")
+	})
+
+	t.Run("compressed advertised size", func(t *testing.T) {

Review Comment:
   All three size-limit subtests hit the early advertised-size rejection; none exercise the `LimitedReader` as the actual backstop, i.e. a frame that advertises a size within the limit but decompresses to more than it. That's the one path that matters if the early check is ever bypassed (see the uint64 issue above), and lz4 won't always validate actual output against the advertised size.
   
   I'd add a subtest that builds a frame with `SizeOption(<= maxFooterSize)` but writes a payload that decompresses past the limit, and assert it still errors with "footer exceeds maximum size".



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org