ClusterFilter.java

/**
 * Copyright (C) 2011 Marko Lavikainen
 *
 * 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 xxx;

import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URLEncoder;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.Enumeration;
import java.util.List;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.apache.commons.codec.digest.DigestUtils;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.NameValuePair;
import org.apache.http.client.entity.UrlEncodedFormEntity;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpRequestBase;
import org.apache.http.cookie.Cookie;
import org.apache.http.impl.client.BasicCookieStore;
import org.apache.http.impl.client.DefaultHttpClient;
import org.apache.http.impl.cookie.BasicClientCookie;
import org.apache.http.message.BasicHeader;
import org.apache.http.message.BasicNameValuePair;
import org.apache.http.protocol.HTTP;

public class ClusterFilter implements Filter {

    private static final String secret = "secret";
    private static final String ipParam = "some_hidden_param";

    private static final String STRING_LOCATION_HEADER = "Location";
    private static final String HEADER_CONTENT_LENGTH = "Content-Length";

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {

        if (request instanceof HttpServletRequest) {
            String cNode = request.getParameter("cNode");
            if (cNode != null
                    && forwardRequest((HttpServletRequest) request, 
                                      (HttpServletResponse) response,
                                      cNode)) {
                return;
            }
        }
        chain.doFilter(request, response);
    }

    private String getForwardAddr(ServletRequest request, String cNode) {
        String addr = getAddr(request);
        String[] splits = cNode.split("/");
        if (splits.length != 2 || !splits[1].equals(getHash(splits[0]))
                || addr.equals(splits[0])) {
            return null;
        } else {
            return splits[0];
        }
    }

    private boolean forwardRequest(HttpServletRequest request, 
                                   HttpServletResponse response,
                                   String cNode)
                  throws IOException, ServletException,
                         MalformedURLException {

        String forwardAddr = getForwardAddr(request, cNode);
        if (forwardAddr != null) {

            StringBuilder url = new StringBuilder("http://");

            url.append(forwardAddr)
                    .append(request.getRequestURI())
                    .append("?cNode=").append(URLEncoder.encode(cNode, "UTF-8"))
                    .append("&").append(ipParam).append("=").append(URLEncoder.encode(
                            request.getRemoteAddr(), "UTF-8"));

            if ("GET".equals(request.getMethod())) {
                forwardGet(request, response, url.toString());
            } else if ("POST".equals(request.getMethod())) {
                forwardPost(request, response, url.toString());
            }
            return true;
        } else {
            return false;
        }
    }

    private void forwardPost(HttpServletRequest request,
                             HttpServletResponse response,
                             String forwardAddr) throws IOException, ServletException {

        HttpPost method = new HttpPost(forwardAddr);
        setProxyRequestHeaders(request, method);

        @SuppressWarnings("unchecked")
        Enumeration<String> parameters = request.getParameterNames();
        List<NameValuePair> valuePairs = new ArrayList<NameValuePair>();
        while (parameters.hasMoreElements()) {
            String name = parameters.nextElement();
            for (String value : request.getParameterValues(name)) {
                valuePairs.add(new BasicNameValuePair(name, value));
            }
        }
        method.setEntity(new UrlEncodedFormEntity(valuePairs, HTTP.UTF_8));
        executeProxyRequest(method, request, response);
    }

    @Override
    public void destroy() {
    }

    public static String getRemoteAddr(HttpServletRequest request) {
        String ip = request.getParameter(ipParam);
        return ip != null ? ip : request.getRemoteAddr();
    }

    public static String getCNodeValue(ServletRequest request) {
        String addr = getAddr(request);
        return addr + "/" + getHash(addr);
    }

    private static String getHash(String addr) {
        return DigestUtils.md5Hex(addr + secret);
    }

    private static String getAddr(ServletRequest request) {
        return request.getLocalAddr() + ":" + request.getLocalPort();
    }

    private void forwardGet(HttpServletRequest request,
                            HttpServletResponse response,
                            String forwardAddr) throws IOException, ServletException {
        HttpGet method = new HttpGet(forwardAddr);
        setProxyRequestHeaders(request, method);
        this.executeProxyRequest(method, request, response);
    }

    private void executeProxyRequest(
            HttpRequestBase method,
            HttpServletRequest httpServletRequest,
            HttpServletResponse response)
                  throws IOException, ServletException {

        DefaultHttpClient httpclient = new DefaultHttpClient();

        javax.servlet.http.Cookie[] servletCookies = httpServletRequest.getCookies();
        BasicCookieStore cookieStore = new BasicCookieStore();
        if (servletCookies != null) {
            for (javax.servlet.http.Cookie servletCookie : httpServletRequest.getCookies()) {
                BasicClientCookie cookie = new BasicClientCookie(servletCookie.getName(),
                        servletCookie.getValue());
                cookie.setComment(servletCookie.getComment());
                cookie.setDomain(servletCookie.getDomain());
                if (servletCookie.getMaxAge() > 0) {
                    Calendar now = Calendar.getInstance();
                    now.add(Calendar.SECOND, servletCookie.getMaxAge());
                    cookie.setExpiryDate(now.getTime());
                }
                cookie.setPath(servletCookie.getPath());
                cookie.setSecure(servletCookie.getSecure());
                cookie.setVersion(servletCookie.getVersion());
            }
        }
        httpclient.setCookieStore(cookieStore);
        HttpResponse httpResponse = httpclient.execute(method);
        HttpEntity entity = httpResponse.getEntity();
        int responseCode = httpResponse.getStatusLine().getStatusCode();
        method.addHeader(new BasicHeader("Content-Type",
                "application/x-www-form-urlencoded; charset=UTF-8"));

        if (responseCode >= HttpServletResponse.SC_MULTIPLE_CHOICES
                        && responseCode < HttpServletResponse.SC_NOT_MODIFIED) {

            String location = httpResponse.getFirstHeader(STRING_LOCATION_HEADER).getValue();

            if (location == null) {
                throw new ServletException(
                        "Recieved status code: " + responseCode
                                + " but no " + STRING_LOCATION_HEADER
                                + " header was found in the response");
            }

            response.sendRedirect(location);
            return;
        } else if (responseCode == HttpServletResponse.SC_NO_CONTENT) {
            response.setIntHeader(HEADER_CONTENT_LENGTH, 0);
            response.setStatus(HttpServletResponse.SC_NO_CONTENT);
            return;
        }

        response.setStatus(responseCode);

        Header[] headerArrayResponse = httpResponse.getAllHeaders();
        for (Header header : headerArrayResponse) {
            response.setHeader(header.getName(), header.getValue());
        }
        
        entity.writeTo(response.getOutputStream());

        for (Cookie cookie : cookieStore.getCookies()) {
            javax.servlet.http.Cookie responseCookie = new javax.servlet.http.Cookie(cookie.getName(), cookie.getValue());
            responseCookie.setComment(cookie.getComment());
            responseCookie.setDomain(cookie.getDomain());
            responseCookie.setPath(cookie.getPath());
            responseCookie.setSecure(cookie.isSecure());
            responseCookie.setVersion(cookie.getVersion());
            if (cookie.getExpiryDate() != null) {
                responseCookie.setMaxAge(
                        (int)(cookie.getExpiryDate().getTime() - System.currentTimeMillis())/1000);
            }
            response.addCookie(responseCookie);
        }
        
        entity.getContent().close();
        response.setCharacterEncoding("UTF-8");
        response.getOutputStream().flush();
        response.getOutputStream().close();
        httpclient.getConnectionManager().shutdown();
    }

    @SuppressWarnings("unchecked")
    private void setProxyRequestHeaders(
            HttpServletRequest request,
            HttpRequestBase method) {

        Enumeration<String> headers = request.getHeaderNames();
        while (headers.hasMoreElements()) {
            String headerName = headers.nextElement();
            if (!headerName.equalsIgnoreCase(HEADER_CONTENT_LENGTH)) {
                Enumeration<String> values = request.getHeaders(headerName);
                while (values.hasMoreElements()) {
                    String headerValue = values.nextElement();
                    method.setHeader(new BasicHeader(headerName, headerValue));
                }
            }
        }
    }
}