1 files changed, 134 insertions, 0 deletions
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/FilteringRequestHandler.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/FilteringRequestHandler.java
new file mode 100644
index 00000000000..de768f979a1
--- /dev/null
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/FilteringRequestHandler.java
@@ -0,0 +1,134 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.jdisc.http.server.jetty;
+
+import com.google.common.base.Preconditions;
+import com.yahoo.jdisc.Request;
+import com.yahoo.jdisc.Response;
+import com.yahoo.jdisc.handler.AbstractRequestHandler;
+import com.yahoo.jdisc.handler.BindingNotFoundException;
+import com.yahoo.jdisc.handler.CompletionHandler;
+import com.yahoo.jdisc.handler.ContentChannel;
+import com.yahoo.jdisc.handler.RequestDeniedException;
+import com.yahoo.jdisc.handler.RequestHandler;
+import com.yahoo.jdisc.handler.ResponseHandler;
+import com.yahoo.jdisc.http.HttpRequest;
+import com.yahoo.jdisc.http.filter.RequestFilter;
+import com.yahoo.jdisc.http.filter.ResponseFilter;
+
+import javax.servlet.http.HttpServletRequest;
+import java.nio.ByteBuffer;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Request handler that invokes request and response filters in addition to the bound request handler.
+ *
+ * @author Øyvind Bakksjø
+ */
+class FilteringRequestHandler extends AbstractRequestHandler {
+
+    private static final ContentChannel COMPLETING_CONTENT_CHANNEL = new ContentChannel() {
+
+        @Override
+        public void write(ByteBuffer buf, CompletionHandler handler) {
+            CompletionHandlers.tryComplete(handler);
+        }
+
+        @Override
+        public void close(CompletionHandler handler) {
+            CompletionHandlers.tryComplete(handler);
+        }
+
+    };
+
+    private final FilterResolver filterResolver;
+    private final HttpServletRequest servletRequest;
+
+    public FilteringRequestHandler(FilterResolver filterResolver, HttpServletRequest servletRequest) {
+        this.filterResolver = filterResolver;
+        this.servletRequest = servletRequest;
+    }
+
+    @Override
+    public ContentChannel handleRequest(Request request, ResponseHandler originalResponseHandler) {
+        Preconditions.checkArgument(request instanceof HttpRequest, "Expected HttpRequest, got " + request);
+        Objects.requireNonNull(originalResponseHandler, "responseHandler");
+
+        RequestFilter requestFilter = filterResolver.resolveRequestFilter(servletRequest, request.getUri())
+                .orElse(null);
+        ResponseFilter responseFilter = filterResolver.resolveResponseFilter(servletRequest, request.getUri())
+                .orElse(null);
+
+        // Not using request.connect() here - it adds logic for error handling that we'd rather leave to the framework.
+        RequestHandler resolvedRequestHandler = request.container().resolveHandler(request);
+
+        if (resolvedRequestHandler == null) {
+            throw new BindingNotFoundException(request.getUri());
+        }
+
+        RequestHandler requestHandler = new ReferenceCountingRequestHandler(resolvedRequestHandler);
+
+        ResponseHandler responseHandler;
+        if (responseFilter != null) {
+            responseHandler = new FilteringResponseHandler(originalResponseHandler, responseFilter, request);
+        } else {
+            responseHandler = originalResponseHandler;
+        }
+
+        if (requestFilter != null) {
+            InterceptingResponseHandler interceptingResponseHandler = new InterceptingResponseHandler(responseHandler);
+            requestFilter.filter(HttpRequest.class.cast(request), interceptingResponseHandler);
+            if (interceptingResponseHandler.hasProducedResponse()) {
+                return COMPLETING_CONTENT_CHANNEL;
+            }
+        }
+
+        ContentChannel contentChannel = requestHandler.handleRequest(request, responseHandler);
+        if (contentChannel == null) {
+            throw new RequestDeniedException(request);
+        }
+        return contentChannel;
+    }
+
+    private static class FilteringResponseHandler implements ResponseHandler {
+
+        private final ResponseHandler delegate;
+        private final ResponseFilter responseFilter;
+        private final Request request;
+
+        public FilteringResponseHandler(ResponseHandler delegate, ResponseFilter responseFilter, Request request) {
+            this.delegate = Objects.requireNonNull(delegate);
+            this.responseFilter = Objects.requireNonNull(responseFilter);
+            this.request = request;
+        }
+
+        @Override
+        public ContentChannel handleResponse(Response response) {
+            responseFilter.filter(response, request);
+            return delegate.handleResponse(response);
+        }
+
+    }
+
+    private static class InterceptingResponseHandler implements ResponseHandler {
+
+        private final ResponseHandler delegate;
+        private AtomicBoolean hasResponded = new AtomicBoolean(false);
+
+        public InterceptingResponseHandler(ResponseHandler delegate) {
+            this.delegate = Objects.requireNonNull(delegate);
+        }
+
+        @Override
+        public ContentChannel handleResponse(Response response) {
+            ContentChannel content = delegate.handleResponse(response);
+            hasResponded.set(true);
+            return content;
+        }
+
+        public boolean hasProducedResponse() {
+            return hasResponded.get();
+        }
+    }
+
+}