From f1e7a76fe7d554e0a3897f01bc2504ef2fd7958f Mon Sep 17 00:00:00 2001
From: Marko Rauhamaa <marko.rauhamaa@f-secure.com>
Date: Thu, 17 Jun 2021 15:36:23 +0300
Subject: [PATCH] Postpone closing plain output until transport is closed.

---
 src/tls_connection.c      | 43 +++++++++++++++++----------------------
 src/tls_openssl.c         |  7 +++++--
 src/tls_securetransport.c |  7 +++++--
 3 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/src/tls_connection.c b/src/tls_connection.c
index 75abb36..1b9bc39 100644
--- a/src/tls_connection.c
+++ b/src/tls_connection.c
@@ -251,8 +251,17 @@ static void encrypted_output_stream_close(void *obj)
     FSTRACE(ASYNCTLS_CONN_ENCRYPTED_OUTPUT_CLOSE, conn->uid);
     assert(!conn->encrypted_output_closed);
     conn->encrypted_output_closed = true;
-    if (conn->state == TLS_CONN_STATE_ZOMBIE)
-        async_wound(conn->async, conn);
+    switch (conn->state) {
+        case TLS_CONN_STATE_SHUT_DOWN_OUTGOING:
+            bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
+            conn->plain_output_stream = drystream;
+            break;
+        case TLS_CONN_STATE_ZOMBIE:
+            async_wound(conn->async, conn);
+            break;
+        default:
+            ;
+    }
 }
 
 FSTRACE_DECL(ASYNCTLS_CONN_ENCRYPTED_OUTPUT_REGISTER,
@@ -473,14 +482,7 @@ FSTRACE_DECL(ASYNCTLS_CONN_CLOSE, "UID=%64u");
 void tls_close(tls_conn_t *conn)
 {
     FSTRACE(ASYNCTLS_CONN_CLOSE, conn->uid);
-    switch (conn->state) {
-        case TLS_CONN_STATE_ZOMBIE:
-            return;
-        case TLS_CONN_STATE_SHUT_DOWN_OUTGOING:
-            break;
-        default:
-            bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
-    }
+    assert(conn->state != TLS_CONN_STATE_ZOMBIE);
     tls_free_underlying_resources(conn);
     bytestream_1_close(conn->encrypted_input_stream);
     if (conn->is_client)
@@ -493,16 +495,13 @@ void tls_close(tls_conn_t *conn)
     tls_set_conn_state(conn, TLS_CONN_STATE_ZOMBIE);
 }
 
+FSTRACE_DECL(ASYNCTLS_CONN_CLOSE_ASYNCTLS_ONLY, "UID=%64u");
+
 void tls_close_asynctls_only(tls_conn_t *conn)
 {
-    switch (conn->state) {
-        case TLS_CONN_STATE_ZOMBIE:
-            return;
-        case TLS_CONN_STATE_SHUT_DOWN_OUTGOING:
-            break;
-        default:
-            bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
-    }
+    FSTRACE(ASYNCTLS_CONN_CLOSE_ASYNCTLS_ONLY, conn->uid);
+    assert(conn->state != TLS_CONN_STATE_ZOMBIE);
+    bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
     fsfree(conn->underlying_tech);
     if (conn->encrypted_output_closed)
         async_wound(conn->async, conn);
@@ -530,12 +529,8 @@ FSTRACE_DECL(ASYNCTLS_CONN_SET_OUTPUT, "UID=%64u OUTPUT=%p")
 void tls_set_plain_output_stream(tls_conn_t *conn, bytestream_1 output_stream)
 {
     FSTRACE(ASYNCTLS_CONN_SET_OUTPUT, conn->uid, output_stream.obj);
-    switch (conn->state) {
-        case TLS_CONN_STATE_ZOMBIE:
-            return;
-        default:
-            bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
-    }
+    assert(conn->state != TLS_CONN_STATE_ZOMBIE);
+    bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
     conn->plain_output_stream = output_stream;
     action_1 plain_output_cb = { conn, (act_1) output_notification };
     bytestream_1_register_callback(output_stream, plain_output_cb);
diff --git a/src/tls_openssl.c b/src/tls_openssl.c
index 509b39d..c0bfc96 100644
--- a/src/tls_openssl.c
+++ b/src/tls_openssl.c
@@ -243,6 +243,7 @@ static int ssl_shutdown(tls_conn_t *conn)
 
 static ssize_t relay_encrypted_output(tls_conn_t *conn, void *buf, size_t count)
 {
+    assert(conn->state == TLS_CONN_STATE_OPEN);
     for (;;) {
         int ret = bio_read(conn, buf, count);
         if (ret > 0)
@@ -276,8 +277,10 @@ static ssize_t relay_encrypted_output(tls_conn_t *conn, void *buf, size_t count)
         buffer_consume(buffer, nn);
     }
     tls_set_conn_state(conn, TLS_CONN_STATE_SHUT_DOWN_OUTGOING);
-    bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
-    conn->plain_output_stream = drystream;
+    if (conn->encrypted_output_closed) {
+        bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
+        conn->plain_output_stream = drystream;
+    }
     int ret = ssl_shutdown(conn);
     if (ret < 0)
         return declare_protocol_error(conn);
diff --git a/src/tls_securetransport.c b/src/tls_securetransport.c
index c4f8ee3..0ec92b0 100644
--- a/src/tls_securetransport.c
+++ b/src/tls_securetransport.c
@@ -160,8 +160,11 @@ static ssize_t relay_encrypted_output(tls_conn_t *conn, void *buf, size_t count)
         if (n == 0) {
             tls_set_conn_state(conn, TLS_CONN_STATE_SHUT_DOWN_OUTGOING);
             FSTRACE(ASYNCTLS_SECTRAN_RELAY_OUTGOING_EXHAUSTED, conn->uid);
-            bytestream_1_close_relaxed(conn->async, conn->plain_output_stream);
-            conn->plain_output_stream = drystream;
+            if (conn->encrypted_output_closed) {
+                bytestream_1_close_relaxed(conn->async,
+                                           conn->plain_output_stream);
+                conn->plain_output_stream = drystream;
+            }
             return shutting_down_outgoing(conn, buf, count);
         }
         if (n < 0) {