Skip to content

API, Core, Spark 4.1: Add key_metadata and encrypted write support for StatisticsFile #17417

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gaborkaszab wants to merge 1 commit into
base: main
Choose a base branch
from
Open

API, Core, Spark 4.1: Add key_metadata and encrypted write support for StatisticsFile #17417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions api/src/main/java/org/apache/iceberg/StatisticsFile.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,11 @@ public interface StatisticsFile {
/** Size of the Puffin footer. */
long fileFooterSizeInBytes();

/** ID of the encryption key used for this file, or null if the file is stored in plain text. */
default String keyId() {
return null;
}

/** List of statistics contained in the file. Never null. */
List<BlobMetadata> blobMetadata();
}
2 changes: 1 addition & 1 deletion core/src/main/java/org/apache/iceberg/BaseManifestListFile.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,6 @@ public String encryptionKeyID() {

@Override
public ByteBuffer decryptKeyMetadata(EncryptionManager em) {
return EncryptionUtil.decryptManifestListKeyMetadata(this, em);
return EncryptionUtil.decryptKeyMetadata(encryptionKeyID, em);
}
}
27 changes: 26 additions & 1 deletion core/src/main/java/org/apache/iceberg/GenericStatisticsFile.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,20 +29,37 @@ public class GenericStatisticsFile implements StatisticsFile {
private final String path;
private final long fileSizeInBytes;
private final long fileFooterSizeInBytes;
private final String keyId;
private final List<BlobMetadata> blobMetadata;

/**
* @deprecated will be removed in 1.13.0. Use {@link #GenericStatisticsFile(long, String, long,
* long, String, List)} instead.
*/
@Deprecated
public GenericStatisticsFile(
long snapshotId,
String path,
long fileSizeInBytes,
long fileFooterSizeInBytes,
List<BlobMetadata> blobMetadata) {
this(snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, null, blobMetadata);
}

public GenericStatisticsFile(
long snapshotId,
String path,
long fileSizeInBytes,
long fileFooterSizeInBytes,
String keyId,
List<BlobMetadata> blobMetadata) {
Preconditions.checkNotNull(path, "path is null");
Preconditions.checkNotNull(blobMetadata, "blobMetadata is null");
this.snapshotId = snapshotId;
this.path = path;
this.fileSizeInBytes = fileSizeInBytes;
this.fileFooterSizeInBytes = fileFooterSizeInBytes;
this.keyId = keyId;
this.blobMetadata = ImmutableList.copyOf(blobMetadata);
}

Expand All @@ -66,6 +83,11 @@ public long fileFooterSizeInBytes() {
return fileFooterSizeInBytes;
}

@Override
public String keyId() {
return keyId;
}

@Override
public List<BlobMetadata> blobMetadata() {
return blobMetadata;
Expand All @@ -84,12 +106,14 @@ public boolean equals(Object o) {
&& fileSizeInBytes == that.fileSizeInBytes
&& fileFooterSizeInBytes == that.fileFooterSizeInBytes
&& Objects.equals(path, that.path)
&& Objects.equals(keyId, that.keyId)
&& Objects.equals(blobMetadata, that.blobMetadata);
}

@Override
public int hashCode() {
return Objects.hash(snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, blobMetadata);
return Objects.hash(
snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, keyId, blobMetadata);
}

@Override
Expand All @@ -99,6 +123,7 @@ public String toString() {
.add("path='" + path + "'")
.add("fileSizeInBytes=" + fileSizeInBytes)
.add("fileFooterSizeInBytes=" + fileFooterSizeInBytes)
.add("keyId=" + keyId)
.add("blobMetadata=" + blobMetadata)
.toString();
}
Expand Down
2 changes: 1 addition & 1 deletion core/src/main/java/org/apache/iceberg/ManifestListWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ public Long nextRowId() {
public ManifestListFile toManifestListFile() {
if (manifestListKeyMetadata != null && manifestListKeyMetadata.encryptionKey() != null) {
String manifestListKeyID =
standardEncryptionManager.addManifestListKeyMetadata(
standardEncryptionManager.addKeyMetadata(
manifestListKeyMetadata.copyWithLength(writer.length()));
return new BaseManifestListFile(outputFile.location(), manifestListKeyID);
} else {
Expand Down
1 change: 1 addition & 0 deletions core/src/main/java/org/apache/iceberg/RewriteTablePathUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,7 @@ private static List<StatisticsFile> updatePathInStatisticsFiles(
newPath(existing.path(), sourcePrefix, targetPrefix),
existing.fileSizeInBytes(),
existing.fileFooterSizeInBytes(),
existing.keyId(),
existing.blobMetadata()))
.collect(Collectors.toList());
}
Expand Down
7 changes: 6 additions & 1 deletion core/src/main/java/org/apache/iceberg/StatisticsFileParser.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ public class StatisticsFileParser {
private static final String STATISTICS_PATH = "statistics-path";
private static final String FILE_SIZE_IN_BYTES = "file-size-in-bytes";
private static final String FILE_FOOTER_SIZE_IN_BYTES = "file-footer-size-in-bytes";
private static final String KEY_ID = "key-id";
private static final String BLOB_METADATA = "blob-metadata";
private static final String TYPE = "type";
private static final String SEQUENCE_NUMBER = "sequence-number";
Expand All @@ -57,6 +58,9 @@ public static void toJson(StatisticsFile statisticsFile, JsonGenerator generator
generator.writeStringField(STATISTICS_PATH, statisticsFile.path());
generator.writeNumberField(FILE_SIZE_IN_BYTES, statisticsFile.fileSizeInBytes());
generator.writeNumberField(FILE_FOOTER_SIZE_IN_BYTES, statisticsFile.fileFooterSizeInBytes());
if (statisticsFile.keyId() != null) {
generator.writeStringField(KEY_ID, statisticsFile.keyId());
}
generator.writeArrayFieldStart(BLOB_METADATA);
for (BlobMetadata blobMetadata : statisticsFile.blobMetadata()) {
toJson(blobMetadata, generator);
Expand All @@ -70,6 +74,7 @@ static StatisticsFile fromJson(JsonNode node) {
String path = JsonUtil.getString(STATISTICS_PATH, node);
long fileSizeInBytes = JsonUtil.getLong(FILE_SIZE_IN_BYTES, node);
long fileFooterSizeInBytes = JsonUtil.getLong(FILE_FOOTER_SIZE_IN_BYTES, node);
String keyId = JsonUtil.getStringOrNull(KEY_ID, node);
ImmutableList.Builder<BlobMetadata> blobMetadata = ImmutableList.builder();
JsonNode blobsJson = node.get(BLOB_METADATA);
Preconditions.checkArgument(
Expand All @@ -80,7 +85,7 @@ static StatisticsFile fromJson(JsonNode node) {
blobMetadata.add(blobMetadataFromJson(blobJson));
}
return new GenericStatisticsFile(
snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, blobMetadata.build());
snapshotId, path, fileSizeInBytes, fileFooterSizeInBytes, keyId, blobMetadata.build());
}

private static void toJson(BlobMetadata blobMetadata, JsonGenerator generator)
Expand Down
37 changes: 25 additions & 12 deletions core/src/main/java/org/apache/iceberg/encryption/EncryptionUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,19 +140,32 @@ public static ByteBuffer setFileLength(ByteBuffer keyMetadata, long fileLength)
* @param manifestList a ManifestListFile
* @param em the table's EncryptionManager
* @return a decrypted key metadata buffer
* @deprecated will be removed in 1.13.0. Use {@link #decryptKeyMetadata(String,
* EncryptionManager)} instead.
*/
@Deprecated
public static ByteBuffer decryptManifestListKeyMetadata(
ManifestListFile manifestList, EncryptionManager em) {
return decryptKeyMetadata(manifestList.encryptionKeyID(), em);
}

/**
* Decrypt an encrypted key metadata.
*
* @param keyId the ID of the encrypted key
* @param em the table's EncryptionManager
* @return a decrypted key metadata buffer
*/
public static ByteBuffer decryptKeyMetadata(String keyId, EncryptionManager em) {
Preconditions.checkState(
em instanceof StandardEncryptionManager,
"Snapshot key metadata encryption requires a StandardEncryptionManager");
"Key metadata decryption requires a StandardEncryptionManager");
StandardEncryptionManager sem = (StandardEncryptionManager) em;
String manifestListKeyId = manifestList.encryptionKeyID();
Map<String, EncryptedKey> encryptionKeys = sem.encryptionKeys();
EncryptedKey manifestListKey = encryptionKeys.get(manifestListKeyId);
ByteBuffer encryptedKeyMetadata = manifestListKey.encryptedKeyMetadata();
String keyEncryptionKeyID = manifestListKey.encryptedById();
ByteBuffer keyEncryptionKey = sem.encryptedByKey(manifestListKeyId);
EncryptedKey encryptedKey = encryptionKeys.get(keyId);
ByteBuffer encryptedKeyMetadata = encryptedKey.encryptedKeyMetadata();
String keyEncryptionKeyID = encryptedKey.encryptedById();
ByteBuffer keyEncryptionKey = sem.encryptedByKey(keyId);
String keyEncryptionKeyTimestamp =
encryptionKeys
.get(keyEncryptionKeyID)
Expand Down Expand Up @@ -182,22 +195,22 @@ public static Map<String, EncryptedKey> encryptionKeys(EncryptionManager em) {
}

/**
* Encrypts the key metadata for a manifest list.
* Encrypts an encryption key metadata.
*
* @param key key encryption key bytes
* @param keyTimestamp timestamp of the key encryption key
* @param mlkMetadata manifest list key metadata
* @param keyMetadata key metadata to encrypt
* @return encrypted key metadata
*/
static ByteBuffer encryptManifestListKeyMetadata(
ByteBuffer key, String keyTimestamp, EncryptionKeyMetadata mlkMetadata) {
static ByteBuffer encryptKeyMetadata(
ByteBuffer key, String keyTimestamp, EncryptionKeyMetadata keyMetadata) {
Ciphers.AesGcmEncryptor encryptor = new Ciphers.AesGcmEncryptor(ByteBuffers.toByteArray(key));
byte[] mlkMetadataBytes = ByteBuffers.toByteArray(mlkMetadata.buffer());
byte[] keyMetadataBytes = ByteBuffers.toByteArray(keyMetadata.buffer());

// Use key encryption key timestamp as AES GCM signature (AAD) of encryption - in order to
// prevent timestamp tampering attacks
byte[] encryptedKeyMetadata =
encryptor.encrypt(mlkMetadataBytes, keyTimestamp.getBytes(StandardCharsets.UTF_8));
encryptor.encrypt(keyMetadataBytes, keyTimestamp.getBytes(StandardCharsets.UTF_8));

return ByteBuffer.wrap(encryptedKeyMetadata);
}
Expand Down
36 changes: 25 additions & 11 deletions core/src/main/java/org/apache/iceberg/encryption/StandardEncryptionManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,36 +193,50 @@ private long currentTimeMillis() {
return System.currentTimeMillis() + testTimeShift;
}

ByteBuffer encryptedByKey(String manifestListKeyID) {
EncryptedKey encryptedKeyMetadata = encryptionKeys.get(manifestListKeyID);
ByteBuffer encryptedByKey(String keyID) {
EncryptedKey encryptedKeyMetadata = encryptionKeys.get(keyID);

Preconditions.checkState(
encryptedKeyMetadata != null,
"Cannot find manifest list key metadata with id %s",
manifestListKeyID);
encryptedKeyMetadata != null, "Cannot find key metadata with id %s", keyID);

Preconditions.checkArgument(
!encryptedKeyMetadata.encryptedById().equals(tableKeyId),
"%s is a key encryption key, not manifest list key metadata",
manifestListKeyID);
"%s is a key encryption key, not key metadata",
keyID);

return unwrappedKeyCache().get(encryptedKeyMetadata.encryptedById());
}

/**
* @deprecated will be removed in 1.13.0. Use {@link #addKeyMetadata(NativeEncryptionKeyMetadata)}
* instead.
*/
@Deprecated
public String addManifestListKeyMetadata(NativeEncryptionKeyMetadata keyMetadata) {
String manifestListKeyID = generateKeyId();
return addKeyMetadata(keyMetadata);
}

/**
* Wraps the given key metadata with the current key encryption key and registers it into this
* encryption manager.
*
* @param keyMetadata the key metadata to wrap
* @return the ID of the wrapped and registered encrypted key
*/
public String addKeyMetadata(NativeEncryptionKeyMetadata keyMetadata) {
String keyID = generateKeyId();
String keyEncryptionKeyID = keyEncryptionKeyID();
String keyEncryptionKeyTimestamp =
encryptionKeys.get(keyEncryptionKeyID).properties().get(KEY_TIMESTAMP);
ByteBuffer encryptedKeyMetadata =
EncryptionUtil.encryptManifestListKeyMetadata(
EncryptionUtil.encryptKeyMetadata(
unwrappedKeyCache().get(keyEncryptionKeyID), keyEncryptionKeyTimestamp, keyMetadata);
BaseEncryptedKey key =
new BaseEncryptedKey(manifestListKeyID, encryptedKeyMetadata, keyEncryptionKeyID, null);
new BaseEncryptedKey(keyID, encryptedKeyMetadata, keyEncryptionKeyID, null);

encryptionKeys.put(key.keyId(), key);

return manifestListKeyID;
return keyID;
}

private String generateKeyId() {
Expand Down
33 changes: 31 additions & 2 deletions ....1/spark/src/main/java/org/apache/iceberg/spark/actions/ComputeTableStatsSparkAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,11 @@
import org.apache.iceberg.TableOperations;
import org.apache.iceberg.actions.ComputeTableStats;
import org.apache.iceberg.actions.ImmutableComputeTableStats;
import org.apache.iceberg.encryption.EncryptedOutputFile;
import org.apache.iceberg.encryption.EncryptingFileIO;
import org.apache.iceberg.encryption.EncryptionKeyMetadata;
import org.apache.iceberg.encryption.NativeEncryptionKeyMetadata;
import org.apache.iceberg.encryption.StandardEncryptionManager;
import org.apache.iceberg.exceptions.RuntimeIOException;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.puffin.Blob;
Expand Down Expand Up @@ -110,21 +115,45 @@ private Result doExecute() {

private StatisticsFile writeStatsFile(List<Blob> blobs) {
LOG.info("Writing stats for table {} for snapshot {}", table.name(), snapshotId());
OutputFile outputFile = table.io().newOutputFile(outputPath());
EncryptingFileIO io = EncryptingFileIO.combine(table.io(), table.encryption());
EncryptedOutputFile encryptedOutputFile = io.newEncryptingOutputFile(outputPath());
OutputFile outputFile = encryptedOutputFile.encryptingOutputFile();

try (PuffinWriter writer = Puffin.write(outputFile).createdBy(appIdentifier()).build()) {
blobs.forEach(writer::add);
writer.finish();
long fileSize = writer.fileSize();
String encryptionKeyId = encryptKeyMetadata(encryptedOutputFile.keyMetadata(), fileSize);
return new GenericStatisticsFile(
snapshotId(),
outputFile.location(),
writer.fileSize(),
fileSize,
writer.footerSize(),
encryptionKeyId,
GenericBlobMetadata.from(writer.writtenBlobsMetadata()));
} catch (IOException e) {
throw new RuntimeIOException(e);
}
}

/**
* Encrypts the key metadata.
*
* @return the encryption key ID of the encrypted key metadata, or null if the file is written in
* plain text
*/
private String encryptKeyMetadata(EncryptionKeyMetadata keyMetadata, long fileSizeInBytes) {
if (!(keyMetadata instanceof NativeEncryptionKeyMetadata nativeKeyMetadata)) {
return null;
}

Preconditions.checkState(
table.encryption() instanceof StandardEncryptionManager,
"Encrypted statistics files require a StandardEncryptionManager");
StandardEncryptionManager em = (StandardEncryptionManager) table.encryption();
return em.addKeyMetadata(nativeKeyMetadata.copyWithLength(fileSizeInBytes));
}

private List<Blob> generateNDVBlobs() {
return NDVSketchUtil.generateBlobs(spark(), table, snapshot, columns());
}
Expand Down
Loading
Loading