From d830e099a6d040b6f03a46cc22e67f959f528043 Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Mon, 11 Feb 2013 02:05:11 +0900
Subject: [PATCH] shrpx: Send pending response data before RST_STREAM in tunnel
 connection

---
 src/shrpx_https_upstream.cc | 14 ++++++++++++++
 src/shrpx_spdy_session.cc   | 19 ++++++++++++++++---
 src/shrpx_spdy_upstream.cc  | 28 +++++++++++++++++++---------
 3 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/src/shrpx_https_upstream.cc b/src/shrpx_https_upstream.cc
index 27e4a52e..415dd9f6 100644
--- a/src/shrpx_https_upstream.cc
+++ b/src/shrpx_https_upstream.cc
@@ -378,6 +378,20 @@ void https_downstream_readcb(bufferevent *bev, void *ptr)
     delete upstream->get_client_handler();
   } else if(rv == 0) {
     if(downstream->get_response_state() == Downstream::MSG_COMPLETE) {
+      if(get_config()->client_mode && downstream->tunnel_established()) {
+        // For tunneled connection, if there is no pending data,
+        // delete handler because on_write will not be called.
+        ClientHandler *handler = upstream->get_client_handler();
+        if(handler->get_pending_write_length() == 0) {
+          delete handler;
+        } else {
+          if(LOG_ENABLED(INFO)) {
+            DLOG(INFO, downstream) << "Tunneled connection has pending data";
+          }
+          handler->set_should_close_after_write(true);
+        }
+        return;
+      }
       if(downstream->get_response_connection_close()) {
         // Connection close
         downstream->set_downstream_connection(0);
diff --git a/src/shrpx_spdy_session.cc b/src/shrpx_spdy_session.cc
index b876889d..4cc8d3c3 100644
--- a/src/shrpx_spdy_session.cc
+++ b/src/shrpx_spdy_session.cc
@@ -697,9 +697,22 @@ void on_ctrl_recv_callback
       if(downstream &&
          downstream->get_downstream_stream_id() ==
          frame->rst_stream.stream_id) {
-        // If we got RST_STREAM, just flag MSG_RESET to indicate
-        // upstream connection must be terminated.
-        downstream->set_response_state(Downstream::MSG_RESET);
+        if(downstream->tunnel_established() &&
+           downstream->get_response_state() == Downstream::HEADER_COMPLETE) {
+          // For tunneled connection, we has to submit RST_STREAM to
+          // upstream *after* whole response body is sent. We just set
+          // MSG_COMPLETE here. Upstream will take care of that.
+          if(LOG_ENABLED(INFO)) {
+            SSLOG(INFO, spdy) << "RST_STREAM against tunneled stream "
+                              << "stream_id="
+                              << frame->rst_stream.stream_id;
+          }
+          downstream->set_response_state(Downstream::MSG_COMPLETE);
+        } else {
+          // If we got RST_STREAM, just flag MSG_RESET to indicate
+          // upstream connection must be terminated.
+          downstream->set_response_state(Downstream::MSG_RESET);
+        }
         call_downstream_readcb(spdy, downstream);
       }
     }
diff --git a/src/shrpx_spdy_upstream.cc b/src/shrpx_spdy_upstream.cc
index 6afcf820..51a482e9 100644
--- a/src/shrpx_spdy_upstream.cc
+++ b/src/shrpx_spdy_upstream.cc
@@ -543,11 +543,12 @@ void spdy_downstream_eventcb(bufferevent *bev, short events, void *ptr)
           ULOG(INFO, upstream) << "Downstream body was ended by EOF";
         }
         downstream->set_response_state(Downstream::MSG_COMPLETE);
-        if(downstream->tunnel_established()) {
-          upstream->rst_stream(downstream, SPDYLAY_INTERNAL_ERROR);
-        } else {
-          upstream->on_downstream_body_complete(downstream);
-        }
+
+        // For tunneled connection, MSG_COMPLETE signals
+        // spdy_data_read_callback to send RST_STREAM after pending
+        // response body is sent. This is needed to ensure that
+        // RST_STREAM is sent after all pending data are sent.
+        upstream->on_downstream_body_complete(downstream);
       } else if(downstream->get_response_state() == Downstream::MSG_COMPLETE) {
         // For SSL tunneling?
         upstream->rst_stream(downstream, SPDYLAY_INTERNAL_ERROR);
@@ -661,11 +662,20 @@ ssize_t spdy_data_read_callback(spdylay_session *session,
   evbuffer *body = downstream->get_response_body_buf();
   assert(body);
   int nread = evbuffer_remove(body, buf, length);
-  // For tunneling, DATA stream is endless
-  if(!downstream->tunnel_established() &&
-     nread == 0 &&
+  if(nread == 0 &&
      downstream->get_response_state() == Downstream::MSG_COMPLETE) {
-    *eof = 1;
+    if(!downstream->tunnel_established()) {
+      *eof = 1;
+    } else {
+      // For tunneling, issue RST_STREAM to finish the stream.
+      SpdyUpstream *upstream;
+      upstream = reinterpret_cast<SpdyUpstream*>(downstream->get_upstream());
+      if(LOG_ENABLED(INFO)) {
+        ULOG(INFO, upstream) << "RST_STREAM to tunneled stream stream_id="
+                             << stream_id;
+      }
+      upstream->rst_stream(downstream, SPDYLAY_CANCEL);
+    }
   }
   if(nread == 0 && *eof != 1) {
     return SPDYLAY_ERR_DEFERRED;