From 9851d37ccbded022ef542f5e65c47cb560a6b6da Mon Sep 17 00:00:00 2001
From: Simon Tatham <anakin@pobox.com>
Date: Mon, 24 May 2021 15:04:35 +0100
Subject: [PATCH] Add test script for simultaneous agent connections.

This script makes 128 connections to your SSH agent at once, and then
sends requests down them in random order to check that the agent is
correctly selecting between all its incoming sockets / named pipes /
whatever.

128 is bigger than MAXIMUM_WAIT_OBJECTS, so a successful run of this
script inside a Windows PuTTY agent-forwarding to a Pageant indicates
that both the PuTTY and the Pageant are managing to handle >64 I/O
subthreads without overloading their event loop.
---
 test/agentmulti.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++
 test/ssh.py        |  1 +
 2 files changed, 57 insertions(+)
 create mode 100755 test/agentmulti.py

diff --git a/test/agentmulti.py b/test/agentmulti.py
new file mode 100755
index 00000000..019bf2b6
--- /dev/null
+++ b/test/agentmulti.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+
+import argparse
+import os
+import random
+import socket
+import sys
+
+from ssh import *
+
+def make_connections(n):
+    connections = []
+
+    for _ in range(n):
+        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        s.connect(os.environ["SSH_AUTH_SOCK"])
+        connections.append(s)
+
+    return connections
+
+def use_connection(s, idstring):
+    print("Trying {}...".format(idstring), end="")
+    sys.stdout.flush()
+
+    s.send(ssh_string(ssh_byte(SSH2_AGENTC_EXTENSION) + ssh_string(
+        b"nonexistent-agent-extension@putty.projects.tartarus.org")))
+    length = ssh_decode_uint32(s.recv(4))
+    assert length < AGENT_MAX_MSGLEN
+    msg = s.recv(length)
+    msgtype = msg[0]
+    msgstring = (
+        "SSH_AGENT_EXTENSION_FAILURE" if msgtype == SSH_AGENT_EXTENSION_FAILURE
+        else "SSH_AGENT_FAILURE" if msgtype == SSH_AGENT_FAILURE
+        else "type {:d}".format(msgtype))
+    print("got", msgstring, "with {:d}-byte payload".format(len(msg)-1))
+
+def randomly_use_connections(connections, iterations):
+    for _ in range(iterations):
+        index = random.randrange(0, len(connections))
+        s = connections[index]
+        use_connection(connections[index], "#{:d}".format(index))
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Test handling of multiple agent connections.')
+    parser.add_argument("--nsockets", type=int, default=128,
+                        help="Number of simultaneous connections to make.")
+    parser.add_argument("--ntries", type=int, default=1024,
+                        help="Number of messages to send in total.")
+    args = parser.parse_args()
+
+    connections = make_connections(args.nsockets)
+    randomly_use_connections(connections, args.ntries)
+
+if __name__ == '__main__':
+    main()
diff --git a/test/ssh.py b/test/ssh.py
index 90eccb7e..c4f2531f 100644
--- a/test/ssh.py
+++ b/test/ssh.py
@@ -88,6 +88,7 @@ SSH1_AGENTC_REMOVE_RSA_IDENTITY = 8
 SSH1_AGENTC_REMOVE_ALL_RSA_IDENTITIES = 9
 SSH_AGENT_FAILURE = 5
 SSH_AGENT_SUCCESS = 6
+SSH_AGENT_EXTENSION_FAILURE = 28
 SSH2_AGENTC_REQUEST_IDENTITIES = 11
 SSH2_AGENT_IDENTITIES_ANSWER = 12
 SSH2_AGENTC_SIGN_REQUEST = 13