Server.java

/*
 * This file is part of the OWASP Proxy, a free intercepting proxy library.
 * Copyright (C) 2008-2010 Rogan Dawes <rogan@dawes.za.net>
 * 
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * 
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to:
 * The Free Software Foundation, Inc., 
 * 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */

package org.owasp.proxy.daemon;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.logging.Logger;

/**
 * This class implements a TCP server. The user is required to provide an implementation of {@link ConnectionHandler}
 * implementing the desired protocol.
 * 
 * The most basic (echo) server might look like:
 * 
 * <code>
 * 	InetSocketAddress address = new InetSocketAddress("localhost", 8008);
 * 	Server echo = new Server(address, new ConnectionHandler() {
 *  	protected void handleConnection(Socket socket) throws IOException {
 *  		InputStream in = socket.getInputStream();
 *  		OutputStream out = socket.getOutputStream();
 *  		byte[] buff = new byte[1024];
 *  		int got;
 *  		while ((got = in.read(buff)) > -1)
 *  			out.write(buff, 0, got);
 *  	}
 *  }); 
 * 	echo.start();
 * 	
 * 	<wait for signal to stop>
 * 	
 * 	if (!echo.stop()) {
 * 		// error stopping server 
 * 	}
 * </code>
 * 
 * @author Rogan Dawes
 * 
 */
public class Server {

    private final static Logger logger = Logger.getLogger(Server.class
            .getName());

    private ServerSocket socket;

    private int socketTimeout = 0;

    private ConnectionHandler connectionHandler;

    private Executor executor;

    public Server(InetSocketAddress listen, ConnectionHandler connectionHandler)
            throws IOException {
        this(listen, null, connectionHandler);
    }

    public Server(final InetSocketAddress listen, Executor executor,
            ConnectionHandler connectionHandler) throws IOException {
        if (listen == null)
            throw new NullPointerException("listen may not be null");
        if (executor == null)
            executor = Executors.newCachedThreadPool(new ThreadFactory() {
                private int threadCount = 0;
                private String prefix = listen + "-";

                private synchronized String getName() {
                    return prefix + (threadCount++);
                }

                /*
                 * (non-Javadoc)
                 * 
                 * @see java.util.concurrent.ThreadFactory#newThread(java.lang.Runnable )
                 */
                public Thread newThread(Runnable r) {
                    Thread t = new Thread(r, getName());
                    t.setDaemon(true);
                    return t;
                }
            });
        if (connectionHandler == null)
            throw new NullPointerException("connectionHandler may not be null");
        socket = new ServerSocket(listen.getPort(), 20, listen.getAddress());
        socket.setReuseAddress(true);
        this.executor = executor;
        this.connectionHandler = connectionHandler;
    }

    /**
     * @return the socketTimeout
     */
    public int getSocketTimeout() {
        return socketTimeout;
    }

    /**
     * @param socketTimeout
     *            the socketTimeout to set
     */
    public void setSocketTimeout(int socketTimeout) {
        this.socketTimeout = socketTimeout;
    }

    private void handleConnection(final Socket socket) throws IOException {
        executor.execute(new SocketHandler(socket));
    }

    private AcceptThread acceptThread = null;

    private class AcceptThread extends Thread {

        public AcceptThread() {
            super("Accept: " + socket.getLocalSocketAddress());
        }

        public void run() {
            try {
                do {
                    handleConnection(socket.accept());
                } while (!socket.isClosed());
            } catch (IOException ioe) {
                if (!socket.isClosed()) {
                    ioe.printStackTrace();
                    logger.warning("Exception listening for connections: "
                            + ioe.getMessage());
                }
            }
            try {
                if (socket != null && !socket.isClosed())
                    socket.close();
            } catch (IOException ioe) {
                logger.warning("Exception closing socket: " + ioe.getMessage());
            }
            synchronized (Server.this) {
                Server.this.notifyAll();
            }
        }
    }

    public synchronized void start() {
        if (acceptThread == null) {
            acceptThread = new AcceptThread();
            acceptThread.setDaemon(true);
        } else if (acceptThread.isAlive()) {
            throw new IllegalStateException(
                    "Already running in another thread!");
        }
        acceptThread.start();
    }

    public synchronized boolean stop() {
        if (!isStopped()) {
            try {
                socket.close();
            } catch (IOException ioe) {
                ioe.printStackTrace();
            }
            while (!isStopped()) {
                int loop = 0;
                try {
                    wait(1000);
                } catch (InterruptedException ie) {
                    loop++;
                    if (loop > 10)
                        return false;
                }
            }
        }
        return true;
    }

    public synchronized boolean isStopped() {
        return acceptThread == null || !acceptThread.isAlive();
    }

    private class SocketHandler implements Runnable {
        private Socket socket;

        public SocketHandler(Socket socket) {
            this.socket = socket;
        }

        public void run() {
            try {
                socket.setSoTimeout(socketTimeout);
                connectionHandler.handleConnection(socket);
            } catch (Exception ignore) {
                ignore.printStackTrace();
            } finally {
                try {
                    if (!socket.isClosed())
                        socket.close();
                } catch (IOException ignore) {
                }
            }
        }
    }
}