sankalpsthakur commented on code in PR #1390:
URL: https://github.com/apache/arrow-go/pull/1390#discussion_r4332815669
##########
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:
done in 4467f2c6: the test now checks every value the way
ReadParquetMRExternalKeyMaterialFile does, integers[row] == row and
strings[row] == prefixes[row % 10] + row with prefixes a through j.
##########
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:
done in 4467f2c6: the padding is now derived from parquet-mr's 100-byte
INIT_MEM_ALLOC_ENCR_BUFFER, with a comment, instead of the bare 76.
--
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]
