gobblin.tunnel.ConnectProxyServer.java Source code

Java tutorial

Introduction

Here is the source code for gobblin.tunnel.ConnectProxyServer.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 gobblin.tunnel;

import org.apache.commons.io.IOUtils;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Due to the lack of a suitable embeddable proxy server (the Jetty version here is too old and MockServer's Proxy
 * expects SSL traffic and breaks for arbitrary bytes) we had to write our own mini CONNECT proxy. This simply gets
 * an HTTP CONNECT request over a socket, opens another socket to the specified remote server and relays bytes between
 * the two connections.
 *
 * @author kkandekar@linkedin.com
 */
class ConnectProxyServer extends MockServer {
    private final boolean mixServerAndProxyResponse;
    private final boolean largeResponse;
    Pattern hostPortPattern;
    int nBytesToCloseSocketAfter;

    /**
     * @param mixServerAndProxyResponse Force proxy to send 200 OK and server response in single write such that both
     *                                  responses reach the tunnel in the same read. This can happen for a multitude of
     *                                  reasons, e.g. the proxy GC's, or the network hiccups, or the tunnel GC's.
     * @param largeResponse Force proxy to send a large response
     * @param nBytesToCloseSocketAfter
     */
    public ConnectProxyServer(boolean mixServerAndProxyResponse, boolean largeResponse,
            int nBytesToCloseSocketAfter) {
        this.mixServerAndProxyResponse = mixServerAndProxyResponse;
        this.largeResponse = largeResponse;
        this.nBytesToCloseSocketAfter = nBytesToCloseSocketAfter;
        hostPortPattern = Pattern.compile("Host: (.*):([0-9]+)");
    }

    @Override
    void handleClientSocket(Socket clientSocket) throws IOException {
        final InputStream clientToProxyIn = clientSocket.getInputStream();
        BufferedReader clientToProxyReader = new BufferedReader(new InputStreamReader(clientToProxyIn));
        final OutputStream clientToProxyOut = clientSocket.getOutputStream();
        String line = clientToProxyReader.readLine();
        String connectRequest = "";
        while (line != null && isServerRunning()) {
            connectRequest += line + "\r\n";
            if (connectRequest.endsWith("\r\n\r\n")) {
                break;
            }
            line = clientToProxyReader.readLine();
        }
        // connect to given host:port
        Matcher matcher = hostPortPattern.matcher(connectRequest);
        if (!matcher.find()) {
            try {
                sendConnectResponse("400 Bad Request", clientToProxyOut, null, 0);
            } finally {
                clientSocket.close();
                stopServer();
            }
            return;
        }
        String host = matcher.group(1);
        int port = Integer.decode(matcher.group(2));

        // connect to server
        Socket serverSocket = new Socket();
        try {
            serverSocket.connect(new InetSocketAddress(host, port));
            addSocket(serverSocket);
            byte[] initialServerResponse = null;
            int nbytes = 0;
            if (mixServerAndProxyResponse) {
                // we want to mix the initial server response with the 200 OK
                initialServerResponse = new byte[64];
                nbytes = serverSocket.getInputStream().read(initialServerResponse);
            }
            sendConnectResponse("200 OK", clientToProxyOut, initialServerResponse, nbytes);
        } catch (IOException e) {
            try {
                sendConnectResponse("404 Not Found", clientToProxyOut, null, 0);
            } finally {
                clientSocket.close();
                stopServer();
            }
            return;
        }
        final InputStream proxyToServerIn = serverSocket.getInputStream();
        final OutputStream proxyToServerOut = serverSocket.getOutputStream();
        _threads.add(new EasyThread() {
            @Override
            void runQuietly() throws Exception {
                try {
                    IOUtils.copy(clientToProxyIn, proxyToServerOut);
                } catch (IOException e) {
                    LOG.warn("Exception " + e.getMessage() + " on " + getServerSocketPort());
                }
            }
        }.startThread());
        try {
            if (nBytesToCloseSocketAfter > 0) {
                // Simulate proxy abruptly closing connection
                int leftToRead = nBytesToCloseSocketAfter;
                byte[] buffer = new byte[leftToRead + 256];
                while (true) {
                    int numRead = proxyToServerIn.read(buffer, 0, leftToRead);
                    if (numRead < 0) {
                        break;
                    }
                    clientToProxyOut.write(buffer, 0, numRead);
                    clientToProxyOut.flush();
                    leftToRead -= numRead;
                    if (leftToRead <= 0) {
                        LOG.warn("Cutting connection after " + nBytesToCloseSocketAfter + " bytes");
                        break;
                    }
                }
            } else {
                IOUtils.copy(proxyToServerIn, clientToProxyOut);
            }
        } catch (IOException e) {
            LOG.warn("Exception " + e.getMessage() + " on " + getServerSocketPort());
        }
        clientSocket.close();
        serverSocket.close();
    }

    private void sendConnectResponse(String statusMessage, OutputStream out, byte[] initialServerResponse,
            int initialServerResponseSize) throws IOException {
        String extraHeader = "";
        if (largeResponse) {
            // this is to force multiple reads while draining the proxy CONNECT response in Tunnel. Normal proxy responses
            // won't be this big (well, unless you annoy squid proxy, which happens sometimes), but a select() call
            // waking up for multiple reads before a buffer is full is normal
            for (int i = 0; i < 260; i++) {
                extraHeader += "a";
            }
        }
        byte[] httpResponse = ("HTTP/1.1 " + statusMessage + "\r\nContent-Length: 0\r\nServer: MockProxy"
                + extraHeader + "\r\n\r\n").getBytes();
        if (initialServerResponse != null) {
            byte[] mixedResponse = new byte[httpResponse.length + initialServerResponseSize];
            System.arraycopy(httpResponse, 0, mixedResponse, 0, httpResponse.length);
            System.arraycopy(initialServerResponse, 0, mixedResponse, httpResponse.length,
                    initialServerResponseSize);
            out.write(mixedResponse);
        } else {
            out.write(httpResponse);
        }
        out.flush();
    }
}