Add authorization support for Jetty Servlets

zacw7 · arhimondr · commit 35e7243e2d98 · 2020-07-08T23:53:24.000-04:00
diff --git a/http-server/src/main/java/com/facebook/airlift/http/server/AuthorizationEnabledServlet.java b/http-server/src/main/java/com/facebook/airlift/http/server/AuthorizationEnabledServlet.java
@@ -0,0 +1,141 @@
+/*
+ * Copyright 2010 Proofpoint, Inc.
+ *
+ * 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.airlift.http.server;
+
+import com.google.common.collect.ImmutableSet;
+
+import javax.annotation.security.RolesAllowed;
+import javax.servlet.Servlet;
+import javax.servlet.ServletException;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.security.Principal;
+import java.util.Optional;
+import java.util.Set;
+
+import static com.facebook.airlift.http.server.HttpServerConfig.AuthorizationPolicy;
+import static com.google.common.io.ByteStreams.copy;
+import static com.google.common.io.ByteStreams.nullOutputStream;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+
+public class AuthorizationEnabledServlet
+        extends HttpServlet
+{
+    private final Servlet delegate;
+    private final Authorizer authorizer;
+    private final AuthorizationPolicy authorizationPolicy;
+    private final Set<String> defaultAllowedRoles;
+    private final Optional<Set<String>> allowedRoles;
+
+    public AuthorizationEnabledServlet(
+            Servlet delegate,
+            Authorizer authorizer,
+            AuthorizationPolicy authorizationPolicy,
+            Set<String> defaultAllowedRoles)
+    {
+        this.delegate = requireNonNull(delegate, "delegate is null");
+        this.authorizer = requireNonNull(authorizer, "authorizer is null");
+        this.authorizationPolicy = requireNonNull(authorizationPolicy, "authorizationPolicy is null");
+        this.defaultAllowedRoles = requireNonNull(defaultAllowedRoles, "defaultAllowedRoles is null");
+        this.allowedRoles = getRolesFromClassMetadata(delegate);
+    }
+
+    @Override
+    public void init()
+            throws ServletException
+    {
+        super.init();
+        delegate.init(this.getServletConfig());
+    }
+
+    @Override
+    public void service(ServletRequest req, ServletResponse res)
+            throws ServletException, IOException
+    {
+        HttpServletRequest request = (HttpServletRequest) req;
+        HttpServletResponse response = (HttpServletResponse) res;
+
+        Principal principal = request.getUserPrincipal();
+        if (principal == null) {
+            abortWithMessage(request, response, "Request principal is missing.");
+            return;
+        }
+
+        Optional<Set<String>> allowedRoles = this.allowedRoles;
+        if (!allowedRoles.isPresent()) {
+            switch (authorizationPolicy) {
+                case ALLOW:
+                    delegate.service(req, res);
+                    return;
+                case DENY:
+                    abortWithMessage(request, response, format("Principal %s is not allowed to access the resource. Reason: denied by default policy",
+                            principal.getName()));
+                    return;
+                case DEFAULT_ROLES:
+                    allowedRoles = Optional.of(defaultAllowedRoles);
+                    break;
+                default:
+            }
+        }
+
+        AuthorizationResult result = authorizer.authorize(principal, allowedRoles.get());
+        if (!result.isAllowed()) {
+            abortWithMessage(request, response, format("Principal %s is not allowed to access the resource. Reason: %s",
+                    principal.getName(),
+                    result.getReason()));
+            return;
+        }
+        delegate.service(req, res);
+    }
+
+    private static void abortWithMessage(HttpServletRequest request, HttpServletResponse response, String message)
+            throws IOException
+    {
+        skipRequestBody(request);
+        response.sendError(HttpServletResponse.SC_FORBIDDEN, format(message));
+    }
+
+    private static void skipRequestBody(HttpServletRequest request)
+            throws IOException
+    {
+        // If we send the challenge without consuming the body of the request,
+        // the server will close the connection after sending the response.
+        // The client may interpret this as a failed request and not resend the
+        // request with the authentication header. We can avoid this behavior
+        // in the client by reading and discarding the entire body of the
+        // unauthenticated request before sending the response.
+        try (InputStream inputStream = request.getInputStream()) {
+            copy(inputStream, nullOutputStream());
+        }
+    }
+
+    private static Optional<Set<String>> getRolesFromClassMetadata(Servlet servlet)
+    {
+        if (servlet.getClass().isAnnotationPresent(RolesAllowed.class)) {
+            return Optional.of(ImmutableSet.copyOf(servlet.getClass().getAnnotation(RolesAllowed.class).value()));
+        }
+        else {
+            return Optional.empty();
+        }
+    }
+}
diff --git a/http-server/src/main/java/com/facebook/airlift/http/server/HttpServer.java b/http-server/src/main/java/com/facebook/airlift/http/server/HttpServer.java
@@ -83,6 +83,7 @@
 
 import static com.facebook.airlift.http.utils.jetty.ConcurrentScheduler.createConcurrentScheduler;
 import static com.google.common.base.MoreObjects.firstNonNull;
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
@@ -118,7 +119,8 @@ public HttpServer(HttpServerInfo httpServerInfo,
             LoginService loginService,
             TraceTokenManager tokenManager,
             RequestStats stats,
-            EventClient eventClient)
+            EventClient eventClient,
+            Authorizer authorizer)
             throws IOException
     {
         requireNonNull(httpServerInfo, "httpServerInfo is null");
@@ -368,7 +370,7 @@ public HttpServer(HttpServerInfo httpServerInfo,
             handlers.addHandler(gzipHandler);
         }
 
-        handlers.addHandler(createServletContext(defaultServlet, servlets, parameters, filters, tokenManager, loginService, "http", "https"));
+        handlers.addHandler(createServletContext(config, defaultServlet, servlets, parameters, filters, tokenManager, loginService, authorizer, "http", "https"));
 
         if (config.isRequestStatsEnabled()) {
             RequestLogHandler statsRecorder = new RequestLogHandler();
@@ -382,7 +384,7 @@ public HttpServer(HttpServerInfo httpServerInfo,
 
         HandlerList rootHandlers = new HandlerList();
         if (theAdminServlet != null && config.isAdminEnabled()) {
-            rootHandlers.addHandler(createServletContext(theAdminServlet, ImmutableMap.of(), adminParameters, adminFilters, tokenManager, loginService, "admin"));
+            rootHandlers.addHandler(createServletContext(config, theAdminServlet, ImmutableMap.of(), adminParameters, adminFilters, tokenManager, loginService, authorizer, "admin"));
         }
         rootHandlers.addHandler(statsHandler);
         server.setHandler(rootHandlers);
@@ -394,12 +396,14 @@ public HttpServer(HttpServerInfo httpServerInfo,
     }
 
     private static ServletContextHandler createServletContext(
+            HttpServerConfig config,
             Servlet defaultServlet,
             Map<String, Servlet> servlets,
             Map<String, String> parameters,
             Set<Filter> filters,
             TraceTokenManager tokenManager,
             LoginService loginService,
+            Authorizer authorizer,
             String... connectorNames)
     {
         ServletContextHandler context = new ServletContextHandler(ServletContextHandler.NO_SESSIONS);
@@ -427,9 +431,22 @@ private static ServletContextHandler createServletContext(
         context.addServlet(servletHolder, "/*");
 
         for (Map.Entry<String, Servlet> servlet : servlets.entrySet()) {
-            ServletHolder holder = new ServletHolder(servlet.getValue());
-            holder.setInitParameters(ImmutableMap.copyOf(parameters));
-            context.addServlet(holder, servlet.getKey());
+            if (config.isAuthorizationEnabled()) {
+                checkArgument(authorizer != null, "when authorization is enabled, authorizer implementation must be provided");
+                AuthorizationEnabledServlet authorizationEnabledServlet = new AuthorizationEnabledServlet(
+                        servlet.getValue(),
+                        authorizer,
+                        config.getDefaultAuthorizationPolicy(),
+                        config.getDefaultAllowedRoles());
+                ServletHolder holder = new ServletHolder(authorizationEnabledServlet);
+                holder.setInitParameters(ImmutableMap.copyOf(parameters));
+                context.addServlet(holder, servlet.getKey());
+            }
+            else {
+                ServletHolder holder = new ServletHolder(servlet.getValue());
+                holder.setInitParameters(ImmutableMap.copyOf(parameters));
+                context.addServlet(holder, servlet.getKey());
+            }
         }
 
         // Starting with Jetty 9 there is no way to specify connectors directly, but
diff --git a/http-server/src/main/java/com/facebook/airlift/http/server/HttpServerProvider.java b/http-server/src/main/java/com/facebook/airlift/http/server/HttpServerProvider.java
@@ -59,6 +59,7 @@ public class HttpServerProvider
     private final Set<Filter> adminFilters;
     private TraceTokenManager traceTokenManager;
     private final EventClient eventClient;
+    private Authorizer authorizer;
 
     @Inject
     public HttpServerProvider(HttpServerInfo httpServerInfo,
@@ -131,6 +132,12 @@ public void setTokenManager(@Nullable TraceTokenManager tokenManager)
         this.traceTokenManager = tokenManager;
     }
 
+    @Inject(optional = true)
+    public void setAuthorizer(@Nullable Authorizer authorizer)
+    {
+        this.authorizer = authorizer;
+    }
+
     @Override
     public HttpServer get()
     {
@@ -150,7 +157,8 @@ public HttpServer get()
                     loginService,
                     traceTokenManager,
                     stats,
-                    eventClient);
+                    eventClient,
+                    authorizer);
             httpServer.start();
             return httpServer;
         }
diff --git a/http-server/src/main/java/com/facebook/airlift/http/server/testing/TestingHttpServer.java b/http-server/src/main/java/com/facebook/airlift/http/server/testing/TestingHttpServer.java
@@ -16,6 +16,7 @@
 package com.facebook.airlift.http.server.testing;
 
 import com.facebook.airlift.event.client.NullEventClient;
+import com.facebook.airlift.http.server.Authorizer;
 import com.facebook.airlift.http.server.HttpServer;
 import com.facebook.airlift.http.server.HttpServerBinder.HttpResourceBinding;
 import com.facebook.airlift.http.server.HttpServerConfig;
@@ -25,14 +26,15 @@
 import com.facebook.airlift.node.NodeInfo;
 import com.facebook.airlift.tracetoken.TraceTokenManager;
 import com.google.common.collect.ImmutableSet;
+import com.google.inject.Inject;
 
-import javax.inject.Inject;
 import javax.servlet.Filter;
 import javax.servlet.Servlet;
 
 import java.io.IOException;
 import java.net.URI;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 public class TestingHttpServer
@@ -46,7 +48,8 @@ public TestingHttpServer(
             HttpServerConfig config,
             @TheServlet Servlet servlet,
             @TheServlet Map<String, Servlet> servlets,
-            @TheServlet Map<String, String> initParameters)
+            @TheServlet Map<String, String> initParameters,
+            Optional<Authorizer> authorizer)
             throws IOException
     {
         this(httpServerInfo,
@@ -56,7 +59,8 @@ public TestingHttpServer(
                 servlets,
                 initParameters,
                 ImmutableSet.of(),
-                ImmutableSet.of());
+                ImmutableSet.of(),
+                authorizer);
     }
 
     @Inject
@@ -68,7 +72,8 @@ public TestingHttpServer(
             @TheServlet Map<String, Servlet> servlets,
             @TheServlet Map<String, String> initParameters,
             @TheServlet Set<Filter> filters,
-            @TheServlet Set<HttpResourceBinding> resources)
+            @TheServlet Set<HttpResourceBinding> resources,
+            Optional<Authorizer> authorizer)
             throws IOException
     {
         super(httpServerInfo,
@@ -86,7 +91,8 @@ public TestingHttpServer(
                 null,
                 new TraceTokenManager(),
                 new RequestStats(),
-                new NullEventClient());
+                new NullEventClient(),
+                authorizer.orElse(null));
         this.httpServerInfo = httpServerInfo;
     }
 
diff --git a/http-server/src/main/java/com/facebook/airlift/http/server/testing/TestingHttpServerModule.java b/http-server/src/main/java/com/facebook/airlift/http/server/testing/TestingHttpServerModule.java
@@ -18,6 +18,7 @@
 import com.facebook.airlift.discovery.client.AnnouncementHttpServerInfo;
 import com.facebook.airlift.http.server.AuthenticationFilter;
 import com.facebook.airlift.http.server.Authenticator;
+import com.facebook.airlift.http.server.Authorizer;
 import com.facebook.airlift.http.server.HttpServer;
 import com.facebook.airlift.http.server.HttpServerConfig;
 import com.facebook.airlift.http.server.HttpServerInfo;
@@ -40,6 +41,7 @@
 import static com.facebook.airlift.http.server.HttpServerBinder.HttpResourceBinding;
 import static com.google.inject.multibindings.MapBinder.newMapBinder;
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
+import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 
 public class TestingHttpServerModule
         implements Module
@@ -79,11 +81,13 @@ public void configure(Binder binder)
         newSetBinder(binder, Filter.class, TheServlet.class).addBinding()
                 .to(AuthenticationFilter.class).in(Scopes.SINGLETON);
         newSetBinder(binder, Authenticator.class);
+        newOptionalBinder(binder, Authorizer.class);
     }
 
     @Provides
     List<Authenticator> getAuthenticatorList(Set<Authenticator> authenticators)
     {
         return ImmutableList.copyOf(authenticators);
+        newOptionalBinder(binder, Authorizer.class);
     }
 }
diff --git a/http-server/src/test/java/com/facebook/airlift/http/server/TestAuthorizationEnabledServletInHttpServer.java b/http-server/src/test/java/com/facebook/airlift/http/server/TestAuthorizationEnabledServletInHttpServer.java
diff --git a/http-server/src/test/java/com/facebook/airlift/http/server/testing/TestTestingHttpServer.java b/http-server/src/test/java/com/facebook/airlift/http/server/testing/TestTestingHttpServer.java
