Gate the handshake drivers on a disconnect

wolfSSH_accept() and wolfSSH_connect() drive the handshake only while the
session is live. Both gate on SendAfterDisconnect() ahead of the pending-send
block, which would otherwise flush a queued disconnect and count it as the
next handshake message; the shutdown paths own that flush.

- The prototype sits ahead of both drivers, since either can be the only one
  built.
- TestDisconnectGatesAccept() and TestDisconnectGatesConnect() cover a local
  disconnect, one from the peer, and a queued short send. One test per
  endpoint, so a single-sided build keeps the coverage that applies to it.
- A received disconnect used to reach the error-state test in wolfSSH_accept()
  and report WS_INVALID_STATE_E; the gate answers WS_FATAL_ERROR first, and
  both tests pin that.
- The contract comments in ssh.h and internal.h drop the ungated note.

Author: ejohnstown

src/ssh.c

@@ -550,6 +550,10 @@ int wolfSSH_CTX_UseTpmHostKey(WOLFSSH_CTX* ctx,
 #endif /* WOLFSSH_TPM */
 
 
+/* Defined below, ahead of both drivers; either can be the only one built. */
+static int SendAfterDisconnect(WOLFSSH* ssh);
+
+
 #ifndef NO_WOLFSSH_SERVER
 
 const char acceptError[] = "accept error: %s, %d";
@@ -563,6 +567,11 @@ int wolfSSH_accept(WOLFSSH* ssh)
     if (ssh == NULL)
         return WS_BAD_ARGUMENT;
 
+    /* No handshake on a session that is over. The pending-send block below
+     * would flush a queued disconnect as the next handshake message. */
+    if (SendAfterDisconnect(ssh))
+        return WS_FATAL_ERROR;
+
     /* clear want read/writes for retry */
     if (ssh->error == WS_WANT_READ || ssh->error == WS_WANT_WRITE || ssh->error == WS_AUTH_PENDING)
         ssh->error = 0;
@@ -826,6 +835,11 @@ int wolfSSH_connect(WOLFSSH* ssh)
     if (ssh == NULL)
         return WS_BAD_ARGUMENT;
 
+    /* See wolfSSH_accept(). No error-state test here, so the peer's
+     * disconnect reaches the state machine like a local one. */
+    if (SendAfterDisconnect(ssh))
+        return WS_FATAL_ERROR;
+
     /* check if data pending to be sent */
     if (ssh->outputBuffer.length > 0 &&
             ssh->connectState < CONNECT_SERVER_CHANNEL_REQUEST_DONE) {

tests/regress.c

@@ -4194,6 +4194,222 @@ static void TestDisconnectOutranksRekey(void)
 }
 
 
+
+#ifndef NO_WOLFSSH_SERVER
+
+/* wolfSSH_accept() drives the handshake, so a session that is already over
+ * must stop it the way it stops every other sender. The short-send case is
+ * the sharp one: the pending-send block at the top would push out a
+ * disconnect left queued by a short send and then count it as the handshake
+ * message the state machine was waiting for. RFC 4253 section 11.1. */
+static void TestDisconnectGatesAccept(void)
+{
+    WOLFSSH_CTX* ctx;
+    WOLFSSH* ssh;
+    MemIo io;
+    byte in[128];
+    byte out[512];
+    word32 inSz;
+    word32 quietSz;
+    byte state;
+    int ret;
+
+    /* A local disconnect leaves ssh->error clear, so the "in error state"
+     * test in wolfSSH_accept() never sees it. */
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
+    AssertNotNull(ctx);
+    wolfSSH_SetIORecv(ctx, MemRecv);
+    wolfSSH_SetIOSend(ctx, MemSend);
+
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+    AddSessionChannel(ssh);
+    /* Not one of the states wolfSSH_accept() holds back, so an unwanted
+     * advance shows up in the assertions below. */
+    ssh->acceptState = ACCEPT_SERVER_USERAUTH_SENT;
+
+    MemIoInit(&io, NULL, 0, out, sizeof(out));
+    wolfSSH_SetIOReadCtx(ssh, &io);
+    wolfSSH_SetIOWriteCtx(ssh, &io);
+
+    AssertIntEQ(wolfSSH_SendDisconnect(ssh, WOLFSSH_DISCONNECT_BY_APPLICATION),
+            WS_SUCCESS);
+    AssertTrue(ssh->disconnected);
+    AssertIntEQ(wolfSSH_get_error(ssh), 0);
+    quietSz = io.outSz;
+    state = ssh->acceptState;
+
+    ret = wolfSSH_accept(ssh);
+    AssertIntEQ(ret, WS_FATAL_ERROR);
+    AssertIntEQ(wolfSSH_get_error(ssh), WS_DISCONNECT);
+    /* No handshake packet, and the state machine did not move. */
+    AssertIntEQ(io.outSz, quietSz);
+    AssertIntEQ(ssh->acceptState, state);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+
+    /* Our disconnect short-sends, so it is sitting in the output buffer
+     * with a flush owed. wolfSSH_shutdown() and wolfSSH_SendDisconnect()
+     * own that flush; wolfSSH_accept() must leave it alone. */
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
+    AssertNotNull(ctx);
+    wolfSSH_SetIORecv(ctx, MemRecv);
+    wolfSSH_SetIOSend(ctx, MemSendWantWrite);
+
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+    AddSessionChannel(ssh);
+    /* Not one of the states wolfSSH_accept() holds back, so an unwanted
+     * advance shows up in the assertions below. */
+    ssh->acceptState = ACCEPT_SERVER_USERAUTH_SENT;
+
+    MemIoInit(&io, NULL, 0, out, sizeof(out));
+    wolfSSH_SetIOReadCtx(ssh, &io);
+    wolfSSH_SetIOWriteCtx(ssh, &io);
+
+    MemSendWantWriteCount = 1;
+    AssertIntEQ(wolfSSH_SendDisconnect(ssh, WOLFSSH_DISCONNECT_BY_APPLICATION),
+            WS_WANT_WRITE);
+    AssertTrue(ssh->disconnected);
+    AssertTrue(ssh->disconnectTxd);
+    AssertTrue(wolfSSH_OutputPending(ssh));
+    AssertIntEQ(io.outSz, 0);
+    state = ssh->acceptState;
+
+    ret = wolfSSH_accept(ssh);
+    AssertIntEQ(ret, WS_FATAL_ERROR);
+    AssertIntEQ(wolfSSH_get_error(ssh), WS_DISCONNECT);
+    /* Still queued, and not mistaken for the awaited handshake message. */
+    AssertIntEQ(io.outSz, 0);
+    AssertTrue(wolfSSH_OutputPending(ssh));
+    AssertIntEQ(ssh->acceptState, state);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+
+    /* The peer's disconnect latches WS_DISCONNECT, which the error-state
+     * test below the gate used to answer with WS_INVALID_STATE_E. */
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
+    AssertNotNull(ctx);
+    wolfSSH_SetIORecv(ctx, MemRecv);
+    wolfSSH_SetIOSend(ctx, MemSend);
+
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+    AddSessionChannel(ssh);
+    ssh->acceptState = ACCEPT_SERVER_USERAUTH_SENT;
+
+    inSz = BuildDisconnectPacket(WOLFSSH_DISCONNECT_BY_APPLICATION,
+            in, sizeof(in));
+    MemIoInit(&io, in, inSz, out, sizeof(out));
+    wolfSSH_SetIOReadCtx(ssh, &io);
+    wolfSSH_SetIOWriteCtx(ssh, &io);
+
+    AssertIntEQ(DoReceive(ssh), WS_FATAL_ERROR);
+    AssertTrue(ssh->disconnected);
+    AssertIntEQ(wolfSSH_get_error(ssh), WS_DISCONNECT);
+    quietSz = io.outSz;
+    state = ssh->acceptState;
+
+    ret = wolfSSH_accept(ssh);
+    AssertIntEQ(ret, WS_FATAL_ERROR);
+    AssertIntEQ(wolfSSH_get_error(ssh), WS_DISCONNECT);
+    AssertIntEQ(io.outSz, quietSz);
+    AssertIntEQ(ssh->acceptState, state);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+}
+
+#endif /* !NO_WOLFSSH_SERVER */
+
+
+#ifndef NO_WOLFSSH_CLIENT
+
+/* The same gate on the client's driver, which has no error-state test of
+ * its own. Split from the accept test so a single-sided build keeps the
+ * coverage that applies to it. */
+static void TestDisconnectGatesConnect(void)
+{
+    WOLFSSH_CTX* ctx;
+    WOLFSSH* ssh;
+    MemIo io;
+    byte in[128];
+    byte out[512];
+    word32 inSz;
+    byte state;
+    int ret;
+
+    /* wolfSSH_connect() has no error-state test of its own, so it reaches
+     * the state machine after a disconnect from either side. */
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
+    AssertNotNull(ctx);
+    wolfSSH_SetIORecv(ctx, MemRecv);
+    wolfSSH_SetIOSend(ctx, MemSendWantWrite);
+
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+    AddSessionChannel(ssh);
+    ssh->connectState = CONNECT_SERVER_USERAUTH_ACCEPT_DONE;
+
+    MemIoInit(&io, NULL, 0, out, sizeof(out));
+    wolfSSH_SetIOReadCtx(ssh, &io);
+    wolfSSH_SetIOWriteCtx(ssh, &io);
+
+    MemSendWantWriteCount = 1;
+    AssertIntEQ(wolfSSH_SendDisconnect(ssh, WOLFSSH_DISCONNECT_BY_APPLICATION),
+            WS_WANT_WRITE);
+    AssertTrue(ssh->disconnected);
+    AssertTrue(wolfSSH_OutputPending(ssh));
+    AssertIntEQ(io.outSz, 0);
+    state = ssh->connectState;
+
+    ret = wolfSSH_connect(ssh);
+    AssertIntEQ(ret, WS_FATAL_ERROR);
+    AssertIntEQ(wolfSSH_get_error(ssh), WS_DISCONNECT);
+    AssertIntEQ(io.outSz, 0);
+    AssertTrue(wolfSSH_OutputPending(ssh));
+    AssertIntEQ(ssh->connectState, state);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+
+    /* The peer's disconnect reaches the same gate. */
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_CLIENT, NULL);
+    AssertNotNull(ctx);
+    wolfSSH_SetIORecv(ctx, MemRecv);
+    wolfSSH_SetIOSend(ctx, MemSend);
+
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+    AddSessionChannel(ssh);
+    ssh->connectState = CONNECT_SERVER_USERAUTH_ACCEPT_DONE;
+
+    inSz = BuildDisconnectPacket(WOLFSSH_DISCONNECT_BY_APPLICATION,
+            in, sizeof(in));
+    MemIoInit(&io, in, inSz, out, sizeof(out));
+    wolfSSH_SetIOReadCtx(ssh, &io);
+    wolfSSH_SetIOWriteCtx(ssh, &io);
+
+    AssertIntEQ(DoReceive(ssh), WS_FATAL_ERROR);
+    AssertTrue(ssh->disconnected);
+    AssertIntEQ(wolfSSH_get_error(ssh), WS_DISCONNECT);
+    state = ssh->connectState;
+
+    ret = wolfSSH_connect(ssh);
+    AssertIntEQ(ret, WS_FATAL_ERROR);
+    AssertIntEQ(wolfSSH_get_error(ssh), WS_DISCONNECT);
+    AssertIntEQ(io.outSz, 0);
+    AssertIntEQ(ssh->connectState, state);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+}
+
+#endif /* !NO_WOLFSSH_CLIENT */
+
+
 /* disconnectTxd means "a flush is owed", not "a disconnect was sent". Once
  * ours has gone out, a teardown call must not push whatever the internal
  * senders queued behind it. */
@@ -8079,6 +8295,12 @@ int main(int argc, char** argv)
     TestDisconnectDrainsBufferedData();
     TestDisconnectBlocksEverySend();
     TestSendDisconnectIsTerminal();
+#ifndef NO_WOLFSSH_SERVER
+    TestDisconnectGatesAccept();
+#endif
+#ifndef NO_WOLFSSH_CLIENT
+    TestDisconnectGatesConnect();
+#endif
     TestDisconnectQuietWindowAdjust();
     TestDisconnectBlocksChannelAndFwdSends();
     TestStreamExitReportsDisconnect();

wolfssh/internal.h

@@ -1105,9 +1105,8 @@ struct WOLFSSH {
     /* Set when a DISCONNECT is sent or received. Gates the public send
      * calls, so nothing more goes out. Reads still hand back what arrived
      * before the disconnect; the head-of-list reads report it once their
-     * buffer runs dry, unless a CHANNEL_EOF arrived first. wolfSSH_worker(),
-     * wolfSSH_accept() and wolfSSH_connect() are not gated; the shutdown
-     * paths pump the worker. */
+     * buffer runs dry, unless a CHANNEL_EOF arrived first.
+     * wolfSSH_worker() is not gated; the shutdown paths pump the worker. */
     byte disconnected;
     /* Set once SendDisconnect() has bundled our own DISCONNECT into the
      * output buffer, and cleared once wolfSSH_SendPacket() drains it, so it

wolfssh/ssh.h

@@ -563,9 +563,10 @@ WOLFSSH_API int wolfSSH_connect(WOLFSSH* ssh);
 WOLFSSH_API int wolfSSH_shutdown(WOLFSSH* ssh);
 /* A disconnect, sent or received, ends the session. Nothing more goes out:
  * wolfSSH_shutdown() above this comment, and every send call below it,
- * report WS_DISCONNECT from then on. wolfSSH_accept() and
- * wolfSSH_connect() are not gated; do not drive the handshake after a
- * disconnect. Reads are not gated either, so channel data that
+ * report WS_DISCONNECT from then on, as do wolfSSH_accept() and
+ * wolfSSH_connect(). Neither driver flushes a disconnect of ours left queued
+ * by a short send; wolfSSH_shutdown() or another wolfSSH_SendDisconnect()
+ * owns that. Reads are not gated, so channel data that
  * arrived before the disconnect can still be drained; wolfSSH_stream_read()
  * and wolfSSH_stream_peek() report WS_DISCONNECT once their buffer runs
  * dry. A CHANNEL_EOF already received outranks that drain: both report