From 9f1a17deb3396d8a5fcfa40822a3cb238dc77bb6 Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Fri, 1 Dec 2023 15:39:22 +0100 Subject: [PATCH 1/3] DRILL-8465. check iceberg input --- .../drill/exec/store/iceberg/IcebergWork.java | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java index 1762d4b81a0..9a839e461ad 100644 --- a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java +++ b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java @@ -33,8 +33,10 @@ import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.InputStream; import java.io.ObjectInputStream; import java.io.ObjectOutputStream; +import java.io.ObjectStreamClass; import java.util.Base64; import java.util.Objects; import java.util.StringJoiner; @@ -92,7 +94,8 @@ public IcebergWorkDeserializer() { public IcebergWork deserialize(JsonParser p, DeserializationContext ctxt) throws IOException { JsonNode node = p.getCodec().readTree(p); String scanTaskString = node.get(IcebergWorkSerializer.SCAN_TASK_FIELD).asText(); - try (ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(Base64.getDecoder().decode(scanTaskString)))) { + try (ObjectInputStream ois = new CombinedScanTaskObjectInputStream( + new ByteArrayInputStream(Base64.getDecoder().decode(scanTaskString)))) { Object scanTask = ois.readObject(); return new IcebergWork((CombinedScanTask) scanTask); } catch (ClassNotFoundException e) { @@ -103,6 +106,22 @@ public IcebergWork deserialize(JsonParser p, DeserializationContext ctxt) throws } } + private static class CombinedScanTaskObjectInputStream extends ObjectInputStream { + + CombinedScanTaskObjectInputStream(InputStream inputStream) throws IOException { + super(inputStream); + } + + @Override + protected Class resolveClass(ObjectStreamClass cls) throws IOException, ClassNotFoundException { + final Class resolvedClass = super.resolveClass(cls); + if (CombinedScanTask.class.isAssignableFrom(resolvedClass)) { + return resolvedClass; + } + throw new IOException("Rejected deserialization of unexpected class: " + cls.getName()); + } + } + /** * Special serializer for {@link IcebergWork} class that serializes * {@code scanTask} field to byte array string created using {@link java.io.Serializable} From 9a2b8f5297e7c5689e1b8ad05f4d0904f0cb53a5 Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Fri, 1 Dec 2023 16:32:44 +0100 Subject: [PATCH 2/3] refactor --- .../drill/exec/store/iceberg/IcebergWork.java | 26 +++++++++++++++---- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java index 9a839e461ad..b6ca0c951da 100644 --- a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java +++ b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java @@ -27,6 +27,7 @@ import com.fasterxml.jackson.databind.deser.std.StdDeserializer; import com.fasterxml.jackson.databind.ser.std.StdSerializer; import org.apache.iceberg.CombinedScanTask; +import org.apache.iceberg.ScanTask; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -94,7 +95,7 @@ public IcebergWorkDeserializer() { public IcebergWork deserialize(JsonParser p, DeserializationContext ctxt) throws IOException { JsonNode node = p.getCodec().readTree(p); String scanTaskString = node.get(IcebergWorkSerializer.SCAN_TASK_FIELD).asText(); - try (ObjectInputStream ois = new CombinedScanTaskObjectInputStream( + try (ObjectInputStream ois = new ScanTaskObjectInputStream( new ByteArrayInputStream(Base64.getDecoder().decode(scanTaskString)))) { Object scanTask = ois.readObject(); return new IcebergWork((CombinedScanTask) scanTask); @@ -106,19 +107,34 @@ public IcebergWork deserialize(JsonParser p, DeserializationContext ctxt) throws } } - private static class CombinedScanTaskObjectInputStream extends ObjectInputStream { + private static class ScanTaskObjectInputStream extends ObjectInputStream { - CombinedScanTaskObjectInputStream(InputStream inputStream) throws IOException { + ScanTaskObjectInputStream(InputStream inputStream) throws IOException { super(inputStream); } @Override protected Class resolveClass(ObjectStreamClass cls) throws IOException, ClassNotFoundException { + final String className = cls.getName(); + if (isValidPackage(className)) { + return super.resolveClass(cls); + } final Class resolvedClass = super.resolveClass(cls); - if (CombinedScanTask.class.isAssignableFrom(resolvedClass)) { + if ((resolvedClass.isArray() && + (resolvedClass.getComponentType().isPrimitive() || + isValidPackage(resolvedClass.getComponentType().getName()) || + ScanTask.class.isAssignableFrom(resolvedClass.getComponentType()))) + || resolvedClass.isPrimitive() + || ScanTask.class.isAssignableFrom(resolvedClass)) { return resolvedClass; } - throw new IOException("Rejected deserialization of unexpected class: " + cls.getName()); + throw new IOException("Rejected deserialization of unexpected class: " + className); + } + + private boolean isValidPackage(final String className) { + return className.startsWith("org.apache.iceberg") || + className.startsWith("org.apache.drill") || + className.startsWith("java"); } } From 7a642d1a5b0dc28e22fd562c710bc7e86ef60d0e Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Sun, 3 Dec 2023 21:58:51 +0100 Subject: [PATCH 3/3] config --- .../drill/exec/store/iceberg/IcebergWork.java | 13 +++++-------- .../format/IcebergFormatPluginConfig.java | 19 +++++++++++++++++-- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java index b6ca0c951da..796efc0fc70 100644 --- a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java +++ b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/IcebergWork.java @@ -27,7 +27,6 @@ import com.fasterxml.jackson.databind.deser.std.StdDeserializer; import com.fasterxml.jackson.databind.ser.std.StdSerializer; import org.apache.iceberg.CombinedScanTask; -import org.apache.iceberg.ScanTask; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -122,19 +121,17 @@ protected Class resolveClass(ObjectStreamClass cls) throws IOException, Class final Class resolvedClass = super.resolveClass(cls); if ((resolvedClass.isArray() && (resolvedClass.getComponentType().isPrimitive() || - isValidPackage(resolvedClass.getComponentType().getName()) || - ScanTask.class.isAssignableFrom(resolvedClass.getComponentType()))) - || resolvedClass.isPrimitive() - || ScanTask.class.isAssignableFrom(resolvedClass)) { + isValidPackage(resolvedClass.getComponentType().getName()))) + || resolvedClass.isPrimitive()) { return resolvedClass; } throw new IOException("Rejected deserialization of unexpected class: " + className); } private boolean isValidPackage(final String className) { - return className.startsWith("org.apache.iceberg") || - className.startsWith("org.apache.drill") || - className.startsWith("java"); + return className.startsWith("org.apache.iceberg.") || + className.startsWith("org.apache.drill.") || + className.startsWith("java."); } } diff --git a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/format/IcebergFormatPluginConfig.java b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/format/IcebergFormatPluginConfig.java index 9e53803883c..5acd9457026 100644 --- a/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/format/IcebergFormatPluginConfig.java +++ b/contrib/format-iceberg/src/main/java/org/apache/drill/exec/store/iceberg/format/IcebergFormatPluginConfig.java @@ -45,6 +45,8 @@ public class IcebergFormatPluginConfig implements FormatPluginConfig { private final Boolean ignoreResiduals; + private final Boolean allowAnyClassToBeLoaded; + private final Long snapshotId; private final Long snapshotAsOfTime; @@ -60,6 +62,7 @@ public IcebergFormatPluginConfig( this.caseSensitive = builder.caseSensitive; this.includeColumnStats = builder.includeColumnStats; this.ignoreResiduals = builder.ignoreResiduals; + this.allowAnyClassToBeLoaded = builder.allowAnyClassToBeLoaded; this.snapshotId = builder.snapshotId; this.snapshotAsOfTime = builder.snapshotAsOfTime; this.fromSnapshotId = builder.fromSnapshotId; @@ -100,6 +103,10 @@ public Boolean getIgnoreResiduals() { return this.ignoreResiduals; } + public Boolean getAllowAnyClassToBeLoaded() { + return this.allowAnyClassToBeLoaded; + } + public Long getSnapshotId() { return this.snapshotId; } @@ -130,6 +137,7 @@ public boolean equals(Object o) { && Objects.equals(caseSensitive, that.caseSensitive) && Objects.equals(includeColumnStats, that.includeColumnStats) && Objects.equals(ignoreResiduals, that.ignoreResiduals) + && Objects.equals(allowAnyClassToBeLoaded, that.allowAnyClassToBeLoaded) && Objects.equals(snapshotId, that.snapshotId) && Objects.equals(snapshotAsOfTime, that.snapshotAsOfTime) && Objects.equals(fromSnapshotId, that.fromSnapshotId) @@ -138,8 +146,8 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(properties, snapshot, caseSensitive, includeColumnStats, - ignoreResiduals, snapshotId, snapshotAsOfTime, fromSnapshotId, toSnapshotId); + return Objects.hash(properties, snapshot, caseSensitive, includeColumnStats, ignoreResiduals, + allowAnyClassToBeLoaded, snapshotId, snapshotAsOfTime, fromSnapshotId, toSnapshotId); } @JsonPOJOBuilder(withPrefix = "") @@ -152,6 +160,8 @@ public static class IcebergFormatPluginConfigBuilder { private Boolean ignoreResiduals; + private Boolean allowAnyClassToBeLoaded; + private Long snapshotId; private Long snapshotAsOfTime; @@ -180,6 +190,11 @@ public IcebergFormatPluginConfigBuilder ignoreResiduals(Boolean ignoreResiduals) return this; } + public IcebergFormatPluginConfigBuilder allowAnyClassToBeLoaded(Boolean allowAnyClassToBeLoaded) { + this.allowAnyClassToBeLoaded = allowAnyClassToBeLoaded; + return this; + } + public IcebergFormatPluginConfigBuilder snapshotId(Long snapshotId) { this.snapshotId = snapshotId; return this;