3 files changed, 24 insertions, 10 deletions

diff --git a/src/or/connection.c b/src/or/connection.c
index c7f3867d0..75804ee02 100644
--- a/src/or/connection.c
+++ b/src/or/connection.c
@@ -199,8 +199,14 @@ _connection_mark_for_close(connection_t *conn, char reason)
       }
       /* No special processing needed. */
       break;
-    case CONN_TYPE_EXIT:
     case CONN_TYPE_AP:
+      if (conn->socks_request->has_finished == 0) {
+        log_fn(LOG_INFO,"Cleaning up AP -- sending socks reject.");
+        connection_ap_handshake_socks_reply(conn, NULL, 0, 0);
+        conn->socks_request->has_finished = 1;
+      }
+      /* fall through, to do things for both ap and exit */
+    case CONN_TYPE_EXIT:
       if (conn->state == EXIT_CONN_STATE_RESOLVING)
         connection_dns_remove(conn);
       if (!conn->has_sent_end && reason &&
@@ -1004,13 +1010,16 @@ void assert_connection_ok(connection_t *conn, time_t now)
     assert(!conn->done_sending);
     assert(!conn->done_receiving);
   } else {
-    if(conn->type == CONN_TYPE_AP && conn->state == AP_CONN_STATE_OPEN)
-      assert(conn->cpath_layer);
-    if(conn->cpath_layer)
-      assert_cpath_layer_ok(conn->cpath_layer);
     /* XXX unchecked: package window, deliver window. */
   }
-  if (conn->type != CONN_TYPE_AP) {
+  if (conn->type == CONN_TYPE_AP) {
+    assert(conn->socks_request);
+    if (conn->state == AP_CONN_STATE_OPEN) {
+      assert(conn->socks_request->has_finished);
+      assert(conn->cpath_layer);
+      assert_cpath_layer_ok(conn->cpath_layer);
+    }
+  } else {
     assert(!conn->socks_request);
   }
 
diff --git a/src/or/connection_edge.c b/src/or/connection_edge.c
index 867fe78c9..6665d7cba 100644
--- a/src/or/connection_edge.c
+++ b/src/or/connection_edge.c
@@ -12,8 +12,6 @@ static int connection_ap_handshake_process_socks(connection_t *conn);
 static int connection_ap_handshake_attach_circuit(connection_t *conn);
 static int connection_ap_handshake_attach_circuit_helper(connection_t *conn);
 static void connection_ap_handshake_send_begin(connection_t *ap_conn, circuit_t *circ);
-static void connection_ap_handshake_socks_reply(connection_t *conn, char *reply,
-                                                int replylen, char success);
 
 static int connection_exit_begin_conn(cell_t *cell, circuit_t *circ);
 static void connection_edge_consider_sending_sendme(connection_t *conn);
@@ -284,6 +282,7 @@ int connection_edge_process_relay_cell(cell_t *cell, circuit_t *circ, connection
              (int)(time(NULL) - conn->timestamp_lastread));
       circuit_log_path(LOG_INFO,circ);
       connection_ap_handshake_socks_reply(conn, NULL, 0, 1);
+      conn->socks_request->has_finished = 1;
       return 0;
     } else {
       log_fn(LOG_WARN,"Got an unexpected relay command %d, in state %d (%s). Closing.",
@@ -704,6 +703,8 @@ static int connection_ap_handshake_process_socks(connection_t *conn) {
     } else {
       log_fn(LOG_DEBUG,"socks handshake not all here yet.");
     }
+    if (sockshere == -1)
+      conn->socks_request->has_finished = 1;
     return sockshere;
   } /* else socks handshake is done, continue processing */
 
@@ -848,8 +849,8 @@ static void connection_ap_handshake_send_begin(connection_t *ap_conn, circuit_t
   return;
 }
 
-static void connection_ap_handshake_socks_reply(connection_t *conn, char *reply,
-                                               int replylen, char success) {
+void connection_ap_handshake_socks_reply(connection_t *conn, char *reply,
+                                         int replylen, char success) {
   char buf[256];
 
   if(replylen) { /* we already have a reply in mind */
diff --git a/src/or/or.h b/src/or/or.h
index feb2d0fd8..f19817e6a 100644
--- a/src/or/or.h
+++ b/src/or/or.h
@@ -554,6 +554,7 @@ struct socks_request_t {
   char socks_version;
   int replylen;
   char reply[MAX_SOCKS_REPLY_LEN];
+  int has_finished; /* has the socks handshake finished? */
   char address[MAX_SOCKS_ADDR_LEN];
   uint16_t port;
 };
@@ -750,6 +751,9 @@ int connection_edge_finished_flushing(connection_t *conn);
 
 int connection_edge_package_raw_inbuf(connection_t *conn);
 
+void connection_ap_handshake_socks_reply(connection_t *conn, char *reply,
+                                         int replylen, char success);
+
 void connection_exit_connect(connection_t *conn);
 int connection_ap_can_use_exit(connection_t *conn, routerinfo_t *exit);
 void connection_ap_expire_beginning(void);