From 1c8f564fdb081526d92b4834a4ee4d0a15c14cc4 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Wed, 21 Jun 2023 13:49:58 +0400
Subject: [PATCH] feat: add logging of ssh connections to agent (#8096)

* feat: adds logging of ssh connections to agent

Signed-off-by: Spike Curtis <spike@coder.com>

* code review improvements

Signed-off-by: Spike Curtis <spike@coder.com>

---------

Signed-off-by: Spike Curtis <spike@coder.com>
---
 agent/agentssh/agentssh.go | 31 ++++++++++++++++++++++++-------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
index a271835df0..3d4ac69c47 100644
--- a/agent/agentssh/agentssh.go
+++ b/agent/agentssh/agentssh.go
@@ -180,21 +180,24 @@ func (s *Server) ConnStats() ConnStats {
 }
 
 func (s *Server) sessionHandler(session ssh.Session) {
+	logger := s.logger.With(slog.F("remote_addr", session.RemoteAddr()), slog.F("local_addr", session.LocalAddr()))
+	logger.Info(session.Context(), "handling ssh session")
+	ctx := session.Context()
 	if !s.trackSession(session, true) {
 		// See (*Server).Close() for why we call Close instead of Exit.
 		_ = session.Close()
+		logger.Info(ctx, "unable to accept new session, server is closing")
 		return
 	}
 	defer s.trackSession(session, false)
 
-	ctx := session.Context()
-
 	extraEnv := make([]string, 0)
 	x11, hasX11 := session.X11()
 	if hasX11 {
 		handled := s.x11Handler(session.Context(), x11)
 		if !handled {
 			_ = session.Exit(1)
+			logger.Error(ctx, "x11 handler failed")
 			return
 		}
 		extraEnv = append(extraEnv, fmt.Sprintf("DISPLAY=:%d.0", x11.ScreenNumber))
@@ -206,7 +209,7 @@ func (s *Server) sessionHandler(session ssh.Session) {
 		s.sftpHandler(session)
 		return
 	default:
-		s.logger.Debug(ctx, "unsupported subsystem", slog.F("subsystem", ss))
+		logger.Warn(ctx, "unsupported subsystem", slog.F("subsystem", ss))
 		_ = session.Exit(1)
 		return
 	}
@@ -214,17 +217,18 @@ func (s *Server) sessionHandler(session ssh.Session) {
 	err := s.sessionStart(session, extraEnv)
 	var exitError *exec.ExitError
 	if xerrors.As(err, &exitError) {
-		s.logger.Warn(ctx, "ssh session returned", slog.Error(exitError))
+		logger.Info(ctx, "ssh session returned", slog.Error(exitError))
 		_ = session.Exit(exitError.ExitCode())
 		return
 	}
 	if err != nil {
-		s.logger.Warn(ctx, "ssh session failed", slog.Error(err))
+		logger.Warn(ctx, "ssh session failed", slog.Error(err))
 		// This exit code is designed to be unlikely to be confused for a legit exit code
 		// from the process.
 		_ = session.Exit(MagicSessionErrorCode)
 		return
 	}
+	logger.Info(ctx, "normal ssh session exit")
 	_ = session.Exit(0)
 }
 
@@ -565,7 +569,12 @@ func (s *Server) CreateCommand(ctx context.Context, script string, env []string)
 	return cmd, nil
 }
 
-func (s *Server) Serve(l net.Listener) error {
+func (s *Server) Serve(l net.Listener) (retErr error) {
+	s.logger.Info(context.Background(), "started serving listener", slog.F("listen_addr", l.Addr()))
+	defer func() {
+		s.logger.Info(context.Background(), "stopped serving listener",
+			slog.F("listen_addr", l.Addr()), slog.Error(retErr))
+	}()
 	defer l.Close()
 
 	s.trackListener(l, true)
@@ -580,15 +589,23 @@ func (s *Server) Serve(l net.Listener) error {
 }
 
 func (s *Server) handleConn(l net.Listener, c net.Conn) {
+	logger := s.logger.With(
+		slog.F("remote_addr", c.RemoteAddr()),
+		slog.F("local_addr", c.LocalAddr()),
+		slog.F("listen_addr", l.Addr()))
 	defer c.Close()
 
 	if !s.trackConn(l, c, true) {
 		// Server is closed or we no longer want
 		// connections from this listener.
-		s.logger.Debug(context.Background(), "received connection after server closed")
+		logger.Info(context.Background(), "received connection after server closed")
 		return
 	}
 	defer s.trackConn(l, c, false)
+	logger.Info(context.Background(), "started serving connection")
+	defer func() {
+		logger.Info(context.Background(), "stopped serving connection")
+	}()
 
 	s.srv.HandleConn(c)
 }