From 4def17e560b66f02b10cc6109f2d8e0ce5777faf Mon Sep 17 00:00:00 2001
From: Anton Miniailo <anton@goteleport.com>
Date: Thu, 28 Sep 2023 10:52:21 -0400
Subject: [PATCH] Add connection information to multiplexer logs so it's easier
 to investigate (#32712)

---
 lib/multiplexer/multiplexer.go | 15 ++++++++++++---
 lib/multiplexer/tls.go         | 15 ++++++++++++---
 lib/multiplexer/web.go         | 15 ++++++++++++---
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/lib/multiplexer/multiplexer.go b/lib/multiplexer/multiplexer.go
index d6cd0533df4d..5fbae24c7243 100644
--- a/lib/multiplexer/multiplexer.go
+++ b/lib/multiplexer/multiplexer.go
@@ -287,7 +287,10 @@ func (m *Mux) detectAndForward(conn net.Conn) {
 	connWrapper, err := m.detect(conn)
 	if err != nil {
 		if trace.Unwrap(err) != io.EOF {
-			m.logLimiter.Log(m.Entry, log.WarnLevel, trace.DebugReport(err))
+			m.logLimiter.Log(m.Entry.WithFields(log.Fields{
+				"src_addr": conn.RemoteAddr(),
+				"dst_addr": conn.LocalAddr(),
+			}), log.WarnLevel, trace.DebugReport(err))
 		}
 		conn.Close()
 		return
@@ -302,9 +305,15 @@ func (m *Mux) detectAndForward(conn net.Conn) {
 	listener := m.protocolListener(connWrapper.protocol)
 	if listener == nil {
 		if connWrapper.protocol == ProtoHTTP {
-			m.Debug("Detected an HTTP request. If this is for a health check, use an HTTPS request instead.")
+			m.WithFields(log.Fields{
+				"src_addr": connWrapper.RemoteAddr(),
+				"dst_addr": connWrapper.LocalAddr(),
+			}).Debug("Detected an HTTP request. If this is for a health check, use an HTTPS request instead.")
 		}
-		m.Debugf("Closing %[1]s connection: %[1]s listener is disabled.", connWrapper.protocol)
+		m.WithFields(log.Fields{
+			"src_addr": connWrapper.RemoteAddr(),
+			"dst_addr": connWrapper.LocalAddr(),
+		}).Debugf("Closing %[1]s connection: %[1]s listener is disabled.", connWrapper.protocol)
 		connWrapper.Close()
 		return
 	}
diff --git a/lib/multiplexer/tls.go b/lib/multiplexer/tls.go
index 3d1d07d48573..9231ec34f15a 100644
--- a/lib/multiplexer/tls.go
+++ b/lib/multiplexer/tls.go
@@ -112,7 +112,10 @@ func (l *TLSListener) Serve() error {
 			tlsConn, ok := conn.(*tls.Conn)
 			if !ok {
 				conn.Close()
-				log.Errorf("Expected tls.Conn, got %T, internal usage error.", conn)
+				l.log.WithFields(log.Fields{
+					"src_addr": conn.RemoteAddr(),
+					"dst_addr": conn.LocalAddr(),
+				}).Errorf("Expected tls.Conn, got %T, internal usage error.", conn)
 				continue
 			}
 			go l.detectAndForward(tlsConn)
@@ -141,7 +144,10 @@ func (l *TLSListener) detectAndForward(conn *tls.Conn) {
 	start := l.cfg.Clock.Now()
 	if err := conn.Handshake(); err != nil {
 		if trace.Unwrap(err) != io.EOF {
-			l.log.WithError(err).Warning("Handshake failed.")
+			l.log.WithFields(log.Fields{
+				"src_addr": conn.RemoteAddr(),
+				"dst_addr": conn.LocalAddr(),
+			}).WithError(err).Warning("Handshake failed.")
 		}
 		conn.Close()
 		return
@@ -167,7 +173,10 @@ func (l *TLSListener) detectAndForward(conn *tls.Conn) {
 		l.httpListener.HandleConnection(l.context, conn)
 	default:
 		conn.Close()
-		l.log.WithError(err).Errorf("unsupported protocol: %v", conn.ConnectionState().NegotiatedProtocol)
+		l.log.WithFields(log.Fields{
+			"src_addr": conn.RemoteAddr(),
+			"dst_addr": conn.LocalAddr(),
+		}).WithError(err).Errorf("unsupported protocol: %v", conn.ConnectionState().NegotiatedProtocol)
 	}
 }
 
diff --git a/lib/multiplexer/web.go b/lib/multiplexer/web.go
index b3de7d3c168e..53f756efb41a 100644
--- a/lib/multiplexer/web.go
+++ b/lib/multiplexer/web.go
@@ -113,7 +113,10 @@ func (l *WebListener) Serve() error {
 
 		tlsConn, ok := conn.(*tls.Conn)
 		if !ok {
-			l.log.Errorf("Expected *tls.Conn, got %T.", conn)
+			l.log.WithFields(logrus.Fields{
+				"src_addr": conn.RemoteAddr(),
+				"dst_addr": conn.LocalAddr(),
+			}).Errorf("Expected *tls.Conn, got %T.", conn)
 			conn.Close()
 			continue
 		}
@@ -132,7 +135,10 @@ func (l *WebListener) detectAndForward(conn *tls.Conn) {
 
 	if err := conn.Handshake(); err != nil {
 		if trace.Unwrap(err) != io.EOF {
-			l.log.WithError(err).Warn("Handshake failed.")
+			l.log.WithFields(logrus.Fields{
+				"src_addr": conn.RemoteAddr(),
+				"dst_addr": conn.LocalAddr(),
+			}).WithError(err).Warn("Handshake failed.")
 		}
 		conn.Close()
 		return
@@ -151,7 +157,10 @@ func (l *WebListener) detectAndForward(conn *tls.Conn) {
 	// tls listener.
 	isDatabaseConnection, err := dbcommon.IsDatabaseConnection(conn.ConnectionState())
 	if err != nil {
-		l.log.WithError(err).Debug("Failed to check if connection is database connection.")
+		l.log.WithFields(logrus.Fields{
+			"src_addr": conn.RemoteAddr(),
+			"dst_addr": conn.LocalAddr(),
+		}).WithError(err).Debug("Failed to check if connection is database connection.")
 	}
 	if isDatabaseConnection {
 		l.dbListener.HandleConnection(l.context, conn)