Add balancer servlet

Balancer Servlet tests

ProxyPassReverse first draft

ProxyPassReverse fix

4 files changed, 705 insertions, 4 deletions

diff --git a/jetty-servlets/src/main/java/org/eclipse/jetty/servlets/BalancerServlet.java b/jetty-servlets/src/main/java/org/eclipse/jetty/servlets/BalancerServlet.java
new file mode 100644
index 0000000000..9ff9536077
--- /dev/null
+++ b/jetty-servlets/src/main/java/org/eclipse/jetty/servlets/BalancerServlet.java
@@ -0,0 +1,417 @@
+// ========================================================================
+// Copyright (c) 2009-2009 Mort Bay Consulting Pty. Ltd.
+// ------------------------------------------------------------------------
+// All rights reserved. This program and the accompanying materials
+// are made available under the terms of the Eclipse Public License v1.0
+// and Apache License v2.0 which accompanies this distribution.
+// The Eclipse Public License is available at 
+// http://www.eclipse.org/legal/epl-v10.html
+// The Apache License v2.0 is available at
+// http://www.opensource.org/licenses/apache2.0.php
+// You may elect to redistribute this code under either of these licenses. 
+// ========================================================================
+
+package org.eclipse.jetty.servlets;
+
+import java.net.MalformedURLException;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import javax.servlet.ServletConfig;
+import javax.servlet.ServletException;
+import javax.servlet.UnavailableException;
+import javax.servlet.http.Cookie;
+import javax.servlet.http.HttpServletRequest;
+
+import org.eclipse.jetty.http.HttpURI;
+import org.eclipse.jetty.server.Request;
+
+/**
+ * 
+ */
+public class BalancerServlet extends ProxyServlet
+{
+
+    private static final class BalancerMember
+    {
+
+        private String _name;
+
+        private String _proxyTo;
+
+        private HttpURI _backendURI;
+
+        public BalancerMember(String name, String proxyTo)
+        {
+            super();
+            _name = name;
+            _proxyTo = proxyTo;
+            _backendURI = new HttpURI(_proxyTo);
+        }
+
+        public String getProxyTo()
+        {
+            return _proxyTo;
+        }
+
+        public HttpURI getBackendURI()
+        {
+            return _backendURI;
+        }
+
+        @Override
+        public String toString()
+        {
+            return "BalancerMember [_name=" + _name + ", _proxyTo=" + _proxyTo + "]";
+        }
+
+        @Override
+        public int hashCode()
+        {
+            final int prime = 31;
+            int result = 1;
+            result = prime * result + ((_name == null)?0:_name.hashCode());
+            return result;
+        }
+
+        @Override
+        public boolean equals(Object obj)
+        {
+            if (this == obj)
+                return true;
+            if (obj == null)
+                return false;
+            if (getClass() != obj.getClass())
+                return false;
+            BalancerMember other = (BalancerMember)obj;
+            if (_name == null)
+            {
+                if (other._name != null)
+                    return false;
+            }
+            else if (!_name.equals(other._name))
+                return false;
+            return true;
+        }
+
+    }
+
+    private static final class RoundRobinIterator implements Iterator<BalancerMember>
+    {
+
+        private BalancerMember[] _balancerMembers;
+
+        private AtomicInteger _index;
+
+        public RoundRobinIterator(Collection<BalancerMember> balancerMembers)
+        {
+            _balancerMembers = (BalancerMember[])balancerMembers.toArray(new BalancerMember[balancerMembers.size()]);
+            _index = new AtomicInteger(-1);
+        }
+
+        public boolean hasNext()
+        {
+            return true;
+        }
+
+        public BalancerMember next()
+        {
+            BalancerMember balancerMember = null;
+            while (balancerMember == null)
+            {
+                int currentIndex = _index.get();
+                int nextIndex = (currentIndex + 1) % _balancerMembers.length;
+                if (_index.compareAndSet(currentIndex,nextIndex))
+                {
+                    balancerMember = _balancerMembers[nextIndex];
+                }
+            }
+            return balancerMember;
+        }
+
+        public void remove()
+        {
+            throw new UnsupportedOperationException();
+        }
+
+    }
+
+    private static final String BALANCER_MEMBER_PREFIX = "BalancerMember.";
+
+    private static final List<String> FORBIDDEN_CONFIG_PARAMETERS;
+    static
+    {
+        List<String> params = new LinkedList<String>();
+        params.add("HostHeader");
+        params.add("whiteList");
+        params.add("blackList");
+        FORBIDDEN_CONFIG_PARAMETERS = Collections.unmodifiableList(params);
+    }
+
+    private static final List<String> REVERSE_PROXY_HEADERS;
+    static
+    {
+        List<String> params = new LinkedList<String>();
+        params.add("Location");
+        params.add("Content-Location");
+        params.add("URI");
+        REVERSE_PROXY_HEADERS = Collections.unmodifiableList(params);
+    }
+
+    private static final String JSESSIONID = "jsessionid";
+
+    private static final String JSESSIONID_URL_PREFIX = JSESSIONID + "=";
+
+    private boolean _stickySessions;
+
+    private Set<BalancerMember> _balancerMembers = new HashSet<BalancerMember>();
+
+    private boolean _proxyPassReverse;
+
+    private RoundRobinIterator _roundRobinIterator;
+
+    @Override
+    public void init(ServletConfig config) throws ServletException
+    {
+        validateConfig(config);
+        super.init(config);
+        initStickySessions(config);
+        initBalancers(config);
+        initProxyPassReverse(config);
+        postInit();
+    }
+
+    private void validateConfig(ServletConfig config) throws ServletException
+    {
+        @SuppressWarnings("unchecked")
+        List<String> initParameterNames = Collections.list(config.getInitParameterNames());
+        for (String initParameterName : initParameterNames)
+        {
+            if (FORBIDDEN_CONFIG_PARAMETERS.contains(initParameterName))
+            {
+                throw new UnavailableException(initParameterName + " not supported in " + getClass().getName());
+            }
+        }
+    }
+
+    private void initStickySessions(ServletConfig config) throws ServletException
+    {
+        _stickySessions = "true".equalsIgnoreCase(config.getInitParameter("StickySessions"));
+    }
+
+    private void initBalancers(ServletConfig config) throws ServletException
+    {
+        Set<String> balancerNames = getBalancerNames(config);
+        for (String balancerName : balancerNames)
+        {
+            String memberProxyToParam = BALANCER_MEMBER_PREFIX + balancerName + ".ProxyTo";
+            String proxyTo = config.getInitParameter(memberProxyToParam);
+            if (proxyTo == null || proxyTo.trim().length() == 0)
+            {
+                throw new UnavailableException(memberProxyToParam + " parameter is empty.");
+            }
+            _balancerMembers.add(new BalancerMember(balancerName,proxyTo));
+        }
+    }
+
+    private void initProxyPassReverse(ServletConfig config)
+    {
+        _proxyPassReverse = "true".equalsIgnoreCase(config.getInitParameter("ProxyPassReverse"));
+    }
+
+    private void postInit()
+    {
+        _roundRobinIterator = new RoundRobinIterator(_balancerMembers);
+    }
+
+    private Set<String> getBalancerNames(ServletConfig config) throws ServletException
+    {
+        Set<String> names = new HashSet<String>();
+        @SuppressWarnings("unchecked")
+        List<String> initParameterNames = Collections.list(config.getInitParameterNames());
+        for (String initParameterName : initParameterNames)
+        {
+            if (!initParameterName.startsWith(BALANCER_MEMBER_PREFIX))
+            {
+                continue;
+            }
+            int endOfNameIndex = initParameterName.lastIndexOf(".");
+            if (endOfNameIndex <= BALANCER_MEMBER_PREFIX.length())
+            {
+                throw new UnavailableException(initParameterName + " parameter does not provide a balancer member name");
+            }
+            names.add(initParameterName.substring(BALANCER_MEMBER_PREFIX.length(),endOfNameIndex));
+        }
+        return names;
+    }
+
+    @Override
+    protected HttpURI proxyHttpURI(HttpServletRequest request, String uri) throws MalformedURLException
+    {
+        BalancerMember balancerMember = selectBalancerMember(request);
+        try
+        {
+            URI dstUri = new URI(balancerMember.getProxyTo() + "/" + uri).normalize();
+            return new HttpURI(dstUri.toString());
+        }
+        catch (URISyntaxException e)
+        {
+            throw new MalformedURLException(e.getMessage());
+        }
+    }
+
+    private BalancerMember selectBalancerMember(HttpServletRequest request)
+    {
+        BalancerMember balancerMember = null;
+        if (_stickySessions)
+        {
+            String name = getBalancerMemberNameFromSessionId(request);
+            if (name != null)
+            {
+                balancerMember = findBalancerMemberByName(name);
+                if (balancerMember != null)
+                {
+                    return balancerMember;
+                }
+            }
+        }
+        return _roundRobinIterator.next();
+    }
+
+    private BalancerMember findBalancerMemberByName(String name)
+    {
+        BalancerMember example = new BalancerMember(name,"");
+        for (BalancerMember balancerMember : _balancerMembers)
+        {
+            if (balancerMember.equals(example))
+            {
+                return balancerMember;
+            }
+        }
+        return null;
+    }
+
+    private String getBalancerMemberNameFromSessionId(HttpServletRequest request)
+    {
+        String name = getBalancerMemberNameFromSessionCookie(request);
+        if (name == null)
+        {
+            name = getBalancerMemberNameFromURL(request);
+        }
+        return name;
+    }
+
+    private String getBalancerMemberNameFromSessionCookie(HttpServletRequest request)
+    {
+        Cookie[] cookies = request.getCookies();
+        String name = null;
+        for (Cookie cookie : cookies)
+        {
+            if (JSESSIONID.equalsIgnoreCase(cookie.getName()))
+            {
+                name = extractBalancerMemberNameFromSessionId(cookie.getValue());
+                break;
+            }
+        }
+        return name;
+    }
+
+    private String getBalancerMemberNameFromURL(HttpServletRequest request)
+    {
+        String name = null;
+        String requestURI = request.getRequestURI();
+        int idx = requestURI.lastIndexOf(";");
+        if (idx != -1)
+        {
+            String requestURISuffix = requestURI.substring(idx);
+            if (requestURISuffix.startsWith(JSESSIONID_URL_PREFIX))
+            {
+                name = extractBalancerMemberNameFromSessionId(requestURISuffix.substring(JSESSIONID_URL_PREFIX.length()));
+            }
+        }
+        return name;
+    }
+
+    private String extractBalancerMemberNameFromSessionId(String sessionId)
+    {
+        String name = null;
+        int idx = sessionId.lastIndexOf(".");
+        if (idx != -1)
+        {
+            String sessionIdSuffix = sessionId.substring(idx + 1);
+            name = (sessionIdSuffix.length() > 0)?sessionIdSuffix:null;
+        }
+        return name;
+    }
+
+    @Override
+    protected String filterResponseHeaderValue(String headerName, String headerValue, HttpServletRequest request)
+    {
+        if (_proxyPassReverse && REVERSE_PROXY_HEADERS.contains(headerName))
+        {
+            HttpURI locationURI = new HttpURI(headerValue);
+            if (isAbsoluteLocation(locationURI) && isBackendLocation(locationURI))
+            {
+                Request jettyRequest = (Request)request;
+                URI reverseUri;
+                try
+                {
+                    reverseUri = new URI(jettyRequest.getRootURL().append(locationURI.getCompletePath()).toString()).normalize();
+                    return reverseUri.toURL().toString();
+                }
+                catch (Exception e)
+                {
+                    _log.warn("Not filtering header response",e);
+                    return headerValue;
+                }
+            }
+        }
+        return headerValue;
+    }
+
+    private boolean isBackendLocation(HttpURI locationURI)
+    {
+        for (BalancerMember balancerMember : _balancerMembers)
+        {
+            HttpURI backendURI = balancerMember.getBackendURI();
+            if (backendURI.getHost().equals(locationURI.getHost()) && backendURI.getScheme().equals(locationURI.getScheme())
+                    && backendURI.getPort() == locationURI.getPort())
+            {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private boolean isAbsoluteLocation(HttpURI locationURI)
+    {
+        return locationURI.getHost() != null;
+    }
+
+    @Override
+    public String getHostHeader()
+    {
+        throw new UnsupportedOperationException("HostHeader not supported in " + getClass().getName());
+    }
+
+    @Override
+    public void setHostHeader(String hostHeader)
+    {
+        throw new UnsupportedOperationException("HostHeader not supported in " + getClass().getName());
+    }
+
+    @Override
+    public boolean validateDestination(String host, String path)
+    {
+        return true;
+    }
+
+}
\ No newline at end of file
diff --git a/jetty-servlets/src/main/java/org/eclipse/jetty/servlets/ProxyServlet.java b/jetty-servlets/src/main/java/org/eclipse/jetty/servlets/ProxyServlet.java
index 7fa971838a..6687903a6a 100644
--- a/jetty-servlets/src/main/java/org/eclipse/jetty/servlets/ProxyServlet.java
+++ b/jetty-servlets/src/main/java/org/eclipse/jetty/servlets/ProxyServlet.java
@@ -483,13 +483,20 @@ public class ProxyServlet implements Servlet
                     @Override
                     protected void onResponseHeader(Buffer name, Buffer value) throws IOException
                     {
-                        String s = name.toString().toLowerCase();
+                        String nameString = name.toString();
+                        String s = nameString.toLowerCase();
                         if (!_DontProxyHeaders.contains(s) || (HttpHeaders.CONNECTION_BUFFER.equals(name) && HttpHeaderValues.CLOSE_BUFFER.equals(value)))
                         {
                             if (debug != 0)
                                 _log.debug(debug + " " + name + ": " + value);
 
-                            response.addHeader(name.toString(),value.toString());
+                            String filteredHeaderValue = filterResponseHeaderValue(nameString,value.toString(),request);
+                            if (filteredHeaderValue != null && filteredHeaderValue.trim().length() > 0)
+                            {
+                                if (debug != 0)
+                                    _log.debug(debug + " " + name + ": (filtered): " + filteredHeaderValue);
+                                response.addHeader(nameString,filteredHeaderValue);
+                            }
                         }
                         else if (debug != 0)
                             _log.debug(debug + " " + name + "! " + value);
@@ -786,8 +793,22 @@ public class ProxyServlet implements Servlet
     }
 
     /**
+     * Extension point for remote server response header filtering. The default implementation returns the header value as is. If null is returned, this header
+     * won't be forwarded back to the client.
+     * 
+     * @param headerName
+     * @param headerValue
+     * @param request
+     * @return filteredHeaderValue
+     */
+    protected String filterResponseHeaderValue(String headerName, String headerValue, HttpServletRequest request)
+    {
+        return headerValue;
+    }
+
+    /**
      * Transparent Proxy.
-     *
+     * 
      * This convenience extension to ProxyServlet configures the servlet as a transparent proxy. The servlet is configured with init parameters:
      * <ul>
      * <li>ProxyTo - a URI like http://host:80/context to which the request is proxied.
@@ -795,7 +816,7 @@ public class ProxyServlet implements Servlet
      * </ul>
      * For example, if a request was received at /foo/bar and the ProxyTo was http://host:80/context and the Prefix was /foo, then the request would be proxied
      * to http://host:80/context/bar
-     *
+     * 
      */
     public static class Transparent extends ProxyServlet
     {
diff --git a/jetty-servlets/src/test/java/org/eclipse/jetty/servlets/AbstractBalancerServletTest.java b/jetty-servlets/src/test/java/org/eclipse/jetty/servlets/AbstractBalancerServletTest.java
new file mode 100644
index 0000000000..f6ab92c172
--- /dev/null
+++ b/jetty-servlets/src/test/java/org/eclipse/jetty/servlets/AbstractBalancerServletTest.java
@@ -0,0 +1,146 @@
+package org.eclipse.jetty.servlets;
+
+import java.io.IOException;
+
+import javax.servlet.http.HttpServlet;
+
+import org.eclipse.jetty.client.ContentExchange;
+import org.eclipse.jetty.client.HttpClient;
+import org.eclipse.jetty.http.HttpCookie;
+import org.eclipse.jetty.http.HttpMethods;
+import org.eclipse.jetty.io.Buffer;
+import org.eclipse.jetty.server.Server;
+import org.eclipse.jetty.server.nio.SelectChannelConnector;
+import org.eclipse.jetty.server.session.HashSessionIdManager;
+import org.eclipse.jetty.servlet.ServletContextHandler;
+import org.eclipse.jetty.servlet.ServletHolder;
+import org.junit.After;
+import org.junit.Before;
+
+/**
+ * @author tsegismont
+ */
+public abstract class AbstractBalancerServletTest
+{
+
+    private boolean _stickySessions;
+
+    private Server _node1;
+
+    private Server _node2;
+
+    private Server _balancerServer;
+
+    private HttpClient _httpClient;
+
+    @Before
+    public void setUp() throws Exception
+    {
+        _httpClient = new HttpClient();
+        _httpClient.registerListener("org.eclipse.jetty.client.RedirectListener");
+        _httpClient.start();
+    }
+
+    @After
+    public void tearDown() throws Exception
+    {
+        stopServer(_node1);
+        stopServer(_node2);
+        stopServer(_balancerServer);
+        _httpClient.stop();
+    }
+
+    private void stopServer(Server server)
+    {
+        try
+        {
+            server.stop();
+        }
+        catch (Exception e)
+        {
+            // Do nothing
+        }
+    }
+
+    protected void setStickySessions(boolean stickySessions)
+    {
+        _stickySessions = stickySessions;
+    }
+
+    protected void startBalancer(Class<? extends HttpServlet> httpServletClass) throws Exception
+    {
+        _node1 = createServer(new ServletHolder(httpServletClass.newInstance()),"/pipo","/molo/*");
+        setSessionIdManager(_node1,"node1");
+        _node1.start();
+
+        _node2 = createServer(new ServletHolder(httpServletClass.newInstance()),"/pipo","/molo/*");
+        setSessionIdManager(_node2,"node2");
+        _node2.start();
+
+        BalancerServlet balancerServlet = new BalancerServlet();
+        ServletHolder balancerServletHolder = new ServletHolder(balancerServlet);
+        balancerServletHolder.setInitParameter("StickySessions",String.valueOf(_stickySessions));
+        balancerServletHolder.setInitParameter("ProxyPassReverse","true");
+        balancerServletHolder.setInitParameter("BalancerMember." + "node1" + ".ProxyTo","http://localhost:" + getServerPort(_node1));
+        balancerServletHolder.setInitParameter("BalancerMember." + "node2" + ".ProxyTo","http://localhost:" + getServerPort(_node2));
+
+        _balancerServer = createServer(balancerServletHolder,"/pipo","/molo/*");
+        _balancerServer.start();
+    }
+
+    private Server createServer(ServletHolder servletHolder, String appContext, String servletUrlPattern)
+    {
+        Server server = new Server();
+        SelectChannelConnector httpConnector = new SelectChannelConnector();
+        server.addConnector(httpConnector);
+
+        ServletContextHandler context = new ServletContextHandler(ServletContextHandler.SESSIONS);
+        context.setContextPath(appContext);
+        server.setHandler(context);
+
+        context.addServlet(servletHolder,servletUrlPattern);
+
+        return server;
+    }
+
+    private void setSessionIdManager(Server node, String nodeName)
+    {
+        HashSessionIdManager sessionIdManager = new HashSessionIdManager();
+        sessionIdManager.setWorkerName(nodeName);
+        node.setSessionIdManager(sessionIdManager);
+    }
+
+    private int getServerPort(Server node)
+    {
+        return node.getConnectors()[0].getLocalPort();
+    }
+
+    protected byte[] sendRequestToBalancer(String requestUri) throws IOException, InterruptedException
+    {
+        ContentExchange exchange = new ContentExchange()
+        {
+            @Override
+            protected void onResponseHeader(Buffer name, Buffer value) throws IOException
+            {
+                // Cookie persistence
+                if (name.toString().equals("Set-Cookie"))
+                {
+                    String cookieVal = value.toString();
+                    if (cookieVal.startsWith("JSESSIONID="))
+                    {
+                        String jsessionid = cookieVal.split(";")[0].substring("JSESSIONID=".length());
+                        _httpClient.getDestination(getAddress(),false).addCookie(new HttpCookie("JSESSIONID",jsessionid));
+                    }
+                }
+            }
+        };
+        exchange.setURL("http://localhost:" + getServerPort(_balancerServer) + "/pipo/molo/" + requestUri);
+        exchange.setMethod(HttpMethods.GET);
+
+        _httpClient.send(exchange);
+        exchange.waitForDone();
+
+        return exchange.getResponseContentBytes();
+    }
+
+}
diff --git a/jetty-servlets/src/test/java/org/eclipse/jetty/servlets/BalancerServletTest.java b/jetty-servlets/src/test/java/org/eclipse/jetty/servlets/BalancerServletTest.java
new file mode 100644
index 0000000000..139e272ab0
--- /dev/null
+++ b/jetty-servlets/src/test/java/org/eclipse/jetty/servlets/BalancerServletTest.java
@@ -0,0 +1,117 @@
+package org.eclipse.jetty.servlets;
+
+import static org.junit.Assert.*;
+
+import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStreamReader;
+
+import javax.servlet.ServletException;
+import javax.servlet.http.HttpServlet;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
+
+import org.junit.Test;
+
+/**
+ * @author tsegismont
+ */
+public class BalancerServletTest extends AbstractBalancerServletTest
+{
+
+    @Test
+    public void testRoundRobinBalancer() throws Exception
+    {
+        setStickySessions(false);
+        startBalancer(CounterServlet.class);
+
+        for (int i = 0; i < 10; i++)
+        {
+            byte[] responseBytes = sendRequestToBalancer("/");
+            String returnedCounter = readFirstLine(responseBytes);
+            // RR : response should increment every other request
+            String expectedCounter = String.valueOf(i / 2);
+            assertEquals(expectedCounter,returnedCounter);
+        }
+    }
+
+    @Test
+    public void testStickySessionsBalancer() throws Exception
+    {
+        setStickySessions(true);
+        startBalancer(CounterServlet.class);
+
+        for (int i = 0; i < 10; i++)
+        {
+            byte[] responseBytes = sendRequestToBalancer("/");
+            String returnedCounter = readFirstLine(responseBytes);
+            // RR : response should increment on each request
+            String expectedCounter = String.valueOf(i);
+            assertEquals(expectedCounter,returnedCounter);
+        }
+    }
+
+    @Test
+    public void testProxyPassReverse() throws Exception
+    {
+        setStickySessions(false);
+        startBalancer(RelocationServlet.class);
+
+        byte[] responseBytes = sendRequestToBalancer("index.html");
+        String msg = readFirstLine(responseBytes);
+        assertEquals("success",msg);
+    }
+
+    private String readFirstLine(byte[] responseBytes) throws IOException
+    {
+        BufferedReader reader = new BufferedReader(new InputStreamReader(new ByteArrayInputStream(responseBytes)));
+        return reader.readLine();
+    }
+
+    @SuppressWarnings("serial")
+    public static final class CounterServlet extends HttpServlet
+    {
+
+        private int counter;
+
+        @Override
+        public void init() throws ServletException
+        {
+            counter = 0;
+        }
+
+        @Override
+        protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException
+        {
+            // Force session creation
+            req.getSession();
+            resp.setContentType("text/plain");
+            resp.getWriter().println(counter++);
+        }
+    }
+
+    @SuppressWarnings("serial")
+    public static final class RelocationServlet extends HttpServlet
+    {
+        @Override
+        protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException
+        {
+            if (req.getRequestURI().endsWith("/index.html"))
+            {
+                resp.sendRedirect("http://localhost:" + req.getLocalPort() + req.getContextPath() + req.getServletPath() + "/other.html?secret=pipo%20molo");
+                return;
+            }
+            resp.setContentType("text/plain");
+            if ("pipo molo".equals(req.getParameter("secret")))
+            {
+                resp.getWriter().println("success");
+            }
+            else
+            {
+                resp.getWriter().println("failure");
+            }
+        }
+    }
+
+}