Handle ALPN handler error for upgrade flow and add some reporting (#54164)

* Handle ALPN handler error for upgrade flow and add some reporting

* fix test and use logutils.StringerAttr

* make lint happy

* rename metrics

* expliclity pass connSource instead of using context

* fix build...

Author: greedy52

lib/service/service.go

@@ -5569,7 +5569,7 @@ func (process *TeleportProcess) initProxyEndpoint(conn *Connector) error {
 		if err != nil {
 			return trace.Wrap(err)
 		}
-		alpnHandlerForWeb.Set(alpnServer.MakeConnectionHandler(alpnTLSConfigForWeb))
+		alpnHandlerForWeb.Set(alpnServer.MakeConnectionHandler(alpnTLSConfigForWeb, alpncommon.ConnHandlerSourceWeb))
 
 		process.RegisterCriticalFunc("proxy.tls.alpn.sni.proxy", func() error {
 			logger.InfoContext(process.ExitContext(), "Starting TLS ALPN SNI proxy server on.", "listen_address", logutils.StringerAttr(listeners.alpn.Addr()))

lib/srv/alpnproxy/common/reporting.go

@@ -0,0 +1,32 @@
+/*
+ * Teleport
+ * Copyright (C) 2025  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package common
+
+// ConnHandlerSource specifies the source of the connection that is passed into
+// the TLS routing connection handler.
+type ConnHandlerSource string
+
+const (
+	// ConnHandlerSourceListener indicates the connection is from the TLS
+	// routing port listener.
+	ConnHandlerSourceListener ConnHandlerSource = "listener"
+	// ConnHandlerSourceWeb indicates the connection is from web handler like
+	// connection upgrades or certain Web UI flows like DB session via Web UI.
+	ConnHandlerSourceWeb ConnHandlerSource = "web"
+)

lib/srv/alpnproxy/conn.go

@@ -21,6 +21,7 @@ package alpnproxy
 import (
 	"io"
 	"net"
+	"sync"
 	"time"
 
 	"github.com/gravitational/teleport/lib/utils"
@@ -99,3 +100,32 @@ func (conn readOnlyConn) RemoteAddr() net.Addr               { return &utils.Net
 func (conn readOnlyConn) SetDeadline(t time.Time) error      { return nil }
 func (conn readOnlyConn) SetReadDeadline(t time.Time) error  { return nil }
 func (conn readOnlyConn) SetWriteDeadline(t time.Time) error { return nil }
+
+type reportingConn struct {
+	net.Conn
+	closeOnce func() error
+}
+
+// newReportingConn returns a net.Conn that wraps the input conn, increments the
+// active connections gauge on creation, and decreases the active connections
+// gauge on close.
+func newReportingConn(conn net.Conn, clientALPN string, connSource string) net.Conn {
+	proxyActiveConnections.WithLabelValues(clientALPN, connSource).Inc()
+	return &reportingConn{
+		Conn: conn,
+		closeOnce: sync.OnceValue(func() error {
+			proxyActiveConnections.WithLabelValues(clientALPN, connSource).Dec()
+			// Do not wrap.
+			return conn.Close()
+		}),
+	}
+}
+
+// NetConn returns the underlying net.Conn.
+func (c *reportingConn) NetConn() net.Conn {
+	return c.Conn
+}
+
+func (c *reportingConn) Close() error {
+	return c.closeOnce()
+}

lib/srv/alpnproxy/proxy.go

@@ -27,23 +27,27 @@ import (
 	"io"
 	"log/slog"
 	"net"
+	"slices"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
+	"github.com/prometheus/client_golang/prometheus"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/constants"
 	"github.com/gravitational/teleport/api/utils/pingconn"
 	"github.com/gravitational/teleport/lib/auth/authclient"
 	"github.com/gravitational/teleport/lib/authz"
 	"github.com/gravitational/teleport/lib/defaults"
+	"github.com/gravitational/teleport/lib/observability/metrics"
 	"github.com/gravitational/teleport/lib/srv/alpnproxy/common"
 	"github.com/gravitational/teleport/lib/srv/db/dbutils"
 	"github.com/gravitational/teleport/lib/tlsca"
 	"github.com/gravitational/teleport/lib/utils"
+	logutils "github.com/gravitational/teleport/lib/utils/log"
 )
 
 // ProxyConfig  is the configuration for an ALPN proxy server.
@@ -298,6 +302,14 @@ func New(cfg ProxyConfig) (*Proxy, error) {
 		return nil, trace.Wrap(err)
 	}
 
+	if err := metrics.RegisterPrometheusCollectors(
+		proxyConnectionsTotal,
+		proxyActiveConnections,
+		proxyConnectionErrorsTotal,
+	); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
 	return &Proxy{
 		cfg:                cfg,
 		log:                cfg.Log,
@@ -332,7 +344,7 @@ func (p *Proxy) Serve(ctx context.Context) error {
 			// For example in ReverseTunnel handles connection asynchronously and closing conn after
 			// service handler returned will break service logic.
 			// https://github.com/gravitational/teleport/blob/master/lib/sshutils/server.go#L397
-			if err := p.handleConn(ctx, clientConn, nil); err != nil {
+			if err := p.handleConn(ctx, clientConn, nil, common.ConnHandlerSourceListener); err != nil {
 				if cerr := clientConn.Close(); cerr != nil && !utils.IsOKNetworkError(cerr) {
 					p.log.WarnContext(ctx, "Failed to close client connection", "error", cerr)
 				}
@@ -372,8 +384,24 @@ type HandlerFuncWithInfo func(ctx context.Context, conn net.Conn, info Connectio
 //  5. For backward compatibility check RouteToDatabase identity field
 //     was set if yes forward to the generic TLS DB handler.
 //  6. Forward connection to the handler obtained in step 2.
-func (p *Proxy) handleConn(ctx context.Context, clientConn net.Conn, defaultOverride *tls.Config) error {
-	hello, conn, err := p.readHelloMessageWithoutTLSTermination(ctx, clientConn)
+func (p *Proxy) handleConn(ctx context.Context, clientConn net.Conn, defaultOverride *tls.Config, connSource common.ConnHandlerSource) (err error) {
+	var hello *tls.ClientHelloInfo
+	var conn net.Conn
+	defer func() {
+		if err != nil {
+			proxyConnectionErrorsTotal.WithLabelValues(
+				getRequestedALPNFromHello(hello),
+				string(connSource),
+			).Inc()
+		}
+	}()
+
+	// Attempt to read TLS hello. Always increment total counters even on failures.
+	hello, conn, err = p.readHelloMessageWithoutTLSTermination(ctx, clientConn, connSource)
+	proxyConnectionsTotal.WithLabelValues(
+		getRequestedALPNFromHello(hello),
+		string(connSource),
+	).Inc()
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -503,7 +531,7 @@ func (p *Proxy) getTLSConfig(desc *HandlerDecs, defaultOverride *tls.Config) *tl
 // readHelloMessageWithoutTLSTermination allows reading a ClientHelloInfo message without termination of
 // incoming TLS connection. After calling readHelloMessageWithoutTLSTermination function a returned
 // net.Conn should be used for further operation.
-func (p *Proxy) readHelloMessageWithoutTLSTermination(ctx context.Context, conn net.Conn) (*tls.ClientHelloInfo, net.Conn, error) {
+func (p *Proxy) readHelloMessageWithoutTLSTermination(ctx context.Context, conn net.Conn, connSource common.ConnHandlerSource) (*tls.ClientHelloInfo, net.Conn, error) {
 	buff := new(bytes.Buffer)
 	var hello *tls.ClientHelloInfo
 	tlsConn := tls.Server(readOnlyConn{reader: io.TeeReader(conn, buff)}, &tls.Config{
@@ -526,7 +554,13 @@ func (p *Proxy) readHelloMessageWithoutTLSTermination(ctx context.Context, conn
 	if err := conn.SetReadDeadline(time.Time{}); err != nil {
 		return nil, nil, trace.Wrap(err)
 	}
-	return hello, newBufferedConn(conn, buff), nil
+	return hello,
+		newReportingConn(
+			newBufferedConn(conn, buff),
+			getRequestedALPNFromHello(hello),
+			string(connSource),
+		),
+		nil
 }
 
 func (p *Proxy) handleDatabaseConnection(ctx context.Context, conn net.Conn, connInfo ConnectionInfo) error {
@@ -630,9 +664,20 @@ func (p *Proxy) Close() error {
 
 // MakeConnectionHandler creates a ConnectionHandler which provides a callback
 // to handle incoming connections by this ALPN proxy server.
-func (p *Proxy) MakeConnectionHandler(defaultOverride *tls.Config) ConnectionHandler {
+//
+// Note that some registered handlers are async. The input client connection
+// will be automatically closed upon handler errors.
+func (p *Proxy) MakeConnectionHandler(defaultOverride *tls.Config, connSource common.ConnHandlerSource) ConnectionHandler {
 	return func(ctx context.Context, conn net.Conn) error {
-		return p.handleConn(ctx, conn, defaultOverride)
+		if err := p.handleConn(ctx, conn, defaultOverride, connSource); err != nil {
+			// Make sure we close the connection on error.
+			if cerr := conn.Close(); cerr != nil && !utils.IsOKNetworkError(cerr) {
+				p.log.WarnContext(ctx, "Failed to close client connection", "error", cerr, "remote_addr", logutils.StringerAttr(conn.RemoteAddr()))
+			}
+			// Still return the error for caller to report/log.
+			return trace.Wrap(err)
+		}
+		return nil
 	}
 }
 
@@ -664,3 +709,56 @@ func isConnRemoteError(err error) bool {
 	var opErr *net.OpError
 	return errors.As(err, &opErr) && opErr.Op == "remote error"
 }
+
+// getRequestedALPNFromHello returns the primary requested ALPN by the client.
+// The server may not always terminate TLS so we won't know the negotiated
+// protocol. Here we just return the first supported ALPN from the client hello
+// and assumes that will likely be the one gets selected. If no supported ALPN
+// found, the function returns "unknown".
+func getRequestedALPNFromHello(hello *tls.ClientHelloInfo) string {
+	if hello == nil {
+		return "unknown"
+	}
+
+	for i := range hello.SupportedProtos {
+		protocol := common.Protocol(hello.SupportedProtos[i])
+		if strings.HasPrefix(string(protocol), string(common.ProtocolAuth)) {
+			protocol = common.ProtocolAuth
+		}
+
+		if slices.Contains(common.SupportedProtocols, protocol) {
+			return string(protocol)
+		}
+	}
+	return "unknown"
+}
+
+var (
+	proxyConnectionsTotal = prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Namespace: teleport.MetricNamespace,
+			Subsystem: "alpn_proxy",
+			Name:      "connections_total",
+			Help:      "Number of total connections handled by TLS routing proxy server.",
+		},
+		[]string{"alpn", "source"},
+	)
+	proxyActiveConnections = prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Namespace: teleport.MetricNamespace,
+			Subsystem: "alpn_proxy",
+			Name:      "active_connections",
+			Help:      "Number of active connections handled by TLS routing proxy server.",
+		},
+		[]string{"alpn", "source"},
+	)
+	proxyConnectionErrorsTotal = prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Namespace: teleport.MetricNamespace,
+			Subsystem: "alpn_proxy",
+			Name:      "connection_errors_total",
+			Help:      "Number of connection handler errors encountered by TLS routing proxy server.",
+		},
+		[]string{"alpn", "source"},
+	)
+)