From fe63b5d57ebbdef43f0903b21b96b22615fddeb3 Mon Sep 17 00:00:00 2001
From: Simon Tatham <anakin@pobox.com>
Date: Sat, 29 Apr 2023 11:34:08 +0100
Subject: [PATCH] Uppity: add a stunt mode --close-after-banner.

A user reported yesterday that PuTTY can fail to print a userauth
banner message if the server sends one and then immediately slams the
connection shut. The first step to fixing this is making a convenient
way to reproduce that server behaviour.

(Apparently the real use case has to do with account expiry - the
server in question presumably doesn't have enough layer violations to
be able to put the text "Your account has expired" into an
SSH_MSG_DISCONNECT, so instead it does the next best thing and sends
it as a userauth banner immediately before disconnection.)
---
 ssh/server.c           |  1 +
 ssh/server.h           |  1 +
 ssh/userauth2-server.c | 20 ++++++++++++++++++++
 unix/uppity.c          |  2 ++
 4 files changed, 24 insertions(+)

diff --git a/ssh/server.c b/ssh/server.c
index 4e315874..d3a129fa 100644
--- a/ssh/server.c
+++ b/ssh/server.c
@@ -502,6 +502,7 @@ void ssh_sw_abort(Ssh *ssh, const char *fmt, ...)
 void ssh_user_close(Ssh *ssh, const char *fmt, ...)
 {
     server *srv = container_of(ssh, server, ssh);
+    ssh_bpp_handle_output(srv->bpp);
     LOG_FORMATTED_MSG(srv->logctx, fmt);
     queue_toplevel_callback(ssh_server_free_callback, srv);
 }
diff --git a/ssh/server.h b/ssh/server.h
index c2b3647b..dc73747c 100644
--- a/ssh/server.h
+++ b/ssh/server.h
@@ -23,6 +23,7 @@ struct SshServerConfig {
     bool stunt_open_unconditional_agent_socket;
     bool stunt_allow_trivial_ki_auth;
     bool stunt_return_success_to_pubkey_offer;
+    bool stunt_close_after_banner;
 };
 
 Plug *ssh_server_plug(
diff --git a/ssh/userauth2-server.c b/ssh/userauth2-server.c
index bfe258ce..a7e8929c 100644
--- a/ssh/userauth2-server.c
+++ b/ssh/userauth2-server.c
@@ -113,6 +113,21 @@ static void ssh2_userauth_server_add_session_id(
     }
 }
 
+static void ssh2_userauth_server_close_after_banner(void *vctx)
+{
+    struct ssh2_userauth_server_state *s =
+        (struct ssh2_userauth_server_state *)vctx;
+
+    if (pq_peek(s->ppl.out_pq)) {
+        /* Don't close the connection until we've passed on our final banner
+         * packet to the lower layer */
+        queue_toplevel_callback(ssh2_userauth_server_close_after_banner, s);
+    } else {
+        ssh_user_close(s->ppl.ssh, "Closing connection on request due to "
+                       "--close-after-banner");
+    }
+}
+
 static void ssh2_userauth_server_process_queue(PacketProtocolLayer *ppl)
 {
     struct ssh2_userauth_server_state *s =
@@ -131,6 +146,11 @@ static void ssh2_userauth_server_process_queue(PacketProtocolLayer *ppl)
         pq_push(s->ppl.out_pq, pktout);
     }
 
+    if (s->ssc->stunt_close_after_banner) {
+        queue_toplevel_callback(ssh2_userauth_server_close_after_banner, s);
+        crReturnV;
+    }
+
     while (1) {
         crMaybeWaitUntilV((pktin = ssh2_userauth_server_pop(s)) != NULL);
         if (pktin->type != SSH2_MSG_USERAUTH_REQUEST) {
diff --git a/unix/uppity.c b/unix/uppity.c
index 36f2c99d..711c2952 100644
--- a/unix/uppity.c
+++ b/unix/uppity.c
@@ -924,6 +924,8 @@ int main(int argc, char **argv)
             ci->ssc.stunt_allow_trivial_ki_auth = true;
         } else if (!strcmp(arg, "--return-success-to-pubkey-offer")) {
             ci->ssc.stunt_return_success_to_pubkey_offer = true;
+        } else if (!strcmp(arg, "--close-after-banner")) {
+            ci->ssc.stunt_close_after_banner = true;
         } else {
             fprintf(stderr, "%s: unrecognised option '%s'\n", appname, arg);
             exit(1);