From e0c2d59ab5a4b6b12782f97fd9698ac551f10e2d Mon Sep 17 00:00:00 2001
From: Jeremy Lakeman <Jeremy.Lakeman@gmail.com>
Date: Fri, 10 Jan 2014 15:35:22 +1030
Subject: [PATCH] Shutdown local tcp connection when remote msp connection has
 shutdown

---
 msp_client.c |  5 +++--
 msp_proxy.c  | 38 ++++++++++++++++++++++++++------------
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/msp_client.c b/msp_client.c
index b683b347..1d1ab13b 100644
--- a/msp_client.c
+++ b/msp_client.c
@@ -421,6 +421,7 @@ int msp_send(struct msp_sock *sock, const uint8_t *payload, size_t len)
 
 int msp_shutdown(struct msp_sock *sock)
 {
+  assert(!(sock->state&MSP_STATE_SHUTDOWN_LOCAL));
   if (sock->tx._tail && sock->tx._tail->sent==0){
     sock->tx._tail->flags |= FLAG_SHUTDOWN;
   }else{
@@ -459,11 +460,11 @@ static int process_sock(struct msp_sock *sock)
   while(p && p->seq == sock->rx.next_seq){
     struct msp_packet *packet=p;
     
-    // process packet flags when we have delivered the packet
+    // process packet flags when we are about to deliver the last packet
     if (packet->flags & FLAG_SHUTDOWN)
       sock->state|=MSP_STATE_SHUTDOWN_REMOTE;
     
-    if (sock->handler && packet->payload){
+    if (sock->handler){
       int r = sock->handler(sock, sock->state, packet->payload, packet->len, sock->context);
       if (r==-1){
 	sock->state |= MSP_STATE_CLOSED;
diff --git a/msp_proxy.c b/msp_proxy.c
index a138a7be..5d53370a 100644
--- a/msp_proxy.c
+++ b/msp_proxy.c
@@ -141,11 +141,26 @@ static void process_msp_asap()
   schedule(&mdp_sock);
 }
 
+static void remote_shutdown(struct connection *conn)
+{
+  struct mdp_sockaddr remote;
+  if (shutdown(conn->alarm_out.poll.fd, SHUT_WR))
+    WARNF_perror("shutdown(%d)", conn->alarm_out.poll.fd);
+  msp_get_remote_adr(conn->sock, &remote);
+  INFOF(" - Connection with %s:%d remote shutdown", alloca_tohex_sid_t(remote.sid), remote.port);
+}
+
+static void local_shutdown(struct connection *conn)
+{
+  struct mdp_sockaddr remote;
+  msp_get_remote_adr(conn->sock, &remote);
+  msp_shutdown(conn->sock);
+  INFOF(" - Connection with %s:%d local shutdown", alloca_tohex_sid_t(remote.sid), remote.port);
+}
+
 static int msp_handler(struct msp_sock *sock, msp_state_t state, const uint8_t *payload, size_t len, void *context)
 {
   struct connection *conn = context;
-  conn->last_state=state;
-  
   if (state & MSP_STATE_ERROR)
     saw_error=1;
     
@@ -168,6 +183,11 @@ static int msp_handler(struct msp_sock *sock, msp_state_t state, const uint8_t *
     conn->alarm_out.function(&conn->alarm_out);
   }
   
+  if ((state & MSP_STATE_SHUTDOWN_REMOTE) && !(conn->last_state & MSP_STATE_SHUTDOWN_REMOTE) && conn->out->limit==0)
+    remote_shutdown(conn);
+  
+  conn->last_state=state;
+  
   if (state & MSP_STATE_CLOSED){
     struct mdp_sockaddr remote;
     msp_get_remote_adr(sock, &remote);
@@ -256,14 +276,6 @@ static void msp_poll(struct sched_ent *alarm)
   }
 }
 
-static void local_shutdown(struct connection *conn)
-{
-  struct mdp_sockaddr remote;
-  msp_get_remote_adr(conn->sock, &remote);
-  msp_shutdown(conn->sock);
-  INFOF(" - Connection with %s:%d local shutdown", alloca_tohex_sid_t(remote.sid), remote.port);
-}
-
 static int try_send(struct connection *conn)
 {
   if (!conn->in->limit)
@@ -274,7 +286,7 @@ static int try_send(struct connection *conn)
   // if this packet was acceptted, clear the read buffer
   conn->in->limit = conn->in->position = 0;
   // hit end of data?
-  if (conn->alarm_in.poll.events==0){
+  if (!(conn->alarm_in.poll.events&POLLIN)){
     local_shutdown(conn);
   }else{
     if (!is_watching(&conn->alarm_in))
@@ -324,6 +336,8 @@ static void io_poll(struct sched_ent *alarm)
       conn->out->limit = conn->out->position = 0;
       if (is_watching(alarm))
 	unwatch(alarm);
+      if (conn->last_state & MSP_STATE_SHUTDOWN_REMOTE)
+	remote_shutdown(conn);
     }
     
     if (conn->out->limit < conn->out->capacity){
@@ -338,7 +352,7 @@ static void io_poll(struct sched_ent *alarm)
   if (alarm->poll.revents & POLLHUP) {
     // EOF? trigger a graceful shutdown
     unwatch(alarm);
-    alarm->poll.events = 0;
+    alarm->poll.events &= ~POLLIN;
     if (!conn->in->limit){
       local_shutdown(conn);
       process_msp_asap();