2 files changed, 92 insertions, 7 deletions

diff --git a/src/stream/stream_api.c b/src/stream/stream_api.c
index 92c1093f7..535850de2 100644
--- a/src/stream/stream_api.c
+++ b/src/stream/stream_api.c
@@ -232,6 +232,11 @@ struct GNUNET_STREAM_Socket
   struct GNUNET_STREAM_IOReadHandle *read_handle;
 
   /**
+   * The shutdown handle associated with this socket
+   */
+  struct GNUNET_STREAM_ShutdownHandle *shutdown_handle;
+
+  /**
    * Buffer for storing received messages
    */
   void *receive_buffer;
@@ -443,6 +448,16 @@ struct GNUNET_STREAM_ShutdownHandle
    * Which operation to shutdown? SHUT_RD, SHUT_WR or SHUT_RDWR
    */
   int operation;
+
+  /**
+   * Shutdown completion callback
+   */
+  GNUNET_STREAM_ShutdownCompletion completion_cb;
+
+  /**
+   * Closure for completion callback
+   */
+  void *completion_cls;
 };
 
 
@@ -1559,6 +1574,10 @@ handle_close (struct GNUNET_STREAM_Socket *socket,
 {
   struct GNUNET_STREAM_MessageHeader *close_ack;
 
+  GNUNET_log (GNUNET_ERROR_TYPE_DEBUG,
+              "%x: Received CLOSE from %x\n",
+              socket->our_id,
+              socket->other_peer);
   close_ack = GNUNET_malloc (sizeof (struct GNUNET_STREAM_MessageHeader));
   close_ack->header.size = htons (sizeof (struct GNUNET_STREAM_MessageHeader));
   close_ack->header.type = htons (GNUNET_MESSAGE_TYPE_STREAM_CLOSE_ACK);
@@ -1625,10 +1644,30 @@ handle_close_ack (struct GNUNET_STREAM_Socket *socket,
                   const struct GNUNET_STREAM_MessageHeader *message,
                   const struct GNUNET_ATS_Information*atsi)
 {
+  struct GNUNET_STREAM_ShutdownHandle *shutdown_handle;
+
+  shutdown_handle = socket->shutdown_handle;
   switch (socket->state)
     {
     case STATE_CLOSE_WAIT:
       socket->state = STATE_CLOSED;
+      if ( (NULL == shutdown_handle) ||
+           (SHUT_RDWR != shutdown_handle->operation) )
+        {
+          GNUNET_log (GNUNET_ERROR_TYPE_DEBUG,
+                      "%x: Received CLOSE_ACK when shutdown handle is NULL or "
+                      "not for SHUT_RDWR\n",
+                      socket->our_id);
+          return GNUNET_OK;
+        }
+      GNUNET_log (GNUNET_ERROR_TYPE_DEBUG,
+                  "%x: Received CLOSE_ACK from %x\n",
+                  socket->our_id,
+                  socket->other_peer);
+      if (NULL != shutdown_handle->completion_cb) /* Shutdown completion */
+        shutdown_handle->completion_cb(shutdown_handle->completion_cls,
+                                       SHUT_RDWR);
+      GNUNET_free (shutdown_handle); /* Free shutdown handle */
       break;
     default:
       GNUNET_log (GNUNET_ERROR_TYPE_DEBUG,
@@ -2515,16 +2554,25 @@ GNUNET_STREAM_shutdown (struct GNUNET_STREAM_Socket *socket,
 {
   struct GNUNET_STREAM_ShutdownHandle *handle;
   struct GNUNET_STREAM_MessageHeader *msg;
+  
+  GNUNET_assert (NULL == socket->shutdown_handle);
 
   handle = GNUNET_malloc (sizeof (struct GNUNET_STREAM_ShutdownHandle));
   handle->socket = socket;
+  handle->completion_cb = completion_cb;
+  handle->completion_cls = completion_cls;
+  socket->shutdown_handle = handle;
+
   msg = GNUNET_malloc (sizeof (struct GNUNET_STREAM_MessageHeader));
   msg->header.size = htons (sizeof (struct GNUNET_STREAM_MessageHeader));
   switch (operation)
     {
     case SHUT_RD:
       handle->operation = SHUT_RD;
-
+      if (NULL != socket->read_handle)
+        GNUNET_log (GNUNET_ERROR_TYPE_WARNING,
+                    "Existing read handle should be cancelled before shutting"
+                    " down reading\n");
       break;
     case SHUT_WR:
       handle->operation = SHUT_WR;
diff --git a/src/stream/test_stream_local.c b/src/stream/test_stream_local.c
index c3fcc6492..125873acf 100644
--- a/src/stream/test_stream_local.c
+++ b/src/stream/test_stream_local.c
@@ -60,6 +60,11 @@ struct PeerData
   struct GNUNET_STREAM_IOReadHandle *io_read_handle;
 
   /**
+   * Peer's shutdown handle
+   */
+  struct GNUNET_STREAM_ShutdownHandle *shutdown_handle;
+
+  /**
    * Our Peer id
    */
   struct GNUNET_PeerIdentity our_id;
@@ -109,7 +114,7 @@ static int reading_success;
  * Check whether peers successfully shut down.
  */
 static void
-shutdown_callback (void *cls, const char *emsg)
+peergroup_shutdown_callback (void *cls, const char *emsg)
 {
   if (emsg != NULL)
   {
@@ -126,10 +131,10 @@ shutdown_callback (void *cls, const char *emsg)
 
 
 /**
- * Shutdown nicely
+ * Close sockets and stop testing deamons nicely
  */
 static void
-do_shutdown (void *cls, const struct GNUNET_SCHEDULER_TaskContext *tc)
+do_close (void *cls, const struct GNUNET_SCHEDULER_TaskContext *tc)
 {
   if (NULL != peer1.socket)
     GNUNET_STREAM_close (peer1.socket);
@@ -149,12 +154,44 @@ do_shutdown (void *cls, const struct GNUNET_SCHEDULER_TaskContext *tc)
   GNUNET_TESTING_daemons_stop (pg,
                                GNUNET_TIME_relative_multiply
                                (GNUNET_TIME_UNIT_SECONDS, 5),
-                               &shutdown_callback,
+                               &peergroup_shutdown_callback,
                                NULL);
 }
 
 
 /**
+ * Completion callback for shutdown
+ *
+ * @param cls the closure from GNUNET_STREAM_shutdown call
+ * @param operation the operation that was shutdown (SHUT_RD, SHUT_WR,
+ *          SHUT_RDWR) 
+ */
+static void 
+shutdown_completion (void *cls,
+                     int operation)
+{
+  GNUNET_log (GNUNET_ERROR_TYPE_DEBUG,
+              "STREAM shutdown successful\n");
+  GNUNET_SCHEDULER_add_now (&do_close,
+                            cls);
+}
+
+
+
+/**
+ * Shutdown sockets gracefully
+ */
+static void
+do_shutdown (void *cls, const struct GNUNET_SCHEDULER_TaskContext *tc)
+{
+  peer1.shutdown_handle = GNUNET_STREAM_shutdown (peer1.socket, 
+                                                  SHUT_RDWR,
+                                                  &shutdown_completion,
+                                                  cls);
+}
+
+
+/**
  * Something went wrong and timed out. Kill everything and set error flag
  */
 static void
@@ -167,7 +204,7 @@ do_abort (void *cls, const struct GNUNET_SCHEDULER_TaskContext *tc)
     }
   result = GNUNET_SYSERR;
   abort_task = 0;
-  do_shutdown (cls, tc);
+  do_close (cls, tc);  
 }
 
 /**
@@ -237,7 +274,7 @@ write_completion (void *cls,
       else
         {
           writing_success = GNUNET_YES;
-          if (GNUNET_YES == reading_success) 
+          if (GNUNET_YES == reading_success)
             GNUNET_SCHEDULER_add_now (&do_shutdown, NULL);
         }
     }