This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git


The following commit(s) were added to refs/heads/main by this push:
     new de0c94e3 fix(parquet/file): read encrypted page headers written by 
parquet-mr (#1390)
de0c94e3 is described below

commit de0c94e39f4c05a997f7d474ccbc5768839d3b36
Author: Sankalp Thakur <[email protected]>
AuthorDate: Thu Oct 22 01:11:46 2026 +0400

    fix(parquet/file): read encrypted page headers written by parquet-mr (#1390)
    
    ### Rationale for this change
    
    arrow-go cannot read a Parquet file that parquet-mr (Spark, Hive,
    Iceberg, anything on parquet-java) wrote with modular encryption.
    `file.OpenParquetFile` succeeds and the first page read fails:
    
    ```
    decrypting data page: encrypted payload length prefix is 4204214195, 
available length is 40: unexpected EOF
    ```
    
    I hit it with `data/external_key_material_java.parquet.encrypted` from
    parquet-testing.
    
    parquet-mr zero-pads every encrypted thrift module to the size of its
    buffer: `Util.java` encrypts the whole backing array of its
    `TMemoryBuffer`, not only the bytes written. So an encrypted page header
    is the thrift struct followed by zeros, all inside one length-prefixed
    module. `readPageHeader` (`parquet/file/page_reader.go:796`) decrypts
    the module, parses the thrift, and then advances the stream by the
    amount of thrift it consumed plus the cipher overhead (line 837). That
    is only right when nothing follows the struct. With padding it stops
    part way through the module, and the next read takes ciphertext for page
    data. Files from arrow-go and from C++ (the ones in parquet-testing) are
    not padded, which is why the existing tests never saw it.
    
    ### What changes are included in this PR?
    
    For an encrypted header, `readPageHeader` now skips the whole framed
    ciphertext (the 4-byte length prefix plus the module, which
    `encryption.FramedCiphertext` already returns). The unencrypted path is
    unchanged.
    
    ### Are these changes tested?
    
    Yes. Both new tests fail on main and pass with the change.
    
    - `TestDecryptParquetMrExternalKeyMaterial`
    (`parquet/encryption_read_config_test.go`) reads the parquet-testing
    file above and expects 100 rows and 2 columns. That file uses external
    key material with double wrapping and arrow-go has no key toolkit, so
    the test carries a small key retriever that unwraps the keys with the
    same master keys the C++ tests use (`kf`, `kc1`, `kc2`). On main it
    fails with the error above.
    - `TestReadEncryptedPageHeaderWithTrailingPadding`
    (`parquet/file/page_reader_internal_test.go`) encrypts a page header
    whose plaintext is the thrift struct plus 76 zero bytes, puts page data
    after it, and checks that `readPageHeader` leaves the stream at the page
    data. On main it is left inside the module.
    
    Commands, with `PARQUET_TEST_DATA=$PWD/parquet-testing/data`:
    
    - `go test ./parquet -run
    'TestFileEncryptionDecryption|TestDecryptParquetMrExternalKeyMaterial'
    -count=1`: ok (the existing C++-written configurations still pass)
    - `go test ./parquet/file -run
    'TestReadEncryptedPageHeaderWithTrailingPadding|TestPageReader'
    -count=1`: ok
    - `go vet ./parquet ./parquet/file` and `gofmt -l` on the changed files:
    clean
    
    The fixture is the one in the parquet-testing submodule this repo
    already pins. Its README says it was encrypted with parquet-mr using
    external key material and the in-memory KMS client. No fixture files
    were added.
    
    ### Are there any user-facing changes?
    
    Yes. Encrypted files written by parquet-mr / parquet-java can be read.
    Nothing changes for files written by arrow-go or C++.
    
    Related to #566. That issue reports that cross-library encryption does
    not work in general; this covers the parquet-mr to Go read path only.
    
    ### AI/LLM disclosure
    
    - AI coding tools (Claude Code) were used to find and write this fix,
    the two tests and this description.
    - The failure was reproduced on main, the change was reviewed, and the
    commands above were run locally before submitting.
---
 parquet/encryption_read_config_test.go    | 87 +++++++++++++++++++++++++++++++
 parquet/file/page_reader.go               | 14 +++--
 parquet/file/page_reader_internal_test.go | 50 ++++++++++++++++++
 3 files changed, 148 insertions(+), 3 deletions(-)

diff --git a/parquet/encryption_read_config_test.go 
b/parquet/encryption_read_config_test.go
index 5bf15f19..06508b7d 100644
--- a/parquet/encryption_read_config_test.go
+++ b/parquet/encryption_read_config_test.go
@@ -18,17 +18,25 @@ package parquet_test
 
 import (
        "bytes"
+       "context"
+       "crypto/aes"
+       "crypto/cipher"
+       "encoding/base64"
        "encoding/binary"
+       "encoding/json"
        "fmt"
        "os"
        "path"
+       "strconv"
        "testing"
 
+       "github.com/apache/arrow-go/v18/arrow/array"
        "github.com/apache/arrow-go/v18/arrow/memory"
        "github.com/apache/arrow-go/v18/parquet"
        "github.com/apache/arrow-go/v18/parquet/file"
        "github.com/apache/arrow-go/v18/parquet/internal/encryption"
        "github.com/apache/arrow-go/v18/parquet/metadata"
+       "github.com/apache/arrow-go/v18/parquet/pqarrow"
        "github.com/stretchr/testify/assert"
        "github.com/stretchr/testify/require"
        "github.com/stretchr/testify/suite"
@@ -607,3 +615,82 @@ func (d *TestDecryptionSuite) TestDecryption() {
                })
        }
 }
+
+// javaKeyRetriever unwraps keys the way parquet-mr's KeyToolkit does for a 
file written
+// with external key material and double wrapping: the KMS master key unwraps 
a key
+// encryption key (KEK), and the KEK unwraps the data key.
+type javaKeyRetriever struct {
+       t        *testing.T
+       material map[string]string // key reference -> key material JSON
+}
+
+func (r javaKeyRetriever) unwrap(wrapped string, key, aad []byte) []byte {
+       raw, err := base64.StdEncoding.DecodeString(wrapped)
+       require.NoError(r.t, err)
+       block, err := aes.NewCipher(key)
+       require.NoError(r.t, err)
+       gcm, err := cipher.NewGCM(block)
+       require.NoError(r.t, err)
+       plain, err := gcm.Open(nil, raw[:gcm.NonceSize()], 
raw[gcm.NonceSize():], aad)
+       require.NoError(r.t, err)
+       return plain
+}
+
+func (r javaKeyRetriever) GetKey(keyMetadata []byte) string {
+       var ref struct{ KeyReference string }
+       require.NoError(r.t, json.Unmarshal(keyMetadata, &ref))
+       var km struct{ MasterKeyID, WrappedDEK, KeyEncryptionKeyID, WrappedKEK 
string }
+       require.NoError(r.t, 
json.Unmarshal([]byte(r.material[ref.KeyReference]), &km))
+
+       masterKeys := map[string]string{"kf": FooterEncryptionKey, "kc1": 
ColumnEncryptionKey1, "kc2": ColumnEncryptionKey2}
+       kek := r.unwrap(km.WrappedKEK, []byte(masterKeys[km.MasterKeyID]), 
[]byte(km.MasterKeyID))
+       kekID, err := base64.StdEncoding.DecodeString(km.KeyEncryptionKeyID)
+       require.NoError(r.t, err)
+       return string(r.unwrap(km.WrappedDEK, kek, kekID))
+}
+
+// parquet-mr zero-pads every encrypted thrift module, page headers included, 
out to its
+// buffer size, so reading its pages needs the reader to skip the whole 
encrypted header.
+func TestDecryptParquetMrExternalKeyMaterial(t *testing.T) {
+       dir := getDataDir()
+       raw, err := os.ReadFile(path.Join(dir, 
"_KEY_MATERIAL_FOR_external_key_material_java.parquet.encrypted.json"))
+       require.NoError(t, err)
+       var material map[string]string
+       require.NoError(t, json.Unmarshal(raw, &material))
+
+       props := parquet.NewReaderProperties(memory.DefaultAllocator)
+       props.FileDecryptProps = 
parquet.NewFileDecryptionProperties(parquet.WithKeyRetriever(javaKeyRetriever{t,
 material}))
+       rdr, err := file.OpenParquetFile(path.Join(dir, 
"external_key_material_java.parquet.encrypted"), false, 
file.WithReadProps(props))
+       require.NoError(t, err)
+       defer rdr.Close()
+
+       fr, err := pqarrow.NewFileReader(rdr, pqarrow.ArrowReadProperties{}, 
memory.DefaultAllocator)
+       require.NoError(t, err)
+       tbl, err := fr.ReadTable(context.Background())
+       require.NoError(t, err)
+       defer tbl.Release()
+       require.EqualValues(t, 100, tbl.NumRows())
+       require.EqualValues(t, 2, tbl.NumCols())
+
+       // Same values the C++ ReadParquetMRExternalKeyMaterialFile test 
checks: column 0 is
+       // the row index and column 1 is prefixes[row % 10] followed by the row 
index.
+       prefixes := []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}
+       row := 0
+       for _, chunk := range tbl.Column(0).Data().Chunks() {
+               ints := chunk.(*array.Int32)
+               for i := 0; i < ints.Len(); i++ {
+                       assert.EqualValues(t, row, ints.Value(i), 
"integers[%d]", row)
+                       row++
+               }
+       }
+       assert.Equal(t, 100, row)
+       row = 0
+       for _, chunk := range tbl.Column(1).Data().Chunks() {
+               strs := chunk.(*array.String)
+               for i := 0; i < strs.Len(); i++ {
+                       assert.Equal(t, 
prefixes[row%len(prefixes)]+strconv.Itoa(row), strs.Value(i), "strings[%d]", 
row)
+                       row++
+               }
+       }
+       assert.Equal(t, 100, row)
+}
diff --git a/parquet/file/page_reader.go b/parquet/file/page_reader.go
index 2b748282..66c33968 100644
--- a/parquet/file/page_reader.go
+++ b/parquet/file/page_reader.go
@@ -805,7 +805,7 @@ func (p *serializedPageReader) readPageHeader(rd 
parquet.BufferedReader, hdr *fo
                        return io.EOF
                }
 
-               extra := 0
+               framedLen := 0
                if p.cryptoCtx.MetaDecryptor != nil {
                        p.updateDecryption(p.cryptoCtx.MetaDecryptor, 
encryption.DictPageHeaderModule, p.dataPageHeaderAad)
                        ciphertext, frameErr := 
encryption.FramedCiphertext(view)
@@ -823,7 +823,7 @@ func (p *serializedPageReader) readPageHeader(rd 
parquet.BufferedReader, hdr *fo
                        if err != nil {
                                return fmt.Errorf("decrypting page header: %w", 
err)
                        }
-                       extra = p.cryptoCtx.MetaDecryptor.CiphertextSizeDelta()
+                       framedLen = len(ciphertext)
                }
 
                remaining, err := thrift.DeserializeThrift(hdr, view)
@@ -834,7 +834,15 @@ func (p *serializedPageReader) readPageHeader(rd 
parquet.BufferedReader, hdr *fo
                        }
                        continue
                }
-               rd.Discard(len(view) - int(remaining) + extra)
+               if p.cryptoCtx.MetaDecryptor != nil {
+                       // An encrypted header is a length-prefixed module, so 
skip the whole frame.
+                       // The decrypted payload can hold bytes after the 
thrift struct (parquet-mr
+                       // zero-pads it to its buffer size), which means the 
amount of thrift read
+                       // does not tell us how many bytes the module takes up 
in the file.
+                       rd.Discard(framedLen)
+               } else {
+                       rd.Discard(len(view) - int(remaining))
+               }
                break
        }
        return nil
diff --git a/parquet/file/page_reader_internal_test.go 
b/parquet/file/page_reader_internal_test.go
index a416fa6a..0f83b7f8 100644
--- a/parquet/file/page_reader_internal_test.go
+++ b/parquet/file/page_reader_internal_test.go
@@ -18,12 +18,17 @@ package file
 
 import (
        "bytes"
+       "context"
        "encoding/binary"
+       "io"
        "math"
        "testing"
 
        "github.com/apache/arrow-go/v18/arrow/memory"
+       "github.com/apache/arrow-go/v18/parquet"
+       "github.com/apache/arrow-go/v18/parquet/internal/encryption"
        format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet"
+       "github.com/apache/arrow-go/v18/parquet/internal/thrift"
        "github.com/stretchr/testify/require"
 )
 
@@ -83,3 +88,48 @@ func TestReadLevelDataBitPacked(t *testing.T) {
        require.Equal(t, data, out.Bytes())
        out.Release()
 }
+
+// parquet-mr zero-pads every encrypted thrift module out to its buffer size, 
so a
+// decrypted page header is the thrift struct followed by zero bytes. The 
reader has
+// to move past the whole encrypted module, not just the thrift it parsed.
+func TestReadEncryptedPageHeaderWithTrailingPadding(t *testing.T) {
+       const key = "0123456789012345"
+       encProps := parquet.NewFileEncryptionProperties(key)
+       enc := encryption.NewFileEncryptor(encProps, 
memory.DefaultAllocator).GetFooterEncryptor()
+       enc.UpdateAad(encryption.CreateModuleAad(encProps.FileAad(), 
encryption.DictPageHeaderModule, 0, 0, -1))
+
+       hdr := format.NewPageHeader()
+       hdr.Type = format.PageType_DICTIONARY_PAGE
+       hdr.CompressedPageSize, hdr.UncompressedPageSize = 10, 10
+       hdr.DictionaryPageHeader = format.NewDictionaryPageHeader()
+       thriftBytes, err := 
thrift.NewThriftSerializer().Write(context.Background(), hdr)
+       require.NoError(t, err)
+
+       // parquet-mr serializes each encrypted module into a TMemoryBuffer 
that starts at
+       // INIT_MEM_ALLOC_ENCR_BUFFER (100) bytes and encrypts the whole 
backing array, so a
+       // small header is followed by zeros up to 100 bytes. Pad the same way.
+       const parquetMrInitBuffer = 100
+       padding := make([]byte, parquetMrInitBuffer-len(thriftBytes))
+
+       var stream bytes.Buffer
+       enc.Encrypt(&stream, append(thriftBytes, padding...))
+       stream.WriteString("page data!")
+
+       dec := 
encryption.NewFileDecryptor(parquet.NewFileDecryptionProperties(parquet.WithFooterKey(key)),
+               encProps.FileAad(), parquet.AesGcm, "", 
memory.DefaultAllocator).GetFooterDecryptor()
+       rd, err := parquet.NewReaderProperties(memory.DefaultAllocator).
+               GetStream(bytes.NewReader(stream.Bytes()), 0, 
int64(stream.Len()))
+       require.NoError(t, err)
+
+       p := &serializedPageReader{
+               maxPageHeaderSize: defaultMaxPageHeaderSize,
+               cryptoCtx:         
CryptoContext{StartDecryptWithDictionaryPage: true, MetaDecryptor: dec},
+       }
+       got := format.NewPageHeader()
+       require.NoError(t, p.readPageHeader(rd, got))
+       require.Equal(t, format.PageType_DICTIONARY_PAGE, got.Type)
+
+       rest, err := io.ReadAll(rd)
+       require.NoError(t, err)
+       require.Equal(t, "page data!", string(rest))
+}

Reply via email to