path: root/src/main/java/org/gnunet/util/Server.java

/*
 This file is part of GNUnet.
 Copyright (C) 2011, 2012 Christian Grothoff (and other contributing authors)

 GNUnet is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published
 by the Free Software Foundation; either version 3, or (at your
 option) any later version.

 GNUnet 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
 General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with GNUnet; see the file COPYING.  If not, write to the
 Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
 Boston, MA 02110-1301, USA.
 */

package org.gnunet.util;

import org.grothoff.Runabout;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.SocketAddress;
import java.nio.channels.SelectionKey;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

/**
 * A server allows to wait for incoming connections from clients and respectively communicate with those clients.
 */
public class Server {
    private static final Logger logger = LoggerFactory
            .getLogger(Server.class);

    /**
     * Default idle timeout for new clients.
     */
    private final RelativeTime idleTimeout;

    /**
     * If true, disconnect a client when it sends a message we do not expect to receive. Otherwise, the unexpected
     * message will just be discarded.
     */
    private final boolean requireFound;

    /**
     * The sockets this server accepts new connections on.
     */
    private List<ServerSocketChannel> listenSockets = new ArrayList<ServerSocketChannel>();

    /**
     * The list of all clients connected to this server.
     */
    private List<ClientHandle> clientHandles = new LinkedList<ClientHandle>();

    /**
     * The runabout that receives received messages, as well as information about the sender of the last
     * received message.
     */
    private MessageRunabout receivedMessageHandler;

    /**
     * Whenever a client is disconnected all disconnect handlers are informed.
     */
    private List<DisconnectHandler> disconnectHandlers = new LinkedList<DisconnectHandler>();

    /**
     * Classes of the messages we expect to receive. If a received message is not in this list, the client
     * will be disconnected, otherwise the message is just ignored.
     */
    private List<Class> expectedMessages = Collections.emptyList();

    /**
     * If true, shut down as soon as all non-monitor clients have finished, and do not allow new connections
     * to be made to this server.
     */
    private boolean inSoftShutdown;

    /**
     * Task that is executed as soon as a connection is ready to be accepted.
     */
    private Cancelable acceptTask;

    /**
     * True if we are destroyed, or in the process of being destroyed with no way back.
     */
    private boolean destroyed;


    /**
     * Interface implemented by disconnect handlers, whose onDisconnect method is called whenever a client
     * is disconnected from the server.
     */
    public interface DisconnectHandler {
        /**
         * Called whenever a client is disconnected from the server.
         *
         * @param clientHandle the handle for the client that was disconnected
         */
        void onDisconnect(ClientHandle clientHandle);
    }

    /**
     * A handle to a (remote) client connected to this server.
     * <p/>
     * Every client handle keeps a reference count..
     * Whenever a part of the programs saves a client handle for further interaction with it, keep() should be called.
     * This prevents the server from disconnecting the client when it is idle.
     * Once this interaction is over, drop() will decrement the reference count and eventually disconnect the client
     * after being idle for long enough.
     */
    public class ClientHandle {
        /**
         * The underlying connection to the client-
         */
        private Connection connection;

        /**
         * When referenceCount==0, the server is allowed to drop the client after a timeout.
         */
        private int referenceCount = 0;

        /**
         * Set to true if the connection to this client should not prevent the server from shutting down.
         */
        private boolean isMonitor;

        /**
         * Iff true, disconnect the client as soon as possible.
         * Disconnecting may sometimes not be possible immediately, for example when the reference count is not zero.
         */
        private boolean disconnectRequested;

        /**
         * Create a client handle.
         *
         * @param sock
         */
        private ClientHandle(SocketChannel sock) {
            connection = new Connection(sock);
            // start receiving
            receiveDone(true);
        }

        /**
         * Notify us when the server has enough space to transmit
         * a message of the given size to the given client.
         *
         * @param size        requested amount of buffer space
         * @param timeout     after how long should we give up (and call
         *                    notify with buf NULL and size 0)?
         * @param transmitter callback
         * @return a handle to onCancel the notification
         */
        public Cancelable notifyTransmitReady(int size, RelativeTime timeout, MessageTransmitter transmitter) {
            return connection.notifyTransmitReady(size, timeout, transmitter);
        }

        /**
         * Convenience method for sending messages.
         *
         * @param timeout when should we give up sending the message, and call cont.cont(false)
         * @param message the message to send
         * @param cont called when the message has been sent successfully or on error
         * @return a handle to onCancel sending the message
         */
        public Cancelable transmitWhenReady(final RelativeTime timeout, final GnunetMessage.Body message, final Continuation cont) {
            return notifyTransmitReady(0, timeout, new MessageTransmitter() {
                @Override
                public void transmit(Connection.MessageSink sink) {
                    sink.send(message);
                    if (cont != null) {
                        cont.cont(true);
                    }
                }

                @Override
                public void handleError() {
                    if (cont != null) {
                        cont.cont(false);
                    }
                }
            });
        }

        /**
         * Resume receiving from this client, we are done processing the
         * current getRequestIdentifier. This function must be called from within each
         * message handler (or its respective continuations).
         * <p/>
         * The server does not automatically continue to receive messages to
         * support flow control.
         *
         * @param stayConnected false if connection to the client should be closed
         */
        public void receiveDone(boolean stayConnected) {
            if (stayConnected) {
                connection.receive(RelativeTime.FOREVER, new MessageReceiver() {
                    @Override
                    public void process(GnunetMessage.Body msg) {
                        if ((msg instanceof UnknownMessageBody) || !expectedMessages.contains(msg.getClass())) {
                            if (requireFound) {
                                logger.info("disconnecting client sending unknown message");
                                disconnect();
                            }
                            // otherwise, just ignore it
                        }
                        if (receivedMessageHandler == null) {
                            throw new AssertionError("received message, but no handler installed");
                        }
                        receivedMessageHandler.setSender(ClientHandle.this);
                        receivedMessageHandler.visitAppropriate(msg);
                    }

                    @Override
                    public void handleError() {
                        logger.warn("error receiving from client");
                        disconnect();
                    }
                });
            } else {
                if (referenceCount > 0) {
                    this.disconnectRequested = true;
                } else {
                    System.out.println("disconnecting " + this.isMonitor);
                    disconnect();
                }
            }

        }

        /**
         * Ask the server to disconnect from the given client.
         * <p/>
         * The client will be disconnected from the server, no matter what the current reference count is.
         */
        public void disconnect() {
            connection.disconnect();
            // if we are in the process of destruction, to not remove, the destruction function will do this,
            // removing the client handle while in destruction would yield a concurrent modification exception
            if (!destroyed) {
                Server.this.clientHandles.remove(this);
            }
            for (DisconnectHandler dh : disconnectHandlers) {
                dh.onDisconnect(this);
            }
            Server.this.testForSoftShutdown();
        }

        /**
         * Prevent the client from being disconnected.
         * For every keep, there should be an additional matching drop.
         */
        public void keep() {
            referenceCount++;
        }


        /**
         * Allow to disconnect this client, if not prevented by previous calls to keep.
         * <p/>
         * A call to drop should be executed for every call to keep.
         * After drop() has been executed for every matching keep(), the next call to drop()
         * allows the server to disconnect the client after a timeout.
         */
        public void drop() {
            assert referenceCount > 0;
            referenceCount--;
            if (referenceCount == 0 && disconnectRequested) {
                disconnect();
            }
        }


        /**
         * Set the 'monitor' flag on this client. Clients which have been
         * marked as 'monitors' won't prevent the server from shutting down
         * once 'GNUNET_SERVER_stop_listening' has been invoked.  The idea is
         * that for "normal" clients we likely want to allow them to process
         * their requests; however, monitor-clients are likely to 'never'
         * disconnect during shutdown and thus will not be considered when
         * determining if the server should continue to exist after
         * 'GNUNET_SERVER_destroy' has been called.
         */
        public void markMonitor() {
            this.isMonitor = true;
        }

        public boolean isMonitor() {
            return isMonitor;
        }
    }


    /**
     * All handlers for receiving messages from clients have to inherit this class.
     * <p/>
     * MessageRunabout is a standard runabout with the added possibility of getting the sender of the message.
     * This is necessary as the runabout's visit methods can have only one parameter.
     */
    public abstract static class MessageRunabout extends Runabout {
        private ClientHandle currentSender;

        /**
         * Allows implementors of MessageRunabout to get the Client that sent the message
         * currently visited.
         * <p/>
         * The return value of getSender() is only valid while executing a visit method.
         *
         * @return handle of the client whose message is currently being visited
         */
        public final ClientHandle getSender() {
            return currentSender;
        }

        /**
         * Private method used to set the sender for the getSender() method.
         *
         * @param clientHandle the client handle to set as the sender
         */
        private void setSender(ClientHandle clientHandle) {
            currentSender = clientHandle;
        }
    }

    /**
     * Create a server listening on all specified addresses.
     *
     * @param addresses    addresses to bind on
     * @param idleTimeout  time after a client will be disconnected if idle
     * @param requireFound allow unknown messages to be received without disconnecting the client in response
     */
    public Server(List<SocketAddress> addresses, RelativeTime idleTimeout, boolean requireFound) {
        this.idleTimeout = idleTimeout;
        this.requireFound = requireFound;
        try {
            for (SocketAddress addr : addresses) {
                ServerSocketChannel socket = ServerSocketChannel.open();
                socket.configureBlocking(false);
                socket.socket().bind(addr);
                logger.debug("socket listening on {}", addr.toString());
                listenSockets.add(socket);
                addAcceptSocket(socket);
            }
        } catch (IOException e) {
            throw new RuntimeException("could not bind", e);
        }
    }

    /**
     * Create a server, not listening on any sockets yet for new connections.
     *
     * @param idleTimeout  time after a client will be disconnected if idle
     * @param requireFound allow unknown messages to be received without disconnecting the client in response
     */
    public Server(RelativeTime idleTimeout, boolean requireFound) {
        this.idleTimeout = idleTimeout;
        this.requireFound = requireFound;
    }

    /**
     * Accept new connections from the given server socket.
     *
     * @param sock the new socket to accept connections from
     */
    public final void addAcceptSocket(final ServerSocketChannel sock) {
        Scheduler.TaskConfiguration b = new Scheduler.TaskConfiguration(RelativeTime.FOREVER,
                new Scheduler.Task() {
                    @Override
                    public void run(Scheduler.RunContext ctx) {
                        acceptTask = null;
                        try {
                            SocketChannel cli = sock.accept();

                            if (cli != null) {
                                logger.debug("client connected");
                                cli.configureBlocking(false);
                                ClientHandle clientHandle = new ClientHandle(cli);
                                clientHandles.add(clientHandle);
                            }

                        } catch (IOException e) {
                            throw new RuntimeException("accept failed", e);
                        }
                        addAcceptSocket(sock);
                    }
                });
        b.addSelectEvent(sock, SelectionKey.OP_ACCEPT);
        acceptTask = b.schedule();
    }

    /**
     * Pass messages that the runabout can handle to it.
     * There can only be one runabout per message type.
     * (Discrepancy with the C-API, could be changed in the future)
     *
     * @param msgRunabout handler
     */
    public void setHandler(MessageRunabout msgRunabout) {
        receivedMessageHandler = msgRunabout;
        expectedMessages = RunaboutUtil.getRunaboutVisitees(msgRunabout);
    }

    /**
     * Ask the server to notify us whenever a client disconnects.
     * This handler is called whenever the actual network connection
     * is closed; the reference count may be zero or larger than zero
     * at this point. Note that the disconnect handler is also called when
     *
     * @param disconnectHandler handler to call on disconnect
     */
    public Cancelable notifyDisconnect(final DisconnectHandler disconnectHandler) {
        this.disconnectHandlers.add(disconnectHandler);
        return new Cancelable() {
            @Override
            public void cancel() {
                Server.this.disconnectHandlers.remove(disconnectHandler);
            }
        };
    }

    /**
     * Stop the listen socket destroy the server as soon as only monitor clients are left.
     */
    public void stopListening() {
        inSoftShutdown = true;
        if (acceptTask != null) {
            acceptTask.cancel();
            acceptTask = null;
        }
        testForSoftShutdown();
    }

    /**
     * Disconnect all clients forcefully from the server and stop listening.
     * <p/>
     * No methods should be called on a server and its client handles after destroy() has been called.
     */
    public void destroy() {
        if (destroyed) {
            return;
        }
        destroyed = true;
        for (ClientHandle h : clientHandles) {
            h.disconnect();
        }
        clientHandles.clear();
        if (acceptTask != null) {
            acceptTask.cancel();
            acceptTask = null;
        }
        for (ServerSocketChannel ssc : listenSockets) {
            try {
                ssc.close();
            } catch (IOException e) {
                logger.error("closing listen socket failed", e);
            }
        }
    }

    /**
     * Test if we should destroy outselves.
     */
    private void testForSoftShutdown() {
        // do this so we don't have many recursive calls to testForSoftShutdown when shutting down
        if (destroyed) {
            return;
        }
        if (inSoftShutdown) {
            System.out.println(""+clientHandles.size());
            boolean done = true;
            for (ClientHandle clientHandle : this.clientHandles) {
                if (!clientHandle.isMonitor) {
                    done = false;
                }
            }
            if (done) {
                destroy();
            }
        }
    }

    @Override
    protected void finalize() throws Throwable {
        super.finalize();
        if (!destroyed) {
            logger.warn("Server instance not destroyed, but finalizer called");
        }
        destroy();
    }
}