zeroshade commented on code in PR #1390:
URL: https://github.com/apache/arrow-go/pull/1390#discussion_r4327780998
##########
parquet/encryption_read_config_test.go:
##########
@@ -607,3 +613,60 @@ 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()
+ assert.EqualValues(t, 100, tbl.NumRows())
Review Comment:
Please also check the values, as the C++ test for this fixture does
(`key_management_test.cc`, `ReadParquetMRExternalKeyMaterialFile`): column 0
(`integers`) equals the row index, and column 1 (`strings`) is `prefixes[row %
10] + strconv.Itoa(row)` with prefixes `"a"` … `"j"`. Only checking 100 rows
and 2 columns would also pass for a reader that lands on the wrong bytes but
still decodes the right number of values.
##########
parquet/file/page_reader_internal_test.go:
##########
@@ -83,3 +88,42 @@ 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)
+
+ var stream bytes.Buffer
+ enc.Encrypt(&stream, append(thriftBytes, make([]byte, 76)...))
Review Comment:
Nit: a comment that the 76 zero bytes bring the plaintext to 100 bytes,
matching parquet-mr's `INIT_MEM_ALLOC_ENCR_BUFFER` initial buffer, would
explain the number.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
