Introducing ClientRequestFilter.java: A New Plugin for Applying Request Headers in the Authentication Filter Class

presto-main/pom.xml

@@ -499,6 +499,11 @@
             <artifactId>ratis-common</artifactId>
             <optional>true</optional>
         </dependency>
+        <dependency>
+            <groupId>com.squareup.okhttp3</groupId>
+            <artifactId>mockwebserver</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

presto-main/src/main/java/com/facebook/presto/ClientRequestFilterManager.java

@@ -0,0 +1,35 @@
+/*
+ * 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 com.facebook.presto;
+
+import com.facebook.presto.spi.ClientRequestFilter;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+public class ClientRequestFilterManager
+{
+    private final List<ClientRequestFilter> clientRequestFilters = new CopyOnWriteArrayList<>();
+
+    public List<ClientRequestFilter> getClientRequestFilters()
+    {
+        return ImmutableList.copyOf(clientRequestFilters);
+    }
+
+    public void registerClientRequestFilter(ClientRequestFilter clientRequestFilter)
+    {
+        clientRequestFilters.add(clientRequestFilter);
+    }
+}

presto-main/src/main/java/com/facebook/presto/ClientRequestFilterModule.java

@@ -0,0 +1,28 @@
+/*
+ * 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 com.facebook.presto;
+
+import com.google.inject.Binder;
+import com.google.inject.Module;
+import com.google.inject.Scopes;
+
+public class ClientRequestFilterModule
+        implements Module
+{
+    @Override
+    public void configure(Binder binder)
+    {
+        binder.bind(ClientRequestFilterManager.class).in(Scopes.SINGLETON);
+    }
+}

presto-main/src/main/java/com/facebook/presto/server/PluginManager.java

@@ -15,6 +15,7 @@
 
 import com.facebook.airlift.log.Logger;
 import com.facebook.airlift.node.NodeInfo;
+import com.facebook.presto.ClientRequestFilterManager;
 import com.facebook.presto.common.block.BlockEncoding;
 import com.facebook.presto.common.block.BlockEncodingManager;
 import com.facebook.presto.common.type.ParametricType;
@@ -27,6 +28,7 @@
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.security.AccessControlManager;
 import com.facebook.presto.server.security.PasswordAuthenticatorManager;
+import com.facebook.presto.spi.ClientRequestFilter;
 import com.facebook.presto.spi.CoordinatorPlugin;
 import com.facebook.presto.spi.Plugin;
 import com.facebook.presto.spi.analyzer.AnalyzerProvider;
@@ -131,6 +133,7 @@ public class PluginManager
     private final AnalyzerProviderManager analyzerProviderManager;
     private final QueryPreparerProviderManager queryPreparerProviderManager;
     private final NodeStatusNotificationManager nodeStatusNotificationManager;
+    private final ClientRequestFilterManager clientRequestFilterManager;
 
     @Inject
     public PluginManager(
@@ -152,7 +155,8 @@ public PluginManager(
             ClusterTtlProviderManager clusterTtlProviderManager,
             HistoryBasedPlanStatisticsManager historyBasedPlanStatisticsManager,
             TracerProviderManager tracerProviderManager,
-            NodeStatusNotificationManager nodeStatusNotificationManager)
+            NodeStatusNotificationManager nodeStatusNotificationManager,
+            ClientRequestFilterManager clientRequestFilterManager)
     {
         requireNonNull(nodeInfo, "nodeInfo is null");
         requireNonNull(config, "config is null");
@@ -184,6 +188,7 @@ public PluginManager(
         this.analyzerProviderManager = requireNonNull(analyzerProviderManager, "analyzerProviderManager is null");
         this.queryPreparerProviderManager = requireNonNull(queryPreparerProviderManager, "queryPreparerProviderManager is null");
         this.nodeStatusNotificationManager = requireNonNull(nodeStatusNotificationManager, "nodeStatusNotificationManager is null");
+        this.clientRequestFilterManager = requireNonNull(clientRequestFilterManager, "clientRequestFilterManager is null");
     }
 
     public void loadPlugins()
@@ -348,6 +353,11 @@ public void installPlugin(Plugin plugin)
             log.info("Registering node status notification provider %s", nodeStatusNotificationProviderFactory.getName());
             nodeStatusNotificationManager.addNodeStatusNotificationProviderFactory(nodeStatusNotificationProviderFactory);
         }
+
+        for (ClientRequestFilter clientRequestFilter : plugin.getClientRequestFilters()) {
+            log.info("Registering client request filter");
+            clientRequestFilterManager.registerClientRequestFilter(clientRequestFilter);
+        }
     }
 
     public void installCoordinatorPlugin(CoordinatorPlugin plugin)

presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java

@@ -30,6 +30,7 @@
 import com.facebook.airlift.tracetoken.TraceTokenModule;
 import com.facebook.drift.server.DriftServer;
 import com.facebook.drift.transport.netty.server.DriftNettyServerTransport;
+import com.facebook.presto.ClientRequestFilterModule;
 import com.facebook.presto.dispatcher.QueryPrerequisitesManager;
 import com.facebook.presto.dispatcher.QueryPrerequisitesManagerModule;
 import com.facebook.presto.eventlistener.EventListenerManager;
@@ -133,7 +134,8 @@ public void run()
                 new TempStorageModule(),
                 new QueryPrerequisitesManagerModule(),
                 new NodeTtlFetcherManagerModule(),
-                new ClusterTtlProviderManagerModule());
+                new ClusterTtlProviderManagerModule(),
+                new ClientRequestFilterModule());
 
         modules.addAll(getAdditionalModules());
 

presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java

@@ -15,9 +15,13 @@
 
 import com.facebook.airlift.http.server.AuthenticationException;
 import com.facebook.airlift.http.server.Authenticator;
+import com.facebook.presto.ClientRequestFilterManager;
+import com.facebook.presto.spi.ClientRequestFilter;
+import com.facebook.presto.spi.PrestoException;
 import com.google.common.base.Joiner;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.net.HttpHeaders;
 
 import javax.inject.Inject;
@@ -35,10 +39,18 @@
 import java.io.InputStream;
 import java.io.PrintWriter;
 import java.security.Principal;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
+import static com.facebook.presto.spi.StandardErrorCode.HEADER_MODIFICATION_ATTEMPT;
 import static com.google.common.io.ByteStreams.copy;
 import static com.google.common.io.ByteStreams.nullOutputStream;
 import static com.google.common.net.HttpHeaders.WWW_AUTHENTICATE;
@@ -52,12 +64,15 @@ public class AuthenticationFilter
     private static final String HTTPS_PROTOCOL = "https";
     private final List<Authenticator> authenticators;
     private final boolean allowForwardedHttps;
+    private final ClientRequestFilterManager clientRequestFilterManager;
+    private final List<String> headersBlockList = ImmutableList.of("X-Presto-Transaction-Id", "X-Presto-Started-Transaction-Id", "X-Presto-Clear-Transaction-Id", "X-Presto-Trace-Token");
 
     @Inject
-    public AuthenticationFilter(List<Authenticator> authenticators, SecurityConfig securityConfig)
+    public AuthenticationFilter(List<Authenticator> authenticators, SecurityConfig securityConfig, ClientRequestFilterManager clientRequestFilterManager)
     {
         this.authenticators = ImmutableList.copyOf(requireNonNull(authenticators, "authenticators is null"));
         this.allowForwardedHttps = requireNonNull(securityConfig, "securityConfig is null").getAllowForwardedHttps();
+        this.clientRequestFilterManager = requireNonNull(clientRequestFilterManager, "clientRequestFilterManager is null");
     }
 
     @Override
@@ -95,9 +110,37 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo
                 e.getAuthenticateHeader().ifPresent(authenticateHeaders::add);
                 continue;
             }
-
             // authentication succeeded
-            nextFilter.doFilter(withPrincipal(request, principal), response);
+            CustomHttpServletRequestWrapper wrappedRequest = new CustomHttpServletRequestWrapper(request);
+            Map<String, String> extraHeadersMap = new HashMap<>();
+            Set<String> globallyAddedHeaders = new HashSet<>();
+
+            for (ClientRequestFilter requestFilter : clientRequestFilterManager.getClientRequestFilters()) {
+                boolean headersPresent = requestFilter.getHeaderNames().stream()
+                        .allMatch(headerName -> request.getHeader(headerName) != null);
+
+                if (!headersPresent) {
+                    Optional<Map<String, String>> extraHeaderValueMap = requestFilter.getExtraHeaders(principal);
+
+                    extraHeaderValueMap.ifPresent(map -> {
+                        for (Map.Entry<String, String> extraHeaderEntry : map.entrySet()) {
+                            String headerKey = extraHeaderEntry.getKey();
+                            if (headersBlockList.contains(headerKey)) {
+                                throw new PrestoException(HEADER_MODIFICATION_ATTEMPT, "Modification attempt detected: The header " + headerKey + " is present in the blocked headers list.");
+                            }
+                            if (globallyAddedHeaders.contains(headerKey)) {
+                                throw new RuntimeException("Header conflict detected: " + headerKey + " already added by another filter.");
+                            }
+                            if (request.getHeader(headerKey) == null && requestFilter.getHeaderNames().contains(headerKey)) {
+                                extraHeadersMap.putIfAbsent(headerKey, extraHeaderEntry.getValue());
+                                globallyAddedHeaders.add(headerKey);
+                            }
+                        }
+                    });
+                }
+            }
+            wrappedRequest.setHeaders(extraHeadersMap);
+            nextFilter.doFilter(withPrincipal(wrappedRequest, principal), response);
             return;
         }
 
@@ -140,7 +183,7 @@ private boolean doesRequestSupportAuthentication(HttpServletRequest request)
         return false;
     }
 
-    private static ServletRequest withPrincipal(HttpServletRequest request, Principal principal)
+    public static ServletRequest withPrincipal(HttpServletRequest request, Principal principal)
     {
         requireNonNull(principal, "principal is null");
         return new HttpServletRequestWrapper(request)
@@ -166,4 +209,51 @@ private static void skipRequestBody(HttpServletRequest request)
             copy(inputStream, nullOutputStream());
         }
     }
+
+    public static class CustomHttpServletRequestWrapper
+            extends HttpServletRequestWrapper
+    {
+        private final Map<String, String> customHeaders = new ConcurrentHashMap<>();
+
+        public CustomHttpServletRequestWrapper(HttpServletRequest request)
+        {
+            super(request);
+        }
+
+        @Override
+        public String getHeader(String name)
+        {
+            String headerValue = customHeaders.get(name);
+            if (headerValue != null) {
+                return headerValue;
+            }
+            return super.getHeader(name);
+        }
+
+        @Override
+        public Enumeration<String> getHeaderNames()
+        {
+            ImmutableSet.Builder<String> headerNamesBuilder = ImmutableSet.builder();
+            headerNamesBuilder.addAll(customHeaders.keySet());
+            Enumeration<String> originalHeaderNames = super.getHeaderNames();
+            while (originalHeaderNames.hasMoreElements()) {
+                headerNamesBuilder.add(originalHeaderNames.nextElement());
+            }
+            return Collections.enumeration(headerNamesBuilder.build());
+        }
+
+        @Override
+        public Enumeration<String> getHeaders(String name)
+        {
+            if (customHeaders.containsKey(name)) {
+                return Collections.enumeration(Collections.singleton(customHeaders.get(name)));
+            }
+            return super.getHeaders(name);
+        }
+
+        public void setHeaders(Map<String, String> headers)
+        {
+            this.customHeaders.putAll(headers);
+        }
+    }
 }

presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java

@@ -32,6 +32,8 @@
 import com.facebook.airlift.tracetoken.TraceTokenModule;
 import com.facebook.drift.server.DriftServer;
 import com.facebook.drift.transport.netty.server.DriftNettyServerTransport;
+import com.facebook.presto.ClientRequestFilterManager;
+import com.facebook.presto.ClientRequestFilterModule;
 import com.facebook.presto.connector.ConnectorManager;
 import com.facebook.presto.cost.StatsCalculator;
 import com.facebook.presto.dispatcher.DispatchManager;
@@ -59,6 +61,7 @@
 import com.facebook.presto.server.ServerMainModule;
 import com.facebook.presto.server.ShutdownAction;
 import com.facebook.presto.server.security.ServerSecurityModule;
+import com.facebook.presto.spi.ClientRequestFilter;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.CoordinatorPlugin;
 import com.facebook.presto.spi.Plugin;
@@ -306,6 +309,7 @@ public TestingPrestoServer(
                 .add(new QueryPrerequisitesManagerModule())
                 .add(new NodeTtlFetcherManagerModule())
                 .add(new ClusterTtlProviderManagerModule())
+                .add(new ClientRequestFilterModule())
                 .add(binder -> {
                     binder.bind(TestingAccessControlManager.class).in(Scopes.SINGLETON);
                     binder.bind(TestingEventListenerManager.class).in(Scopes.SINGLETON);
@@ -839,4 +843,11 @@ private static int driftServerPort(DriftServer server)
     {
         return ((DriftNettyServerTransport) server.getServerTransport()).getPort();
     }
+
+    public static ClientRequestFilterManager getClientRequestFilterManager(List<ClientRequestFilter> requestFilters)
+    {
+        ClientRequestFilterManager manager = new ClientRequestFilterManager();
+        requestFilters.forEach(manager::registerClientRequestFilter);
+        return manager;
+    }
 }

presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.testing;
 
 import com.facebook.airlift.node.NodeInfo;
+import com.facebook.presto.ClientRequestFilterManager;
 import com.facebook.presto.GroupByHashPageIndexerFactory;
 import com.facebook.presto.PagesIndexPageSorter;
 import com.facebook.presto.Session;
@@ -515,7 +516,8 @@ private LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig,
                 new ThrowingClusterTtlProviderManager(),
                 historyBasedPlanStatisticsManager,
                 new TracerProviderManager(new TracingConfig()),
-                new NodeStatusNotificationManager());
+                new NodeStatusNotificationManager(),
+                new ClientRequestFilterManager());
 
         connectorManager.addConnectorFactory(globalSystemConnectorFactory);
         connectorManager.createConnection(GlobalSystemConnector.NAME, GlobalSystemConnector.NAME, ImmutableMap.of());