Merge pull request #76 from AlbinTheander/albin_theander

Closes all existing connections at shutdown.

2 files changed, 86 insertions, 0 deletions

diff --git a/core/src/main/java/fi/iki/elonen/NanoHTTPD.java b/core/src/main/java/fi/iki/elonen/NanoHTTPD.java
index 5ce08e7..424856b 100644
--- a/core/src/main/java/fi/iki/elonen/NanoHTTPD.java
+++ b/core/src/main/java/fi/iki/elonen/NanoHTTPD.java
@@ -73,6 +73,7 @@ public abstract class NanoHTTPD {
     private final String hostname;
     private final int myPort;
     private ServerSocket myServerSocket;
+    private Set<Socket> openConnections = new HashSet<Socket>();
     private Thread myThread;
     /**
      * Pluggable strategy for asynchronously executing requests.
@@ -142,10 +143,12 @@ public abstract class NanoHTTPD {
                 do {
                     try {
                         final Socket finalAccept = myServerSocket.accept();
+                        registerConnection(finalAccept);
                         finalAccept.setSoTimeout(SOCKET_READ_TIMEOUT);
                         final InputStream inputStream = finalAccept.getInputStream();
                         if (inputStream == null) {
                             safeClose(finalAccept);
+                            unRegisterConnection(finalAccept);
                         } else {
                             asyncRunner.exec(new Runnable() {
                                 @Override
@@ -168,6 +171,7 @@ public abstract class NanoHTTPD {
                                         safeClose(outputStream);
                                         safeClose(inputStream);
                                         safeClose(finalAccept);
+                                        unRegisterConnection(finalAccept);
                                     }
                                 }
                             });
@@ -188,12 +192,41 @@ public abstract class NanoHTTPD {
     public void stop() {
         try {
             safeClose(myServerSocket);
+            closeAllConnections();
             myThread.join();
         } catch (Exception e) {
             e.printStackTrace();
         }
     }
 
+    /**
+     * Registers that a new connection has been set up.
+     *
+     * @param socket
+     *            the {@link Socket} for the connection.
+     */
+    public synchronized void registerConnection(Socket socket) {
+        openConnections.add(socket);
+    }
+
+    /**
+     * Registers that a connection has been closed
+     *
+     * @param socket
+     *            the {@link Socket} for the connection.
+     */
+    public synchronized void unRegisterConnection(Socket socket) {
+        openConnections.remove(socket);
+    }
+
+    /**
+     * Forcibly closes all connections that are open.
+     */
+    public synchronized void closeAllConnections() {
+        for (Socket socket : openConnections) {
+            safeClose(socket);
+        }
+    }
 
     public final int getListeningPort() {
         return myServerSocket == null ? -1 : myServerSocket.getLocalPort();
diff --git a/core/src/test/java/fi/iki/elonen/integration/ShutdownTest.java b/core/src/test/java/fi/iki/elonen/integration/ShutdownTest.java
new file mode 100644
index 0000000..0fcb275
--- /dev/null
+++ b/core/src/test/java/fi/iki/elonen/integration/ShutdownTest.java
@@ -0,0 +1,53 @@
+package fi.iki.elonen.integration;
+
+import static org.junit.Assert.*;
+import fi.iki.elonen.NanoHTTPD;
+
+import org.junit.Test;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.HttpURLConnection;
+import java.net.MalformedURLException;
+import java.net.URL;
+
+public class ShutdownTest {
+
+    @Test
+    public void connectionsAreClosedWhenServerStops() throws IOException {
+        TestServer server = new TestServer();
+        server.start();
+        makeRequest();
+        server.stop();
+        try {
+            makeRequest();
+            fail("Connection should be closed!");
+        } catch (IOException e) {
+            // Expected exception
+        }
+    }
+
+    private void makeRequest() throws MalformedURLException, IOException {
+        HttpURLConnection connection = (HttpURLConnection) new URL("http://localhost:8092/").openConnection();
+        // Keep-alive seems to be on by default, but just in case that changes.
+        connection.addRequestProperty("Connection", "keep-alive");
+        InputStream in = connection.getInputStream();
+        while (in.available() > 0) {
+            in.read();
+        }
+        in.close();
+    }
+
+    private class TestServer extends NanoHTTPD {
+
+        public TestServer() {
+            super(8092);
+        }
+
+        @Override
+        public Response serve(IHTTPSession session) {
+            return new Response("Whatever");
+        }
+    }
+
+}