feat: add ConnectionCompleteCallback

spikecurtis · spikecurtis · commit c92d70594c77 · 2023-06-21T05:39:55.000Z
Signed-off-by: Spike Curtis &lt;spike@coder.com&gt;
diff --git a/server.go b/server.go
@@ -51,7 +51,10 @@ type Server struct {
 	ServerConfigCallback          ServerConfigCallback          // callback for configuring detailed SSH options
 	SessionRequestCallback        SessionRequestCallback        // callback for allowing or denying SSH sessions
 
-	ConnectionFailedCallback ConnectionFailedCallback // callback to report connection failures
+	// server calls Failed callback for connections that fail initial handshake, and Complete callback for those that
+	// succeed, never both.
+	ConnectionFailedCallback   ConnectionFailedCallback   // callback to report connection failures
+	ConnectionCompleteCallback ConnectionCompleteCallback // callback to report connection completion
 
 	IdleTimeout time.Duration // connection timeout when no activity, none if empty
 	MaxTimeout  time.Duration // absolute connection timeout, none if empty
@@ -295,6 +298,11 @@ func (srv *Server) HandleConn(newConn net.Conn) {
 		}
 		return
 	}
+	if srv.ConnectionCompleteCallback != nil {
+		defer func() {
+			srv.ConnectionCompleteCallback(sshConn, sshConn.Wait())
+		}()
+	}
 
 	srv.trackConn(sshConn, true)
 	defer srv.trackConn(sshConn, false)
diff --git a/ssh.go b/ssh.go
@@ -71,6 +71,14 @@ type ServerConfigCallback func(ctx Context) *gossh.ServerConfig
 // Please note: the net.Conn is likely to be closed at this point
 type ConnectionFailedCallback func(conn net.Conn, err error)
 
+// ConnectionCompleteCallback is a hook for reporting connections that
+// complete.  The included error is from the underlying SSH transport
+// protocol mux (golang.org/x/crypto/ssh), and is non-nil, even for
+// normal termination.
+//
+// Please note: the ServerConn is closed at this point
+type ConnectionCompleteCallback func(conn *gossh.ServerConn, err error)
+
 // Window represents the size of a PTY window.
 //
 // See https://datatracker.ietf.org/doc/html/rfc4254#section-6.2
