From fe2df405e02825c1b0745ea665d378de32da2d61 Mon Sep 17 00:00:00 2001
From: Kenneth Jenkins <51246568+kenjenkins@users.noreply.github.com>
Date: Fri, 11 Apr 2025 13:46:47 -0700
Subject: [PATCH] test that upstream banner is sent to client

---
 integration2/ssh_int_test.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/integration2/ssh_int_test.go b/integration2/ssh_int_test.go
index ebe8bf102..91153c80d 100644
--- a/integration2/ssh_int_test.go
+++ b/integration2/ssh_int_test.go
@@ -28,6 +28,7 @@ func TestSSH(t *testing.T) {
 
 	// ssh client setup
 	var ki scenarios.EmptyKeyboardInteractiveChallenge
+	var bannerReceived string
 	clientConfig := &ssh.ClientConfig{
 		User: "demo@example",
 		Auth: []ssh.AuthMethod{
@@ -35,6 +36,7 @@ func TestSSH(t *testing.T) {
 			ssh.KeyboardInteractive(ki.Do),
 		},
 		HostKeyCallback: ssh.FixedHostKey(newPublicKey(t, serverHostKey.Public())),
+		BannerCallback:  func(message string) error { bannerReceived = message; return nil },
 	}
 
 	// pomerium + upstream setup
@@ -56,6 +58,9 @@ func TestSSH(t *testing.T) {
 	up := upstreams.SSH(
 		upstreams.WithHostKeys(newSignerFromKey(t, serverHostKey)),
 		upstreams.WithPublicKeyCallback(certChecker.Authenticate),
+		upstreams.WithBannerCallback(func(_ ssh.ConnMetadata) string {
+			return "UPSTREAM BANNER"
+		}),
 	)
 	up.SetServerConnCallback(echoShell{t}.handleConnection)
 	r := up.Route().
@@ -70,6 +75,8 @@ func TestSSH(t *testing.T) {
 	require.NoError(t, err)
 	defer client.Close()
 
+	assert.Equal(t, "UPSTREAM BANNER", bannerReceived)
+
 	sess, err := client.NewSession()
 	require.NoError(t, err)
 	defer sess.Close()