1 files changed, 294 insertions, 71 deletions

diff --git a/src/org/gnunet/util/Server.java b/src/org/gnunet/util/Server.java
index f2db539..3d65754 100644
--- a/src/org/gnunet/util/Server.java
+++ b/src/org/gnunet/util/Server.java
@@ -20,6 +20,7 @@
 
 package org.gnunet.util;
 
+import org.gnunet.construct.Construct;
 import org.grothoff.Runabout;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -29,47 +30,138 @@ import java.net.SocketAddress;
 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;
-    private List<ServerSocketChannel> listenSockets;
-    private List<ClientHandle> clients = new LinkedList<ClientHandle>();
 
-    private MessageRunabout receivedMessagehandler;
+    /**
+     * 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>();
-    private ArrayList<Class> expectedMessages;
 
-    private boolean shutdownRequested;
+    /**
+     * 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 {
-        private RelativeTime clientTimeout;
+        /**
+         * The underlying connection to the client-
+         */
         private Connection connection;
 
+        /**
+         * When do we disconnect the client after it has been idle?
+         */
+        private RelativeTime clientTimeout;
+
+        /**
+         * When referenceCount==0, the server is allowed to drop the client after a timeout.
+         */
         private int referenceCount = 0;
-        private Connection.ReceiveHandle currentReceive;
 
+        /**
+         * Handle for canceling the receive process of this client, null if no receive is currently going on.
+         */
+        private Cancelable currentReceive;
+
+        /**
+         * Set to true if the connection to this client should not prevent the server from shutting down.
+         */
         private boolean isMonitor;
 
-        private ClientHandle(SocketChannel accept) {
-            connection = new Connection(accept);
+        /**
+         * 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);
             clientTimeout = idleTimeout;
             // start receiving
             receiveDone(true);
         }
 
-
         /**
          * Notify us when the server has enough space to transmit
          * a message of the given size to the given client.
@@ -81,36 +173,63 @@ public class Server {
          * @return a handle to cancel the notification
          */
         public Cancelable notifyTransmitReady(int size, RelativeTime timeout, MessageTransmitter transmitter) {
-            return connection.notifyTransmitReady(0, timeout, 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 cancel 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 request.  This function must be called from within each
+         * current request. 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 keepClient false if connection to the client should be closed
+         * @param stayConnected false if connection to the client should be closed
          */
-        public void receiveDone(boolean keepClient) {
-            if (keepClient) {
+        public void receiveDone(boolean stayConnected) {
+            if (stayConnected) {
                 currentReceive = connection.receive(RelativeTime.FOREVER, new MessageReceiver() {
                     @Override
                     public void process(GnunetMessage.Body msg) {
-                        if (msg instanceof UnknownMessageBody) {
+                        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) {
+                        if (receivedMessageHandler == null) {
                             throw new AssertionError("received message, but no handler installed");
                         }
-                        receivedMessagehandler.setSender(ClientHandle.this);
-                        receivedMessagehandler.visitAppropriate(msg);
+                        receivedMessageHandler.setSender(ClientHandle.this);
+                        receivedMessageHandler.visitAppropriate(msg);
                     }
 
                     @Override
@@ -120,7 +239,12 @@ public class Server {
                     }
                 });
             } else {
-                disconnect();
+                if (referenceCount > 0) {
+                    this.disconnectRequested = true;
+                } else {
+                    System.out.println("disconnecting " + this.isMonitor);
+                    disconnect();
+                }
             }
 
         }
@@ -129,48 +253,86 @@ public class Server {
          * Change the idle timeout of this particular client.
          */
         public void setTimeout(RelativeTime newTimeout) {
+            this.clientTimeout = newTimeout;
         }
 
+        /**
+         * 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();
-            Server.this.clients.remove(this);
+            // 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();
         }
 
         /**
-         * Disable the warning the server issues if a message is not acknowledged
-         * in a timely fashion.  Use this call if a client is intentionally delayed
-         * for a while.  Only applies to the current message.
+         * Prevent the client from being disconnected.
+         * For every keep, there should be an additional matching drop.
          */
-        public void disableReceiveDoneWarning() {
-            // todo
-        }
-
         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 && shutdownRequested) {
+            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;
+        }
     }
 
 
-    abstract static class MessageRunabout extends Runabout {
+    /**
+     * 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
          */
@@ -178,39 +340,16 @@ public class Server {
             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;
         }
     }
 
-
-    private void doAccept(final ServerSocketChannel srv) {
-        Scheduler.TaskConfiguration b = new Scheduler.TaskConfiguration(RelativeTime.FOREVER,
-                new Scheduler.Task() {
-                    @Override
-                    public void run(Scheduler.RunContext ctx) {
-                        acceptTask = null;
-                        try {
-                            SocketChannel cli = srv.accept();
-
-                            if (cli != null) {
-                                logger.debug("client connected");
-                                cli.configureBlocking(false);
-                                ClientHandle clientHandle = new ClientHandle(cli);
-                                clients.add(clientHandle);
-                            }
-
-                        } catch (IOException e) {
-                            throw new RuntimeException("accept failed", e);
-                        }
-                        doAccept(srv);
-                    }
-                });
-        b.selectAccept(srv);
-        acceptTask = b.schedule();
-    }
-
-
     /**
      * Create a server listening on all specified addresses.
      *
@@ -221,7 +360,6 @@ public class Server {
     public Server(List<SocketAddress> addresses, RelativeTime idleTimeout, boolean requireFound) {
         this.idleTimeout = idleTimeout;
         this.requireFound = requireFound;
-        listenSockets = new ArrayList<ServerSocketChannel>(addresses.size());
         try {
             for (SocketAddress addr : addresses) {
                 ServerSocketChannel socket = ServerSocketChannel.open();
@@ -229,20 +367,53 @@ public class Server {
                 socket.socket().bind(addr);
                 logger.debug("socket listening on {}", addr.toString());
                 listenSockets.add(socket);
-                doAccept(socket);
+                addAcceptSocket(socket);
             }
         } catch (IOException e) {
-            throw new RuntimeException("could not bind");
+            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;
     }
 
-    public void addAcceptSocket(ServerSocketChannel sock) {
-        doAccept(sock);
+    /**
+     * 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.selectAccept(sock);
+        acceptTask = b.schedule();
     }
 
     /**
@@ -253,10 +424,18 @@ public class Server {
      * @param msgRunabout handler
      */
     public void setHandler(MessageRunabout msgRunabout) {
-        receivedMessagehandler = 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() {
@@ -268,19 +447,63 @@ public class Server {
     }
 
     /**
-     * Stop the listen socket and get ready to shutdown the server
-     * once only 'monitor' clients are left.
+     * Stop the listen socket destroy the server as soon as only monitor clients are left.
      */
     public void stopListening() {
-        shutdownRequested = true;
-        // todo: shut down if only monitor clients left
+        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() {
-        for (ClientHandle h : new ArrayList<ClientHandle>(clients)) {
+        if (destroyed) {
+            return;
+        }
+        destroyed = true;
+        for (ClientHandle h : clientHandles) {
             h.disconnect();
         }
-        acceptTask.cancel();
+        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();
+            }
+        }
+    }
 }