auditing - initial changes

common/src/main/java/software/amazon/s3/analyticsaccelerator/request/StreamAuditContext.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * 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 software.amazon.s3.analyticsaccelerator.request;
+
+import lombok.AccessLevel;
+import lombok.Builder;
+import lombok.Getter;
+
+/**
+ * This is a class to contain any stream specific information, such as what is the higher level
+ * operation that this stream belongs to? Useful for the integration with S3A, it allows S3A to pass
+ * in the spanId and operation name with which a stream is opened. This is used by S3A's audit
+ * logging, which is of the form "3eb3c657-3e59-4f20-b484-e215c90c49f2-00000011 Executing op_open
+ * with {action_http_head_request....", where 3eb3c657-3e59-4f20-b484-e215c90c49f2-00000011 is the
+ * span_id for the open() operation, and ` op_open` is the operation name. See <a
+ * href="https://github.com/apache/hadoop/blob/trunk/hadoop-tools/hadoop-aws/src/site/markdown/tools/hadoop-aws/auditing.md">...</a>
+ * for more details.
+ */
+@Builder(access = AccessLevel.PUBLIC)
+@Getter
+public class StreamAuditContext {
+  private final String spanId;
+  private final String operationName;
+}

common/src/main/java/software/amazon/s3/analyticsaccelerator/util/OpenStreamInformation.java

@@ -19,7 +19,7 @@
 import lombok.Builder;
 import lombok.Getter;
 import software.amazon.s3.analyticsaccelerator.request.ObjectMetadata;
-import software.amazon.s3.analyticsaccelerator.request.StreamContext;
+import software.amazon.s3.analyticsaccelerator.request.StreamAuditContext;
 
 /**
  * Open stream information, useful for allowing the stream opening application to pass down known
@@ -38,7 +38,7 @@
 @Builder(access = AccessLevel.PUBLIC)
 @Getter
 public class OpenStreamInformation {
-  private final StreamContext streamContext;
+  private final StreamAuditContext streamAuditContext;
   private final ObjectMetadata objectMetadata;
   private final InputPolicy inputPolicy;
 

common/src/test/java/software/amazon/s3/analyticsaccelerator/util/OpenStreamInformationTest.java

@@ -20,7 +20,7 @@
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 import software.amazon.s3.analyticsaccelerator.request.ObjectMetadata;
-import software.amazon.s3.analyticsaccelerator.request.StreamContext;
+import software.amazon.s3.analyticsaccelerator.request.StreamAuditContext;
 
 public class OpenStreamInformationTest {
 
@@ -29,63 +29,63 @@ public void testDefaultInstance() {
     OpenStreamInformation info = OpenStreamInformation.DEFAULT;
 
     assertNotNull(info, "Default instance should not be null");
-    assertNull(info.getStreamContext(), "Default streamContext should be null");
+    assertNull(info.getStreamAuditContext(), "Default streamContext should be null");
     assertNull(info.getObjectMetadata(), "Default objectMetadata should be null");
     assertNull(info.getInputPolicy(), "Default inputPolicy should be null");
   }
 
   @Test
   public void testBuilderWithAllFields() {
-    StreamContext mockContext = Mockito.mock(StreamContext.class);
+    StreamAuditContext mockContext = Mockito.mock(StreamAuditContext.class);
     ObjectMetadata mockMetadata = Mockito.mock(ObjectMetadata.class);
     InputPolicy mockPolicy = Mockito.mock(InputPolicy.class);
 
     OpenStreamInformation info =
         OpenStreamInformation.builder()
-            .streamContext(mockContext)
+            .streamAuditContext(mockContext)
             .objectMetadata(mockMetadata)
             .inputPolicy(mockPolicy)
             .build();
 
-    assertSame(mockContext, info.getStreamContext(), "StreamContext should match");
+    assertSame(mockContext, info.getStreamAuditContext(), "StreamContext should match");
     assertSame(mockMetadata, info.getObjectMetadata(), "ObjectMetadata should match");
     assertSame(mockPolicy, info.getInputPolicy(), "InputPolicy should match");
   }
 
   @Test
   public void testBuilderWithPartialFields() {
-    StreamContext mockContext = Mockito.mock(StreamContext.class);
+    StreamAuditContext mockContext = Mockito.mock(StreamAuditContext.class);
     ObjectMetadata mockMetadata = Mockito.mock(ObjectMetadata.class);
 
     OpenStreamInformation info =
         OpenStreamInformation.builder()
-            .streamContext(mockContext)
+            .streamAuditContext(mockContext)
             .objectMetadata(mockMetadata)
             .build();
 
-    assertSame(mockContext, info.getStreamContext(), "StreamContext should match");
+    assertSame(mockContext, info.getStreamAuditContext(), "StreamContext should match");
     assertSame(mockMetadata, info.getObjectMetadata(), "ObjectMetadata should match");
     assertNull(info.getInputPolicy(), "InputPolicy should be null");
   }
 
   @Test
   public void testBuilderFieldRetention() {
     // Create mocks
-    StreamContext mockContext = Mockito.mock(StreamContext.class);
+    StreamAuditContext mockContext = Mockito.mock(StreamAuditContext.class);
     ObjectMetadata mockMetadata = Mockito.mock(ObjectMetadata.class);
     InputPolicy mockPolicy = Mockito.mock(InputPolicy.class);
 
     // Build object
     OpenStreamInformation info =
         OpenStreamInformation.builder()
-            .streamContext(mockContext)
+            .streamAuditContext(mockContext)
             .objectMetadata(mockMetadata)
             .inputPolicy(mockPolicy)
             .build();
 
     // Verify field retention
     assertNotNull(info, "Built object should not be null");
-    assertNotNull(info.getStreamContext(), "StreamContext should be retained");
+    assertNotNull(info.getStreamAuditContext(), "StreamContext should be retained");
     assertNotNull(info.getObjectMetadata(), "ObjectMetadata should be retained");
     assertNotNull(info.getInputPolicy(), "InputPolicy should be retained");
   }
@@ -94,12 +94,12 @@ public void testBuilderFieldRetention() {
   public void testNullFields() {
     OpenStreamInformation info =
         OpenStreamInformation.builder()
-            .streamContext(null)
+            .streamAuditContext(null)
             .objectMetadata(null)
             .inputPolicy(null)
             .build();
 
-    assertNull(info.getStreamContext(), "StreamContext should be null");
+    assertNull(info.getStreamAuditContext(), "StreamContext should be null");
     assertNull(info.getObjectMetadata(), "ObjectMetadata should be null");
     assertNull(info.getInputPolicy(), "InputPolicy should be null");
   }

object-client/src/main/java/software/amazon/s3/analyticsaccelerator/S3SdkObjectClient.java

@@ -15,6 +15,11 @@
  */
 package software.amazon.s3.analyticsaccelerator;
 
+import static software.amazon.s3.analyticsaccelerator.request.Constants.HEADER_REFERER;
+import static software.amazon.s3.analyticsaccelerator.request.Constants.HEADER_USER_AGENT;
+import static software.amazon.s3.analyticsaccelerator.request.Constants.OPERATION_NAME;
+import static software.amazon.s3.analyticsaccelerator.request.Constants.SPAN_ID;
+
 import java.io.UncheckedIOException;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -38,8 +43,6 @@
 
 /** Object client, based on AWS SDK v2 */
 public class S3SdkObjectClient implements ObjectClient {
-  private static final String HEADER_USER_AGENT = "User-Agent";
-  private static final String HEADER_REFERER = "Referer";
 
   @Getter @NonNull private final S3AsyncClient s3AsyncClient;
   @NonNull private final Telemetry telemetry;
@@ -117,11 +120,17 @@ public CompletableFuture<ObjectMetadata> headObject(
             .bucket(headRequest.getS3Uri().getBucket())
             .key(headRequest.getS3Uri().getKey());
 
-    // Add User-Agent header to the request.
-    builder.overrideConfiguration(
-        AwsRequestOverrideConfiguration.builder()
-            .putHeader(HEADER_USER_AGENT, this.userAgent.getUserAgent())
-            .build());
+    AwsRequestOverrideConfiguration.Builder requestOverrideConfigurationBuilder =
+        AwsRequestOverrideConfiguration.builder();
+
+    requestOverrideConfigurationBuilder.putHeader(HEADER_USER_AGENT, this.userAgent.getUserAgent());
+
+    if (openStreamInformation.getStreamAuditContext() != null) {
+      attachStreamContextToExecutionAttributes(
+          requestOverrideConfigurationBuilder, openStreamInformation.getStreamAuditContext());
+    }
+
+    builder.overrideConfiguration(requestOverrideConfigurationBuilder.build());
 
     return this.telemetry
         .measureCritical(
@@ -154,19 +163,17 @@ public CompletableFuture<ObjectContent> getObject(
     final String range = getRequest.getRange().toHttpString();
     builder.range(range);
 
-    final String referrerHeader;
-    if (openStreamInformation.getStreamContext() != null) {
-      referrerHeader =
-          openStreamInformation.getStreamContext().modifyAndBuildReferrerHeader(getRequest);
-    } else {
-      referrerHeader = getRequest.getReferrer().toString();
+    AwsRequestOverrideConfiguration.Builder requestOverrideConfigurationBuilder =
+        AwsRequestOverrideConfiguration.builder()
+            .putHeader(HEADER_REFERER, getRequest.getReferrer().toString())
+            .putHeader(HEADER_USER_AGENT, this.userAgent.getUserAgent());
+
+    if (openStreamInformation.getStreamAuditContext() != null) {
+      attachStreamContextToExecutionAttributes(
+          requestOverrideConfigurationBuilder, openStreamInformation.getStreamAuditContext());
     }
 
-    builder.overrideConfiguration(
-        AwsRequestOverrideConfiguration.builder()
-            .putHeader(HEADER_REFERER, referrerHeader)
-            .putHeader(HEADER_USER_AGENT, this.userAgent.getUserAgent())
-            .build());
+    builder.overrideConfiguration(requestOverrideConfigurationBuilder.build());
 
     return this.telemetry.measureCritical(
         () ->
@@ -195,4 +202,12 @@ private <T> Function<Throwable, T> handleException(S3URI s3Uri) {
       throw new UncheckedIOException(ExceptionHandler.toIOException(cause, s3Uri));
     };
   }
+
+  private void attachStreamContextToExecutionAttributes(
+      AwsRequestOverrideConfiguration.Builder requestOverrideConfigurationBuilder,
+      StreamAuditContext streamAuditContext) {
+    requestOverrideConfigurationBuilder
+        .putExecutionAttribute(SPAN_ID, streamAuditContext.getSpanId())
+        .putExecutionAttribute(OPERATION_NAME, streamAuditContext.getOperationName());
+  }
 }

object-client/src/main/java/software/amazon/s3/analyticsaccelerator/request/Constants.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * 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 software.amazon.s3.analyticsaccelerator.request;
+
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
+
+/** Class for request related constants. */
+public final class Constants {
+  private Constants() {}
+
+  public static final String HEADER_USER_AGENT = "User-Agent";
+  public static final String HEADER_REFERER = "Referer";
+
+  // These execution attributes are specific to the hadoop S3A integration, and are required for
+  // S3A's auditing feature. These execution attributes can be set per request, which are then
+  // picked up
+  // by execution interceptors in S3A. See S3A's LoggingAuditor for usage.
+  public static final ExecutionAttribute<String> SPAN_ID = new ExecutionAttribute<>("span");
+  public static final ExecutionAttribute<String> OPERATION_NAME =
+      new ExecutionAttribute<>("operation");
+}