1 files changed, 48 insertions, 9 deletions

diff --git a/adb/sockets.c b/adb/sockets.c
index f17608b..7f54ad9 100644
--- a/adb/sockets.c
+++ b/adb/sockets.c
@@ -59,17 +59,22 @@ static asocket local_socket_closing_list = {
     .prev = &local_socket_closing_list,
 };
 
-asocket *find_local_socket(unsigned id)
+// Parse the global list of sockets to find one with id |local_id|.
+// If |peer_id| is not 0, also check that it is connected to a peer
+// with id |peer_id|. Returns an asocket handle on success, NULL on failure.
+asocket *find_local_socket(unsigned local_id, unsigned peer_id)
 {
     asocket *s;
     asocket *result = NULL;
 
     adb_mutex_lock(&socket_list_lock);
     for (s = local_socket_list.next; s != &local_socket_list; s = s->next) {
-        if (s->id == id) {
+        if (s->id != local_id)
+            continue;
+        if (peer_id == 0 || (s->peer && s->peer->id == peer_id)) {
             result = s;
-            break;
         }
+        break;
     }
     adb_mutex_unlock(&socket_list_lock);
 
@@ -91,6 +96,11 @@ void install_local_socket(asocket *s)
     adb_mutex_lock(&socket_list_lock);
 
     s->id = local_socket_next_id++;
+
+    // Socket ids should never be 0.
+    if (local_socket_next_id == 0)
+      local_socket_next_id = 1;
+
     insert_local_socket(s, &local_socket_list);
 
     adb_mutex_unlock(&socket_list_lock);
@@ -230,6 +240,12 @@ static void local_socket_close_locked(asocket *s)
     if(s->peer) {
         D("LS(%d): closing peer. peer->id=%d peer->fd=%d\n",
           s->id, s->peer->id, s->peer->fd);
+        /* Note: it's important to call shutdown before disconnecting from
+         * the peer, this ensures that remote sockets can still get the id
+         * of the local socket they're connected to, to send a CLOSE()
+         * protocol event. */
+        if (s->peer->shutdown)
+          s->peer->shutdown(s->peer);
         s->peer->peer = 0;
         // tweak to avoid deadlock
         if (s->peer->close == local_socket_close) {
@@ -397,6 +413,7 @@ asocket *create_local_socket(int fd)
     s->fd = fd;
     s->enqueue = local_socket_enqueue;
     s->ready = local_socket_ready;
+    s->shutdown = NULL;
     s->close = local_socket_close;
     install_local_socket(s);
 
@@ -485,21 +502,29 @@ static void remote_socket_ready(asocket *s)
     send_packet(p, s->transport);
 }
 
-static void remote_socket_close(asocket *s)
+static void remote_socket_shutdown(asocket *s)
 {
-    D("entered remote_socket_close RS(%d) CLOSE fd=%d peer->fd=%d\n",
+    D("entered remote_socket_shutdown RS(%d) CLOSE fd=%d peer->fd=%d\n",
       s->id, s->fd, s->peer?s->peer->fd:-1);
     apacket *p = get_apacket();
     p->msg.command = A_CLSE;
     if(s->peer) {
         p->msg.arg0 = s->peer->id;
+    }
+    p->msg.arg1 = s->id;
+    send_packet(p, s->transport);
+}
+
+static void remote_socket_close(asocket *s)
+{
+    if (s->peer) {
         s->peer->peer = 0;
         D("RS(%d) peer->close()ing peer->id=%d peer->fd=%d\n",
           s->id, s->peer->id, s->peer->fd);
         s->peer->close(s->peer);
     }
-    p->msg.arg1 = s->id;
-    send_packet(p, s->transport);
+    D("entered remote_socket_close RS(%d) CLOSE fd=%d peer->fd=%d\n",
+      s->id, s->fd, s->peer?s->peer->fd:-1);
     D("RS(%d): closed\n", s->id);
     remove_transport_disconnect( s->transport, &((aremotesocket*)s)->disconnect );
     free(s);
@@ -519,15 +544,24 @@ static void remote_socket_disconnect(void*  _s, atransport*  t)
     free(s);
 }
 
+/* Create an asocket to exchange packets with a remote service through transport
+  |t|. Where |id| is the socket id of the corresponding service on the other
+   side of the transport (it is allocated by the remote side and _cannot_ be 0).
+   Returns a new non-NULL asocket handle. */
 asocket *create_remote_socket(unsigned id, atransport *t)
 {
-    asocket *s = calloc(1, sizeof(aremotesocket));
-    adisconnect*  dis = &((aremotesocket*)s)->disconnect;
+    asocket* s;
+    adisconnect* dis;
+
+    if (id == 0) fatal("invalid remote socket id (0)");
+    s = calloc(1, sizeof(aremotesocket));
+    dis = &((aremotesocket*)s)->disconnect;
 
     if (s == NULL) fatal("cannot allocate socket");
     s->id = id;
     s->enqueue = remote_socket_enqueue;
     s->ready = remote_socket_ready;
+    s->shutdown = remote_socket_shutdown;
     s->close = remote_socket_close;
     s->transport = t;
 
@@ -562,6 +596,7 @@ void connect_to_remote(asocket *s, const char *destination)
 static void local_socket_ready_notify(asocket *s)
 {
     s->ready = local_socket_ready;
+    s->shutdown = NULL;
     s->close = local_socket_close;
     adb_write(s->fd, "OKAY", 4);
     s->ready(s);
@@ -573,6 +608,7 @@ static void local_socket_ready_notify(asocket *s)
 static void local_socket_close_notify(asocket *s)
 {
     s->ready = local_socket_ready;
+    s->shutdown = NULL;
     s->close = local_socket_close;
     sendfailmsg(s->fd, "closed");
     s->close(s);
@@ -767,6 +803,7 @@ static int smart_socket_enqueue(asocket *s, apacket *p)
         adb_write(s->peer->fd, "OKAY", 4);
 
         s->peer->ready = local_socket_ready;
+        s->peer->shutdown = NULL;
         s->peer->close = local_socket_close;
         s->peer->peer = s2;
         s2->peer = s->peer;
@@ -806,6 +843,7 @@ static int smart_socket_enqueue(asocket *s, apacket *p)
         ** tear down
         */
     s->peer->ready = local_socket_ready_notify;
+    s->peer->shutdown = NULL;
     s->peer->close = local_socket_close_notify;
     s->peer->peer = 0;
         /* give him our transport and upref it */
@@ -851,6 +889,7 @@ static asocket *create_smart_socket(void)
     if (s == NULL) fatal("cannot allocate socket");
     s->enqueue = smart_socket_enqueue;
     s->ready = smart_socket_ready;
+    s->shutdown = NULL;
     s->close = smart_socket_close;
 
     D("SS(%d)\n", s->id);