Improve the design of SecurityResponse

With the latest Java, we can have a more modern definition for `SecurityResponse` as a sealed hierarchy with explicit Success and Failure types. This makes it easier to understand and use properly.
trinodb · Jul 24, 2024 · ee7dcb0 · ee7dcb0
1 parent c3245af
commit ee7dcb0
Show file tree

Hide file tree

Showing 9 changed files with 71 additions and 21 deletions.
diff --git a/pom.xml b/pom.xml
@@ -201,4 +201,16 @@
             </dependency>
         </dependencies>
     </dependencyManagement>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-checkstyle-plugin</artifactId>
+                <configuration>
+                    <violationIgnore>InterfaceIsType</violationIgnore>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
 </project>
diff --git a/...aws-proxy-spi/src/main/java/io/trino/aws/proxy/spi/security/S3DatabaseSecurityFacade.java b/...aws-proxy-spi/src/main/java/io/trino/aws/proxy/spi/security/S3DatabaseSecurityFacade.java
@@ -15,6 +15,8 @@
 
 import java.util.Optional;
 
+import static io.trino.aws.proxy.spi.security.SecurityResponse.SUCCESS;
+
 public interface S3DatabaseSecurityFacade
 {
     S3DatabaseSecurityFacade DEFAULT = new S3DatabaseSecurityFacade() {};
@@ -26,11 +28,11 @@ default Optional<String> tableName(Optional<String> lowercaseAction)
 
     default SecurityResponse nonTableOperation(Optional<String> lowercaseAction)
     {
-        return SecurityResponse.DEFAULT;
+        return SUCCESS;
     }
 
     default SecurityResponse tableOperation(String tableName, Optional<String> lowercaseAction)
     {
-        return SecurityResponse.DEFAULT;
+        return SUCCESS;
     }
 }
diff --git a/trino-aws-proxy-spi/src/main/java/io/trino/aws/proxy/spi/security/SecurityResponse.java b/trino-aws-proxy-spi/src/main/java/io/trino/aws/proxy/spi/security/SecurityResponse.java
@@ -17,12 +17,32 @@
 
 import static java.util.Objects.requireNonNull;
 
-public record SecurityResponse(boolean canProceed, Optional<String> error)
+public sealed interface SecurityResponse
 {
-    public static final SecurityResponse DEFAULT = new SecurityResponse(true, Optional.empty());
+    SecurityResponse SUCCESS = new Success();
+    SecurityResponse FAILURE = new Failure();
 
-    public SecurityResponse
+    record Success()
+            implements SecurityResponse
     {
-        requireNonNull(error, "error is null");
+    }
+
+    record Failure(Optional<String> error)
+            implements SecurityResponse
+    {
+        public Failure
+        {
+            requireNonNull(error, "error is null");
+        }
+
+        public Failure()
+        {
+            this(Optional.empty());
+        }
+
+        public Failure(String error)
+        {
+            this(Optional.of(error));
+        }
     }
 }
diff --git a/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/TrinoAwsProxyServerModule.java b/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/TrinoAwsProxyServerModule.java
@@ -42,7 +42,6 @@
 import io.trino.aws.proxy.spi.credentials.CredentialsProvider;
 import io.trino.aws.proxy.spi.security.S3DatabaseSecurityFacadeProvider;
 import io.trino.aws.proxy.spi.security.S3SecurityFacadeProvider;
-import io.trino.aws.proxy.spi.security.SecurityResponse;
 import io.trino.aws.proxy.spi.signing.SigningController;
 import io.trino.aws.proxy.spi.signing.SigningServiceType;
 import org.glassfish.jersey.server.model.Resource;
@@ -56,6 +55,7 @@
 import static io.airlift.http.client.HttpClientBinder.httpClientBinder;
 import static io.airlift.http.server.HttpServerBinder.httpServerBinder;
 import static io.airlift.jaxrs.JaxrsBinder.jaxrsBinder;
+import static io.trino.aws.proxy.spi.security.SecurityResponse.SUCCESS;
 
 public class TrinoAwsProxyServerModule
         extends AbstractConfigurationAwareModule
@@ -114,7 +114,7 @@ public XmlMapper newXmlMapper()
     protected void moduleSpecificBinding(Binder binder)
     {
         binder.bind(S3SecurityController.class).in(Scopes.SINGLETON);
-        newOptionalBinder(binder, S3SecurityFacadeProvider.class).setDefault().toInstance(_ -> _ -> SecurityResponse.DEFAULT);
+        newOptionalBinder(binder, S3SecurityFacadeProvider.class).setDefault().toInstance(_ -> _ -> SUCCESS);
         binder.bind(RemoteS3Facade.class).to(VirtualHostStyleRemoteS3Facade.class).in(Scopes.SINGLETON);
     }
 

diff --git a/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/TrinoS3ProxyClient.java b/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/rest/TrinoS3ProxyClient.java
@@ -24,6 +24,7 @@
 import io.trino.aws.proxy.spi.rest.ParsedS3Request;
 import io.trino.aws.proxy.spi.rest.RequestContent;
 import io.trino.aws.proxy.spi.security.SecurityResponse;
+import io.trino.aws.proxy.spi.security.SecurityResponse.Failure;
 import io.trino.aws.proxy.spi.signing.SigningContext;
 import io.trino.aws.proxy.spi.signing.SigningController;
 import io.trino.aws.proxy.spi.signing.SigningMetadata;
@@ -93,11 +94,11 @@ public void proxyRequest(SigningMetadata signingMetadata, ParsedS3Request reques
         URI remoteUri = remoteS3Facade.buildEndpoint(uriBuilder(request.queryParameters()), request.rawPath(), request.bucketName(), request.requestAuthorization().region());
 
         SecurityResponse securityResponse = s3SecurityController.apply(request);
-        if (!securityResponse.canProceed()) {
+        if (securityResponse instanceof Failure(var error)) {
             log.debug("SecurityController check failed. AccessKey: %s, Request: %s, SecurityResponse: %s", signingMetadata.credentials().emulated().accessKey(), request, securityResponse);
             requestLoggingSession.logError("request.security.fail.credentials", signingMetadata.credentials().emulated());
             requestLoggingSession.logError("request.security.fail.request", request);
-            requestLoggingSession.logError("request.security.fail.response", securityResponse);
+            requestLoggingSession.logError("request.security.fail.error", error);
 
             throw new WebApplicationException(Response.Status.UNAUTHORIZED);
         }

diff --git a/...-proxy/src/main/java/io/trino/aws/proxy/server/security/S3DatabaseSecurityController.java b/...-proxy/src/main/java/io/trino/aws/proxy/server/security/S3DatabaseSecurityController.java
@@ -22,6 +22,8 @@
 import io.trino.aws.proxy.spi.security.S3SecurityFacade;
 import io.trino.aws.proxy.spi.security.S3SecurityFacadeProvider;
 import io.trino.aws.proxy.spi.security.SecurityResponse;
+import io.trino.aws.proxy.spi.security.SecurityResponse.Failure;
+import io.trino.aws.proxy.spi.security.SecurityResponse.Success;
 import jakarta.ws.rs.WebApplicationException;
 
 import java.util.Optional;
@@ -64,18 +66,30 @@ private SecurityResponse applyToRequest(S3DatabaseSecurityFacade securityFacade,
                     SecurityResponse securityResponse = securityFacade.tableOperation(tableName, lowercaseAction);
 
                     requestLoggingSession.logProperty("response.database.table-operation.table-name", tableName);
-                    requestLoggingSession.logProperty("response.database.table-operation.can-proceed", securityResponse.canProceed());
-                    requestLoggingSession.logProperty("response.database.table-operation.error", securityResponse.error());
+                    addLogging(requestLoggingSession, "response.database.table-operation.", securityResponse);
 
                     return securityResponse;
                 })
                 .orElseGet(() -> {
                     SecurityResponse securityResponse = securityFacade.nonTableOperation(lowercaseAction);
 
-                    requestLoggingSession.logProperty("response.database.non-table-operation.can-proceed", securityResponse.canProceed());
-                    requestLoggingSession.logProperty("response.database.non-table-operation.error", securityResponse.error());
+                    addLogging(requestLoggingSession, "response.database.non-table-operation.", securityResponse);
 
                     return securityResponse;
                 });
     }
+
+    private void addLogging(RequestLoggingSession requestLoggingSession, String prefix, SecurityResponse securityResponse)
+    {
+        switch (securityResponse) {
+            case Success _ -> addLogging(requestLoggingSession, prefix, true, Optional.empty());
+            case Failure(var error) -> addLogging(requestLoggingSession, prefix, false, error);
+        }
+    }
+
+    private void addLogging(RequestLoggingSession requestLoggingSession, String prefix, boolean success, Optional<String> error)
+    {
+        requestLoggingSession.logProperty(prefix + "success", success);
+        requestLoggingSession.logProperty(prefix + "error", error);
+    }
 }
diff --git a/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/security/S3SecurityController.java b/trino-aws-proxy/src/main/java/io/trino/aws/proxy/server/security/S3SecurityController.java
@@ -27,11 +27,12 @@
 import java.util.Optional;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static io.trino.aws.proxy.spi.security.SecurityResponse.SUCCESS;
 import static java.util.Objects.requireNonNull;
 
 public class S3SecurityController
 {
-    private static final S3SecurityFacadeProvider DEFAULT_SECURITY_FACADE_PROVIDER = _ -> _ -> SecurityResponse.DEFAULT;
+    private static final S3SecurityFacadeProvider DEFAULT_SECURITY_FACADE_PROVIDER = _ -> _ -> SUCCESS;
 
     private final S3SecurityFacadeProvider s3SecurityFacadeProvider;
 

diff --git a/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/TestDatabaseSecurity.java b/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/TestDatabaseSecurity.java
@@ -34,6 +34,7 @@
 import static io.trino.aws.proxy.server.testing.containers.DockerAttachUtil.clearInputStreamAndClose;
 import static io.trino.aws.proxy.server.testing.containers.DockerAttachUtil.inputToContainerStdin;
 import static io.trino.aws.proxy.spi.TrinoAwsProxyBinder.trinoAwsProxyBinder;
+import static io.trino.aws.proxy.spi.security.SecurityResponse.FAILURE;
 import static java.util.Objects.requireNonNull;
 
 @TrinoAwsProxyTest(filters = TestDatabaseSecurity.Filter.class)
@@ -76,7 +77,7 @@ public Optional<String> tableName(Optional<String> lowercaseAction)
                 public SecurityResponse tableOperation(String tableName, Optional<String> lowercaseAction)
                 {
                     if (disallowGets.get() && request.httpVerb().equalsIgnoreCase("GET")) {
-                        return new SecurityResponse(false, Optional.empty());
+                        return FAILURE;
                     }
                     return S3DatabaseSecurityFacade.super.tableOperation(tableName, lowercaseAction);
                 }

diff --git a/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/TestS3SecurityController.java b/trino-aws-proxy/src/test/java/io/trino/aws/proxy/server/TestS3SecurityController.java
@@ -17,7 +17,6 @@
 import io.trino.aws.proxy.server.testing.TestingS3SecurityController;
 import io.trino.aws.proxy.server.testing.harness.TrinoAwsProxyTest;
 import io.trino.aws.proxy.server.testing.harness.TrinoAwsProxyTestCommonModules.WithConfiguredBuckets;
-import io.trino.aws.proxy.spi.security.SecurityResponse;
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
@@ -27,8 +26,8 @@
 import software.amazon.awssdk.services.s3.model.ListObjectsResponse;
 import software.amazon.awssdk.services.s3.model.S3Exception;
 
-import java.util.Optional;
-
+import static io.trino.aws.proxy.spi.security.SecurityResponse.FAILURE;
+import static io.trino.aws.proxy.spi.security.SecurityResponse.SUCCESS;
 import static java.util.Objects.requireNonNull;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -66,9 +65,9 @@ public void testDisableListObjects()
         // set facade that disallows list objects on bucket "one"
         securityController.setDelegate(request -> _ -> {
             if ("one".equals(request.bucketName()) && request.httpVerb().equalsIgnoreCase("get")) {
-                return new SecurityResponse(false, Optional.empty());
+                return FAILURE;
             }
-            return SecurityResponse.DEFAULT;
+            return SUCCESS;
         });
 
         // list objects on "one" now disallowed