Add ability to mask object key prefixes in logs

Aiven-Open · Sep 28, 2023 · b024730 · b024730
1 parent f273491
commit b024730
Show file tree

Hide file tree

Showing 12 changed files with 204 additions and 70 deletions.
diff --git a/core/src/integration-test/java/io/aiven/kafka/tieredstorage/RemoteStorageManagerTest.java b/core/src/integration-test/java/io/aiven/kafka/tieredstorage/RemoteStorageManagerTest.java
@@ -480,7 +480,7 @@ void testFetchingSegmentFileNonExistent() throws IOException {
         );
         rsm.configure(config);
 
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("", false);
 
         // Ensure the manifest exists.
         writeManifest(objectKeyFactory);
@@ -510,7 +510,7 @@ void testFetchingSegmentManifestNotFound() {
         rsm.configure(config);
 
         // Make sure the exception is connected to the manifest file.
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("", false);
         final String expectedMessage = "Key "
             + objectKeyFactory.key(REMOTE_LOG_METADATA, ObjectKeyFactory.Suffix.MANIFEST)
             + " does not exists in storage";
@@ -535,7 +535,7 @@ void testFetchingIndexNonExistent(final IndexType indexType) throws IOException
         );
         rsm.configure(config);
 
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("", false);
 
         // Ensure the manifest exists.
         writeManifest(objectKeyFactory);
@@ -562,7 +562,7 @@ void testFetchingIndexManifestNotFound(final IndexType indexType) throws Storage
         rsm.configure(config);
 
         // Make sure the exception is connected to the manifest file.
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("", false);
         final String expectedMessage = "Key "
             + objectKeyFactory.key(REMOTE_LOG_METADATA, ObjectKeyFactory.Suffix.MANIFEST)
             + " does not exists in storage";

diff --git a/core/src/main/java/io/aiven/kafka/tieredstorage/ObjectKeyFactory.java b/core/src/main/java/io/aiven/kafka/tieredstorage/ObjectKeyFactory.java
@@ -19,6 +19,7 @@
 import java.text.NumberFormat;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.BiFunction;
 
 import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.server.log.remote.storage.RemoteLogSegmentId;
@@ -69,12 +70,17 @@ static Suffix fromIndexType(final RemoteStorageManager.IndexType indexType) {
     }
 
     private final String prefix;
+    private final BiFunction<String, String, ObjectKey> objectKeyConstructor;
 
     /**
      * @param prefix the prefix to add to all created keys.
+     * @param maskPrefix whether to mask the prefix in {@code toString()}.
      */
-    public ObjectKeyFactory(final String prefix) {
+    public ObjectKeyFactory(final String prefix, final boolean maskPrefix) {
         this.prefix = prefix == null ? "" : prefix;
+        this.objectKeyConstructor = maskPrefix
+            ? ObjectKeyWithMaskedPrefix::new
+            : PlainObjectKey::new;
     }
 
     /**
@@ -93,7 +99,7 @@ public ObjectKey key(final RemoteLogSegmentMetadata remoteLogSegmentMetadata, fi
         Objects.requireNonNull(remoteLogSegmentMetadata, "remoteLogSegmentMetadata cannot be null");
         Objects.requireNonNull(suffix, "suffix cannot be null");
 
-        return new PlainObjectKey(prefix, mainPath(remoteLogSegmentMetadata) + "." + suffix.value);
+        return objectKeyConstructor.apply(prefix, mainPath(remoteLogSegmentMetadata) + "." + suffix.value);
     }
 
     /**
@@ -115,7 +121,7 @@ public ObjectKey key(final Map<Integer, Object> fields,
 
         final var prefix = (String) fields.getOrDefault(OBJECT_PREFIX.index(), this.prefix);
         final var main = (String) fields.getOrDefault(OBJECT_KEY.index(), mainPath(remoteLogSegmentMetadata));
-        return new PlainObjectKey(prefix, main + "." + suffix.value);
+        return objectKeyConstructor.apply(prefix, main + "." + suffix.value);
     }
 
     /**
@@ -160,8 +166,8 @@ private static String filenamePrefixFromOffset(final long offset) {
      * <p>Its string representation is identical to its value.
      */
     static class PlainObjectKey implements ObjectKey {
-        private final String prefix;
-        private final String mainPathAndSuffix;
+        protected final String prefix;
+        protected final String mainPathAndSuffix;
 
         PlainObjectKey(final String prefix, final String mainPathAndSuffix) {
             this.prefix = Objects.requireNonNull(prefix, "prefix cannot be null");
@@ -198,4 +204,20 @@ public String toString() {
             return value();
         }
     }
+
+    /**
+     * The object key that consists of a prefix and main part + suffix (as the parent class {@link PlainObjectKey}).
+     *
+     * <p>In its string representation, the prefix is masked with {@code <prefix>/}.
+     */
+    static class ObjectKeyWithMaskedPrefix extends PlainObjectKey {
+        ObjectKeyWithMaskedPrefix(final String prefix, final String mainPathAndSuffix) {
+            super(prefix, mainPathAndSuffix);
+        }
+
+        @Override
+        public String toString() {
+            return "<prefix>/" + mainPathAndSuffix;
+        }
+    }
 }
diff --git a/core/src/main/java/io/aiven/kafka/tieredstorage/RemoteStorageManager.java b/core/src/main/java/io/aiven/kafka/tieredstorage/RemoteStorageManager.java
@@ -139,7 +139,7 @@ public void configure(final Map<String, ?> configs) {
             .recordLevel(Sensor.RecordingLevel.forName(config.getString(METRICS_RECORDING_LEVEL_CONFIG)));
         metrics = new Metrics(time, metricConfig);
         setStorage(config.storage());
-        objectKeyFactory = new ObjectKeyFactory(config.keyPrefix());
+        objectKeyFactory = new ObjectKeyFactory(config.keyPrefix(), config.keyPrefixMask());
         encryptionEnabled = config.encryptionEnabled();
         if (encryptionEnabled) {
             final Map<String, KeyPair> keyRing = new HashMap<>();

diff --git a/core/src/main/java/io/aiven/kafka/tieredstorage/config/RemoteStorageManagerConfig.java b/core/src/main/java/io/aiven/kafka/tieredstorage/config/RemoteStorageManagerConfig.java
@@ -48,6 +48,9 @@ public class RemoteStorageManagerConfig extends AbstractConfig {
     private static final String OBJECT_KEY_PREFIX_CONFIG = "key.prefix";
     private static final String OBJECT_KEY_PREFIX_DOC = "The object storage path prefix";
 
+    private static final String OBJECT_KEY_PREFIX_MASK_CONFIG = "key.prefix.mask";
+    private static final String OBJECT_KEY_PREFIX_MASK_DOC = "Whether to mask path prefix in logs";
+
     private static final String SEGMENT_MANIFEST_CACHE_PREFIX = "segment.manifest.cache.";
     private static final String SEGMENT_MANIFEST_CACHE_SIZE_CONFIG = SEGMENT_MANIFEST_CACHE_PREFIX + "size";
     private static final Long SEGMENT_MANIFEST_CACHE_SIZE_DEFAULT = 1000L;  // TODO consider a better default
@@ -115,6 +118,14 @@ public class RemoteStorageManagerConfig extends AbstractConfig {
             OBJECT_KEY_PREFIX_DOC
         );
 
+        CONFIG.define(
+            OBJECT_KEY_PREFIX_MASK_CONFIG,
+            ConfigDef.Type.BOOLEAN,
+            false,
+            ConfigDef.Importance.LOW,
+            OBJECT_KEY_PREFIX_MASK_DOC
+        );
+
         CONFIG.define(
             SEGMENT_MANIFEST_CACHE_SIZE_CONFIG,
             ConfigDef.Type.LONG,
@@ -335,6 +346,10 @@ public String keyPrefix() {
         return getString(OBJECT_KEY_PREFIX_CONFIG);
     }
 
+    public boolean keyPrefixMask() {
+        return getBoolean(OBJECT_KEY_PREFIX_MASK_CONFIG);
+    }
+
     public int chunkSize() {
         return getInt(CHUNK_SIZE_CONFIG);
     }

diff --git a/core/src/test/java/io/aiven/kafka/tieredstorage/ObjectKeyEqualsTest.java b/core/src/test/java/io/aiven/kafka/tieredstorage/ObjectKeyEqualsTest.java
@@ -0,0 +1,58 @@
+/*
+ * Copyright 2023 Aiven Oy
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.aiven.kafka.tieredstorage;
+
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class ObjectKeyEqualsTest {
+    @ParameterizedTest
+    @ValueSource(classes = {ObjectKeyFactory.PlainObjectKey.class, ObjectKeyFactory.ObjectKeyWithMaskedPrefix.class})
+    void identical(final Class<ObjectKeyFactory.PlainObjectKey> keyClass) throws ReflectiveOperationException {
+        final var constructor = keyClass.getDeclaredConstructor(String.class, String.class);
+        final var k1 = constructor.newInstance("prefix", "mainPathAndSuffix");
+        final var k2 = constructor.newInstance("prefix", "mainPathAndSuffix");
+        assertThat(k1).isEqualTo(k2);
+        assertThat(k2).isEqualTo(k1);
+        assertThat(k1).hasSameHashCodeAs(k2);
+    }
+
+    @ParameterizedTest
+    @ValueSource(classes = {ObjectKeyFactory.PlainObjectKey.class, ObjectKeyFactory.ObjectKeyWithMaskedPrefix.class})
+    void differentPrefix(final Class<ObjectKeyFactory.PlainObjectKey> keyClass) throws ReflectiveOperationException {
+        final var constructor = keyClass.getDeclaredConstructor(String.class, String.class);
+        final var k1 = constructor.newInstance("prefix1", "mainPathAndSuffix");
+        final var k2 = constructor.newInstance("prefix2", "mainPathAndSuffix");
+        assertThat(k1).isNotEqualTo(k2);
+        assertThat(k2).isNotEqualTo(k1);
+        assertThat(k1).doesNotHaveSameHashCodeAs(k2);
+    }
+
+    @ParameterizedTest
+    @ValueSource(classes = {ObjectKeyFactory.PlainObjectKey.class, ObjectKeyFactory.ObjectKeyWithMaskedPrefix.class})
+    void differentMainPathAndSuffix(final Class<ObjectKeyFactory.PlainObjectKey> keyClass)
+        throws ReflectiveOperationException {
+        final var constructor = keyClass.getDeclaredConstructor(String.class, String.class);
+        final var k1 = constructor.newInstance("prefix", "mainPathAndSuffix1");
+        final var k2 = constructor.newInstance("prefix", "mainPathAndSuffix2");
+        assertThat(k1).isNotEqualTo(k2);
+        assertThat(k2).isNotEqualTo(k1);
+        assertThat(k1).doesNotHaveSameHashCodeAs(k2);
+    }
+}
diff --git a/core/src/test/java/io/aiven/kafka/tieredstorage/ObjectKeyFactoryTest.java b/core/src/test/java/io/aiven/kafka/tieredstorage/ObjectKeyFactoryTest.java
@@ -42,7 +42,7 @@ class ObjectKeyFactoryTest {
 
     @Test
     void test() {
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/", false);
         assertThat(objectKeyFactory.key(REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).value())
             .isEqualTo(
                 "prefix/topic-AAAAAAAAAAAAAAAAAAAAAQ/7/"
@@ -76,7 +76,7 @@ void test() {
 
     @Test
     void withCustomFieldsEmpty() {
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/", false);
         final Map<Integer, Object> fields = Map.of();
         assertThat(
             objectKeyFactory.key(fields, REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).value()
@@ -118,7 +118,7 @@ void withCustomFieldsEmpty() {
 
     @Test
     void withCustomFieldsOnlyPrefix() {
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/", false);
         final Map<Integer, Object> fields = Map.of(SegmentCustomMetadataField.OBJECT_PREFIX.index(), "other/");
         assertThat(
             objectKeyFactory.key(fields, REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).value()
@@ -160,7 +160,7 @@ void withCustomFieldsOnlyPrefix() {
 
     @Test
     void withCustomFieldsOnlyKey() {
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/", false);
         final Map<Integer, Object> fields = Map.of(SegmentCustomMetadataField.OBJECT_KEY.index(), "topic/7/file");
         assertThat(
             objectKeyFactory.key(fields, REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).value()
@@ -188,7 +188,7 @@ void withCustomFieldsOnlyKey() {
 
     @Test
     void withCustomFieldsAll() {
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/");
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/", false);
         final Map<Integer, Object> fields = Map.of(
             SegmentCustomMetadataField.OBJECT_PREFIX.index(), "other/",
             SegmentCustomMetadataField.OBJECT_KEY.index(), "topic/7/file");
@@ -218,7 +218,7 @@ void withCustomFieldsAll() {
 
     @Test
     void nullPrefix() {
-        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory(null);
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory(null, false);
         assertThat(objectKeyFactory.key(REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).value())
             .isEqualTo(
                 "topic-AAAAAAAAAAAAAAAAAAAAAQ/7/00000000000000001234-AAAAAAAAAAAAAAAAAAAAAA.log");
@@ -247,4 +247,49 @@ void suffixForIndexTypes() {
             .extracting("value")
             .isEqualTo("leader-epoch-checkpoint");
     }
+
+    @Test
+    void prefixMasking() {
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("prefix/", true);
+        assertThat(objectKeyFactory.key(REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).toString())
+            .isEqualTo("<prefix>/topic-AAAAAAAAAAAAAAAAAAAAAQ/7/00000000000000001234-AAAAAAAAAAAAAAAAAAAAAA.log");
+    }
+
+    @Test
+    void prefixMaskingWithCustomFieldsEmpty() {
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("real-prefix/", true);
+        final Map<Integer, Object> fields = Map.of();
+        assertThat(
+            objectKeyFactory.key(fields, REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).toString()
+        ).isEqualTo("<prefix>/topic-AAAAAAAAAAAAAAAAAAAAAQ/7/00000000000000001234-AAAAAAAAAAAAAAAAAAAAAA.log");
+    }
+
+    @Test
+    void prefixMaskingWithCustomFieldsOnlyPrefix() {
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("real-prefix/", true);
+        final Map<Integer, Object> fields = Map.of(SegmentCustomMetadataField.OBJECT_PREFIX.index(), "other/");
+        assertThat(
+            objectKeyFactory.key(fields, REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).toString()
+        ).isEqualTo("<prefix>/topic-AAAAAAAAAAAAAAAAAAAAAQ/7/00000000000000001234-AAAAAAAAAAAAAAAAAAAAAA.log");
+    }
+
+    @Test
+    void prefixMaskingWithCustomFieldsOnlyKey() {
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("real-prefix/", true);
+        final Map<Integer, Object> fields = Map.of(SegmentCustomMetadataField.OBJECT_KEY.index(), "topic/7/file");
+        assertThat(
+            objectKeyFactory.key(fields, REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).toString()
+        ).isEqualTo("<prefix>/topic/7/file.log");
+    }
+
+    @Test
+    void prefixMaskingWithCustomFieldsAll() {
+        final ObjectKeyFactory objectKeyFactory = new ObjectKeyFactory("real-prefix/", true);
+        final Map<Integer, Object> fields = Map.of(
+            SegmentCustomMetadataField.OBJECT_PREFIX.index(), "other/",
+            SegmentCustomMetadataField.OBJECT_KEY.index(), "topic/7/file");
+        assertThat(
+            objectKeyFactory.key(fields, REMOTE_LOG_SEGMENT_METADATA, ObjectKeyFactory.Suffix.LOG).toString()
+        ).isEqualTo("<prefix>/topic/7/file.log");
+    }
 }
diff --git a/core/src/test/java/io/aiven/kafka/tieredstorage/PlainObjectKeyEqualsTest.java b/core/src/test/java/io/aiven/kafka/tieredstorage/PlainObjectKeyEqualsTest.java
diff --git a/core/src/test/java/io/aiven/kafka/tieredstorage/config/RemoteStorageManagerConfigTest.java b/core/src/test/java/io/aiven/kafka/tieredstorage/config/RemoteStorageManagerConfigTest.java
@@ -51,6 +51,7 @@ void minimalConfig() {
         assertThat(config.encryptionKeyPairId()).isNull();
         assertThat(config.encryptionKeyRing()).isNull();
         assertThat(config.keyPrefix()).isEmpty();
+        assertThat(config.keyPrefixMask()).isFalse();
         assertThat(config.customMetadataKeysIncluded()).isEmpty();
     }
 
@@ -338,4 +339,16 @@ void invalidCustomMetadataFields() {
             .hasMessage("Invalid value unknown for configuration custom.metadata.fields.include: "
                 + "String must be one of: REMOTE_SIZE, OBJECT_PREFIX, OBJECT_KEY");
     }
+
+    @Test
+    void keyPrefixMasking() {
+        final var config = new RemoteStorageManagerConfig(
+            Map.of(
+                "storage.backend.class", NoopStorageBackend.class.getCanonicalName(),
+                "chunk.size", "123",
+                "key.prefix.mask", "true"
+            )
+        );
+        assertThat(config.keyPrefixMask()).isTrue();
+    }
 }