Add tracing instrumentation for ssh clients/servers (#12434)

* Add tracing instrumentation for ssh clients/servers

Add tracing context to the existing ProxyHelloSignature to provide
span information across ssh connections. To add span context per
ssh session on top of new connections, the same tracing context is
passed in the first global request of the session.

In order to ensure that tracing context is pulled from and inserted
into the proper context.Context, some interfaces and methods were
changed to take one as the first argument.

api/client/contextdialer.go

@@ -25,6 +25,7 @@ import (
 	"github.com/gravitational/teleport/api/client/proxy"
 	"github.com/gravitational/teleport/api/client/webclient"
 	"github.com/gravitational/teleport/api/constants"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/utils/sshutils"
 
 	"github.com/gravitational/trace"
@@ -93,7 +94,7 @@ func newTunnelDialer(ssh ssh.ClientConfig, keepAlivePeriod, dialTimeout time.Dur
 			return nil, trace.Wrap(err)
 		}
 
-		sconn, err := sshConnect(conn, ssh, dialTimeout, addr)
+		sconn, err := sshConnect(ctx, conn, ssh, dialTimeout, addr)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
@@ -133,7 +134,7 @@ func newTLSRoutingTunnelDialer(ssh ssh.ClientConfig, keepAlivePeriod, dialTimeou
 			return nil, trace.Wrap(err)
 		}
 
-		sconn, err := sshConnect(tlsConn, ssh, dialTimeout, tunnelAddr)
+		sconn, err := sshConnect(ctx, tlsConn, ssh, dialTimeout, tunnelAddr)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
@@ -142,9 +143,9 @@ func newTLSRoutingTunnelDialer(ssh ssh.ClientConfig, keepAlivePeriod, dialTimeou
 }
 
 // sshConnect upgrades the underling connection to ssh and connects to the Auth service.
-func sshConnect(conn net.Conn, ssh ssh.ClientConfig, dialTimeout time.Duration, addr string) (net.Conn, error) {
+func sshConnect(ctx context.Context, conn net.Conn, ssh ssh.ClientConfig, dialTimeout time.Duration, addr string) (net.Conn, error) {
 	ssh.Timeout = dialTimeout
-	sconn, err := sshutils.NewClientConnWithDeadline(conn, addr, &ssh)
+	sconn, err := tracessh.NewClientConnWithDeadline(ctx, conn, addr, &ssh)
 	if err != nil {
 		return nil, trace.NewAggregate(err, conn.Close())
 	}

api/go.mod

@@ -12,6 +12,8 @@ require (
 	github.com/stretchr/testify v1.7.1
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.31.0
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.31.0
+	go.opentelemetry.io/otel v1.6.1
+	go.opentelemetry.io/otel/trace v1.6.1
 	golang.org/x/crypto v0.0.0-20220126234351-aa10faf2a1f8
 	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd
 	google.golang.org/grpc v1.45.0

api/observability/tracing/ssh/ssh.go

@@ -0,0 +1,121 @@
+// Copyright 2022 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ssh
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/gravitational/teleport/api/observability/tracing"
+	"github.com/gravitational/teleport/api/utils/sshutils"
+
+	"github.com/gravitational/trace"
+	log "github.com/sirupsen/logrus"
+	oteltrace "go.opentelemetry.io/otel/trace"
+	"golang.org/x/crypto/ssh"
+)
+
+const (
+	// TracingRequest is sent by clients to server to pass along tracing context.
+	TracingRequest = "tracing@goteleport.com"
+)
+
+// Client is a wrapper around ssh.Client that adds tracing support.
+type Client struct {
+	*ssh.Client
+}
+
+// NewClient creates a new Client.
+func NewClient(c ssh.Conn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request) *Client {
+	return &Client{Client: ssh.NewClient(c, chans, reqs)}
+}
+
+// NewSession creates a new SSH session that is passed tracing context so that spans may be correlated
+// properly over the ssh connection.
+func (c *Client) NewSession(ctx context.Context) (*ssh.Session, error) {
+	session, err := c.Client.NewSession()
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	span := oteltrace.SpanFromContext(ctx)
+	if !span.IsRecording() {
+		return session, nil
+	}
+
+	traceCtx := tracing.PropagationContextFromContext(ctx)
+	if len(traceCtx) == 0 {
+		return session, nil
+	}
+
+	payload, err := json.Marshal(traceCtx)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if _, err := session.SendRequest(TracingRequest, false, payload); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return session, nil
+}
+
+// NewClientConn creates a new SSH client connection that is passed tracing context so that spans may be correlated
+// properly over the ssh connection.
+func NewClientConn(ctx context.Context, conn net.Conn, addr string, config *ssh.ClientConfig) (ssh.Conn, <-chan ssh.NewChannel, <-chan *ssh.Request, error) {
+	hp := &sshutils.HandshakePayload{
+		TracingContext: tracing.PropagationContextFromContext(ctx),
+	}
+
+	if len(hp.TracingContext) > 0 {
+		payloadJSON, err := json.Marshal(hp)
+		if err == nil {
+			payload := fmt.Sprintf("%s%s\x00", sshutils.ProxyHelloSignature, payloadJSON)
+			_, err = conn.Write([]byte(payload))
+			if err != nil {
+				log.WithError(err).Warnf("Failed to pass along tracing context to proxy %v", addr)
+			}
+		}
+	}
+
+	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
+	if err != nil {
+		return nil, nil, nil, trace.Wrap(err)
+	}
+
+	return c, chans, reqs, nil
+}
+
+// NewClientConnWithDeadline establishes new client connection with specified deadline
+func NewClientConnWithDeadline(ctx context.Context, conn net.Conn, addr string, config *ssh.ClientConfig) (*Client, error) {
+	if config.Timeout > 0 {
+		if err := conn.SetReadDeadline(time.Now().Add(config.Timeout)); err != nil {
+			return nil, trace.Wrap(err)
+		}
+	}
+	c, chans, reqs, err := NewClientConn(ctx, conn, addr, config)
+	if err != nil {
+		return nil, err
+	}
+	if config.Timeout > 0 {
+		if err := conn.SetReadDeadline(time.Time{}); err != nil {
+			return nil, trace.Wrap(err)
+		}
+	}
+	return NewClient(c, chans, reqs), nil
+}

api/observability/tracing/tracing.go

@@ -0,0 +1,39 @@
+// Copyright 2022 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tracing
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+)
+
+// PropagationContext contains tracing information to be passed across service boundaries
+type PropagationContext map[string]string
+
+// PropagationContextFromContext creates a PropagationContext from the given context.Context. If the context
+// does not contain any tracing information, the PropagationContext will be empty.
+func PropagationContextFromContext(ctx context.Context) PropagationContext {
+	carrier := propagation.MapCarrier{}
+	otel.GetTextMapPropagator().Inject(ctx, &carrier)
+	return PropagationContext(carrier)
+}
+
+// WithPropagationContext injects any tracing information from the given PropagationContext into the
+// given context.Context.
+func WithPropagationContext(ctx context.Context, pc PropagationContext) context.Context {
+	return otel.GetTextMapPropagator().Extract(ctx, propagation.MapCarrier(pc))
+}

api/utils/sshutils/conn.go

@@ -21,30 +21,14 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"net"
-	"time"
 
 	"github.com/gravitational/teleport/api/constants"
 	"github.com/gravitational/teleport/api/types"
+
 	"github.com/gravitational/trace"
 	"golang.org/x/crypto/ssh"
 )
 
-// NewClientConnWithDeadline establishes new client connection with specified deadline
-func NewClientConnWithDeadline(conn net.Conn, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
-	if config.Timeout > 0 {
-		conn.SetReadDeadline(time.Now().Add(config.Timeout))
-	}
-	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
-	if err != nil {
-		return nil, err
-	}
-	if config.Timeout > 0 {
-		conn.SetReadDeadline(time.Time{})
-	}
-	return ssh.NewClient(c, chans, reqs), nil
-}
-
 // ConnectProxyTransport opens a channel over the remote tunnel and connects
 // to the requested host.
 func ConnectProxyTransport(sconn ssh.Conn, req *DialReq, exclusive bool) (*ChConn, bool, error) {

api/utils/sshutils/ssh.go

@@ -27,12 +27,30 @@ import (
 
 	"github.com/gravitational/teleport/api/constants"
 	"github.com/gravitational/teleport/api/defaults"
-
 	"github.com/gravitational/trace"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 )
 
+const (
+	// ProxyHelloSignature is a string which Teleport proxy will send
+	// right after the initial SSH "handshake/version" message if it detects
+	// talking to a Teleport server.
+	ProxyHelloSignature = "Teleport-Proxy"
+)
+
+// HandshakePayload structure is sent as a JSON blob by the teleport
+// proxy to every SSH server who identifies itself as Teleport server
+//
+// It allows teleport proxies to communicate additional data to server
+type HandshakePayload struct {
+	// ClientAddr is the IP address of the remote client
+	ClientAddr string `json:"clientAddr,omitempty"`
+	// TracingContext contains tracing information so that spans can be correlated
+	// across ssh boundaries
+	TracingContext map[string]string `json:"tracingContext,omitempty"`
+}
+
 // ParseCertificate parses an SSH certificate from the authorized_keys format.
 func ParseCertificate(buf []byte) (*ssh.Certificate, error) {
 	k, _, _, _, err := ssh.ParseAuthorizedKey(buf)

go.mod

@@ -93,6 +93,7 @@ require (
 	go.mozilla.org/pkcs7 v0.0.0-20210826202110-33d05740a352
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.31.0
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.31.0
+	go.opentelemetry.io/otel v1.6.1
 	go.uber.org/atomic v1.7.0
 	golang.org/x/crypto v0.0.0-20220126234351-aa10faf2a1f8
 	golang.org/x/mod v0.4.2
@@ -245,7 +246,6 @@ require (
 	github.com/yuin/gopher-lua v0.0.0-20200816102855-ee81675732da // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.1 // indirect
 	go.opencensus.io v0.23.0 // indirect
-	go.opentelemetry.io/otel v1.6.1 // indirect
 	go.opentelemetry.io/otel/metric v0.28.0 // indirect
 	go.opentelemetry.io/otel/trace v1.6.1 // indirect
 	go.starlark.net v0.0.0-20200306205701-8dd3e2ee1dd5 // indirect

integration/integration_test.go

@@ -45,6 +45,7 @@ import (
 
 	"github.com/gravitational/teleport"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/profile"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
@@ -4516,7 +4517,7 @@ func testWindowChange(t *testing.T, suite *integrationTestSuite) {
 		cl.Stdin = personB
 
 		// Change the size of the window immediately after it is created.
-		cl.OnShellCreated = func(s *ssh.Session, c *ssh.Client, terminal io.ReadWriteCloser) (exit bool, err error) {
+		cl.OnShellCreated = func(s *ssh.Session, c *tracessh.Client, terminal io.ReadWriteCloser) (exit bool, err error) {
 			err = s.WindowChange(48, 160)
 			if err != nil {
 				return true, trace.Wrap(err)

lib/auth/init_test.go

@@ -948,7 +948,7 @@ func TestIdentityChecker(t *testing.T) {
 			require.NoError(t, err)
 
 			dialer := proxy.DialerFromEnvironment(sshServer.Addr())
-			sconn, err := dialer.Dial("tcp", sshServer.Addr(), sshClientConfig)
+			sconn, err := dialer.Dial(ctx, "tcp", sshServer.Addr(), sshClientConfig)
 			if test.err {
 				require.Error(t, err)
 			} else {

lib/client/api.go

@@ -39,14 +39,12 @@ import (
 	"time"
 	"unicode/utf8"
 
-	"golang.org/x/crypto/ssh"
-	"golang.org/x/crypto/ssh/agent"
-
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/client/proto"
 	"github.com/gravitational/teleport/api/client/webclient"
 	"github.com/gravitational/teleport/api/constants"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/profile"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/api/types/wrappers"
@@ -70,11 +68,12 @@ import (
 	"github.com/gravitational/teleport/lib/utils/prompt"
 	"github.com/gravitational/teleport/lib/utils/proxy"
 
-	"github.com/gravitational/trace"
-
 	"github.com/duo-labs/webauthn/protocol"
+	"github.com/gravitational/trace"
 	"github.com/jonboulle/clockwork"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
 )
 
 const (
@@ -1348,7 +1347,7 @@ type TeleportClient struct {
 // hasn't begun yet.
 //
 // It allows clients to cancel SSH action
-type ShellCreatedCallback func(s *ssh.Session, c *ssh.Client, terminal io.ReadWriteCloser) (exit bool, err error)
+type ShellCreatedCallback func(s *ssh.Session, c *tracessh.Client, terminal io.ReadWriteCloser) (exit bool, err error)
 
 // NewClient creates a TeleportClient object and fully configures it
 func NewClient(c *Config) (tc *TeleportClient, err error) {
@@ -1691,7 +1690,7 @@ func (tc *TeleportClient) SSH(ctx context.Context, command []string, runLocally
 		return trace.Wrap(err)
 	}
 	defer proxyClient.Close()
-	siteInfo, err := proxyClient.currentCluster()
+	siteInfo, err := proxyClient.currentCluster(ctx)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -1956,7 +1955,7 @@ func (tc *TeleportClient) ExecuteSCP(ctx context.Context, cmd scp.Command) (err
 	}
 	defer proxyClient.Close()
 
-	clusterInfo, err := proxyClient.currentCluster()
+	clusterInfo, err := proxyClient.currentCluster(ctx)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -2025,7 +2024,7 @@ func (tc *TeleportClient) SCP(ctx context.Context, args []string, port int, flag
 	// helper function connects to the src/target node:
 	connectToNode := func(addr, hostLogin string) (*NodeClient, error) {
 		// determine which cluster we're connecting to:
-		siteInfo, err := proxyClient.currentCluster()
+		siteInfo, err := proxyClient.currentCluster(ctx)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
@@ -2525,7 +2524,7 @@ func (tc *TeleportClient) connectToProxy(ctx context.Context) (*ProxyClient, err
 //    In case of proxy web address check if the proxy supports TLS Routing and connect to the proxy with TLSWrapper
 // 3) Dial sshProxyAddr with raw SSH Dialer where sshProxyAddress is proxy ssh address or JumpHost address if
 //    JumpHost address was provided.
-func makeProxySSHClient(ctx context.Context, tc *TeleportClient, sshConfig *ssh.ClientConfig) (*ssh.Client, error) {
+func makeProxySSHClient(ctx context.Context, tc *TeleportClient, sshConfig *ssh.ClientConfig) (*tracessh.Client, error) {
 	// Use TLS Routing dialer only if proxy support TLS Routing and JumpHost was not set.
 	if tc.Config.TLSRoutingEnabled && len(tc.JumpHosts) == 0 {
 		log.Infof("Connecting to proxy=%v login=%q using TLS Routing", tc.Config.WebProxyAddr, sshConfig.User)
@@ -2557,7 +2556,7 @@ func makeProxySSHClient(ctx context.Context, tc *TeleportClient, sshConfig *ssh.
 	}
 
 	log.Infof("Connecting to proxy=%v login=%q", sshProxyAddr, sshConfig.User)
-	client, err := makeProxySSHClientDirect(tc, sshConfig, sshProxyAddr)
+	client, err := makeProxySSHClientDirect(ctx, tc, sshConfig, sshProxyAddr)
 	if err != nil {
 		return nil, trace.Wrap(err, "failed to authenticate with proxy %v", sshProxyAddr)
 	}
@@ -2565,12 +2564,12 @@ func makeProxySSHClient(ctx context.Context, tc *TeleportClient, sshConfig *ssh.
 	return client, nil
 }
 
-func makeProxySSHClientDirect(tc *TeleportClient, sshConfig *ssh.ClientConfig, proxyAddr string) (*ssh.Client, error) {
+func makeProxySSHClientDirect(ctx context.Context, tc *TeleportClient, sshConfig *ssh.ClientConfig, proxyAddr string) (*tracessh.Client, error) {
 	dialer := proxy.DialerFromEnvironment(tc.Config.SSHProxyAddr)
-	return dialer.Dial("tcp", proxyAddr, sshConfig)
+	return dialer.Dial(ctx, "tcp", proxyAddr, sshConfig)
 }
 
-func makeProxySSHClientWithTLSWrapper(ctx context.Context, tc *TeleportClient, sshConfig *ssh.ClientConfig, proxyAddr string) (*ssh.Client, error) {
+func makeProxySSHClientWithTLSWrapper(ctx context.Context, tc *TeleportClient, sshConfig *ssh.ClientConfig, proxyAddr string) (*tracessh.Client, error) {
 	tlsConfig, err := tc.loadTLSConfig()
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -2578,7 +2577,7 @@ func makeProxySSHClientWithTLSWrapper(ctx context.Context, tc *TeleportClient, s
 
 	tlsConfig.NextProtos = []string{string(alpncommon.ProtocolProxySSH)}
 	dialer := proxy.DialerFromEnvironment(tc.Config.WebProxyAddr, proxy.WithALPNDialer(tlsConfig))
-	return dialer.Dial("tcp", proxyAddr, sshConfig)
+	return dialer.Dial(ctx, "tcp", proxyAddr, sshConfig)
 }
 
 func (tc *TeleportClient) rootClusterName() (string, error) {

lib/client/client.go

@@ -30,13 +30,11 @@ import (
 	"strings"
 	"time"
 
-	"golang.org/x/crypto/ssh"
-	"golang.org/x/crypto/ssh/agent"
-
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/client"
 	"github.com/gravitational/teleport/api/client/proto"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/auth"
 	"github.com/gravitational/teleport/lib/defaults"
@@ -45,16 +43,18 @@ import (
 	"github.com/gravitational/teleport/lib/sshutils/scp"
 	"github.com/gravitational/teleport/lib/utils"
 	"github.com/gravitational/teleport/lib/utils/socks"
-	"github.com/moby/term"
 
 	"github.com/gravitational/trace"
+	"github.com/moby/term"
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
 )
 
 // ProxyClient implements ssh client to a teleport proxy
 // It can provide list of nodes or connect to nodes
 type ProxyClient struct {
 	teleportClient  *TeleportClient
-	Client          *ssh.Client
+	Client          *tracessh.Client
 	hostLogin       string
 	proxyAddress    string
 	proxyPrincipal  string
@@ -68,7 +68,7 @@ type ProxyClient struct {
 // NodeClient can run shell and commands or upload and download files.
 type NodeClient struct {
 	Namespace string
-	Client    *ssh.Client
+	Client    *tracessh.Client
 	Proxy     *ProxyClient
 	TC        *TeleportClient
 	OnMFA     func()
@@ -77,8 +77,8 @@ type NodeClient struct {
 // GetSites returns list of the "sites" (AKA teleport clusters) connected to the proxy
 // Each site is returned as an instance of its auth server
 //
-func (proxy *ProxyClient) GetSites() ([]types.Site, error) {
-	proxySession, err := proxy.Client.NewSession()
+func (proxy *ProxyClient) GetSites(ctx context.Context) ([]types.Site, error) {
+	proxySession, err := proxy.Client.NewSession(ctx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -770,7 +770,7 @@ func (proxy *ProxyClient) ListResources(ctx context.Context, namespace, resource
 // and could be cached based on the access policy
 func (proxy *ProxyClient) CurrentClusterAccessPoint(ctx context.Context, quiet bool) (auth.ClientI, error) {
 	// get the current cluster:
-	cluster, err := proxy.currentCluster()
+	cluster, err := proxy.currentCluster(ctx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -796,7 +796,7 @@ func (proxy *ProxyClient) ClusterAccessPoint(ctx context.Context, clusterName st
 // if 'quiet' is set to true, no errors will be printed to stdout, otherwise
 // any connection errors are visible to a user.
 func (proxy *ProxyClient) ConnectToCurrentCluster(ctx context.Context, quiet bool) (auth.ClientI, error) {
-	cluster, err := proxy.currentCluster()
+	cluster, err := proxy.currentCluster(ctx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -978,7 +978,7 @@ func (proxy *ProxyClient) dialAuthServer(ctx context.Context, clusterName string
 		return nil, trace.Wrap(err)
 	}
 
-	proxySession, err := proxy.Client.NewSession()
+	proxySession, err := proxy.Client.NewSession(ctx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -1096,7 +1096,7 @@ func (proxy *ProxyClient) ConnectToNode(ctx context.Context, nodeAddress NodeAdd
 		return nil, trace.Wrap(err)
 	}
 
-	proxySession, err := proxy.Client.NewSession()
+	proxySession, err := proxy.Client.NewSession(ctx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -1129,7 +1129,7 @@ func (proxy *ProxyClient) ConnectToNode(ctx context.Context, nodeAddress NodeAdd
 		if proxy.teleportClient.localAgent == nil {
 			return nil, trace.BadParameter("cluster is in proxy recording mode and requires agent forwarding for connections, but no agent was initialized")
 		}
-		err = agent.ForwardToAgent(proxy.Client, proxy.teleportClient.localAgent.Agent)
+		err = agent.ForwardToAgent(proxy.Client.Client, proxy.teleportClient.localAgent.Agent)
 		if err != nil && !strings.Contains(err.Error(), "agent: already have handler for") {
 			return nil, trace.Wrap(err)
 		}
@@ -1181,10 +1181,8 @@ func (proxy *ProxyClient) ConnectToNode(ctx context.Context, nodeAddress NodeAdd
 	emptyCh := make(chan *ssh.Request)
 	close(emptyCh)
 
-	client := ssh.NewClient(conn, chans, emptyCh)
-
 	nc := &NodeClient{
-		Client:    client,
+		Client:    tracessh.NewClient(conn, chans, emptyCh),
 		Proxy:     proxy,
 		Namespace: apidefaults.Namespace,
 		TC:        proxy.teleportClient,
@@ -1223,7 +1221,7 @@ func (proxy *ProxyClient) PortForwardToNode(ctx context.Context, nodeAddress Nod
 		if proxy.teleportClient.localAgent == nil {
 			return nil, trace.BadParameter("cluster is in proxy recording mode and requires agent forwarding for connections, but no agent was initialized")
 		}
-		err = agent.ForwardToAgent(proxy.Client, proxy.teleportClient.localAgent.Agent)
+		err = agent.ForwardToAgent(proxy.Client.Client, proxy.teleportClient.localAgent.Agent)
 		if err != nil && !strings.Contains(err.Error(), "agent: already have handler for") {
 			return nil, trace.Wrap(err)
 		}
@@ -1253,10 +1251,8 @@ func (proxy *ProxyClient) PortForwardToNode(ctx context.Context, nodeAddress Nod
 	emptyCh := make(chan *ssh.Request)
 	close(emptyCh)
 
-	client := ssh.NewClient(conn, chans, emptyCh)
-
 	nc := &NodeClient{
-		Client:    client,
+		Client:    tracessh.NewClient(conn, chans, emptyCh),
 		Proxy:     proxy,
 		Namespace: apidefaults.Namespace,
 		TC:        proxy.teleportClient,
@@ -1365,7 +1361,7 @@ func (c *NodeClient) ExecuteSCP(ctx context.Context, cmd scp.Command) error {
 		return trace.Wrap(err)
 	}
 
-	s, err := c.Client.NewSession()
+	s, err := c.Client.NewSession(ctx)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -1620,8 +1616,8 @@ func (c *NodeClient) Close() error {
 }
 
 // currentCluster returns the connection to the API of the current cluster
-func (proxy *ProxyClient) currentCluster() (*types.Site, error) {
-	sites, err := proxy.GetSites()
+func (proxy *ProxyClient) currentCluster(ctx context.Context) (*types.Site, error) {
+	sites, err := proxy.GetSites(ctx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}

lib/client/client_test.go

@@ -25,6 +25,7 @@ import (
 	"strings"
 	"time"
 
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/lib/sshutils"
 	"github.com/gravitational/trace"
 	"golang.org/x/crypto/ssh"
@@ -185,8 +186,10 @@ func (s *ClientTestSuite) TestProxyConnection(c *check.C) {
 
 func (s *ClientTestSuite) TestListenAndForwardCancel(c *check.C) {
 	client := &NodeClient{
-		Client: &ssh.Client{
-			Conn: &fakeSSHConn{},
+		Client: &tracessh.Client{
+			Client: &ssh.Client{
+				Conn: &fakeSSHConn{},
+			},
 		},
 	}
 

lib/client/session.go

@@ -210,7 +210,7 @@ func (ns *NodeSession) regularSession(ctx context.Context, callback func(s *ssh.
 type interactiveCallback func(serverSession *ssh.Session, shell io.ReadWriteCloser) error
 
 func (ns *NodeSession) createServerSession(ctx context.Context) (*ssh.Session, error) {
-	sess, err := ns.nodeClient.Client.NewSession()
+	sess, err := ns.nodeClient.Client.NewSession(ctx)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -247,7 +247,7 @@ func (ns *NodeSession) createServerSession(ctx context.Context) (*ssh.Session, e
 
 	if targetAgent != nil {
 		log.Debugf("Forwarding Selected Key Agent")
-		err = agent.ForwardToAgent(ns.nodeClient.Client, targetAgent)
+		err = agent.ForwardToAgent(ns.nodeClient.Client.Client, targetAgent)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
@@ -494,7 +494,7 @@ func (ns *NodeSession) runShell(ctx context.Context, mode types.SessionParticipa
 		}
 		// call the client-supplied callback
 		if callback != nil {
-			exit, err := callback(s, ns.NodeClient().Client, shell)
+			exit, err := callback(s, ns.nodeClient.Client, shell)
 			if exit {
 				return trace.Wrap(err)
 			}

lib/client/x11_session.go

@@ -127,7 +127,7 @@ func (ns *NodeSession) setXAuthData(ctx context.Context, display x11.Display) er
 
 // serveX11Channels serves incoming X11 channels by starting X11 forwarding with the session.
 func (ns *NodeSession) serveX11Channels(ctx context.Context, sess *ssh.Session) error {
-	err := x11.ServeChannelRequests(ctx, ns.nodeClient.Client, func(ctx context.Context, nch ssh.NewChannel) {
+	err := x11.ServeChannelRequests(ctx, ns.nodeClient.Client.Client, func(ctx context.Context, nch ssh.NewChannel) {
 		if !ns.x11RefuseTime.IsZero() && time.Now().After(ns.x11RefuseTime) {
 			nch.Reject(ssh.Prohibited, "rejected X11 channel request after ForwardX11Timeout")
 			log.Warn("rejected X11 forwarding attempt after the ForwardX11Timeout")
@@ -195,7 +195,7 @@ func (ns *NodeSession) serveX11Channels(ctx context.Context, sess *ssh.Session)
 
 // rejectX11Channels rejects any incomign X11 channels for this node session.
 func (ns *NodeSession) rejectX11Channels(ctx context.Context) error {
-	err := x11.ServeChannelRequests(ctx, ns.nodeClient.Client, func(_ context.Context, nch ssh.NewChannel) {
+	err := x11.ServeChannelRequests(ctx, ns.nodeClient.Client.Client, func(_ context.Context, nch ssh.NewChannel) {
 		// According to RFC 4254, client "implementations MUST reject any X11 channel
 		// open requests if they have not requested X11 forwarding. Following openssh's
 		// example, we treat such a request as a break in attempt and warn the user.

lib/reversetunnel/agent.go

@@ -31,6 +31,7 @@ import (
 	"github.com/gravitational/teleport/api/client/webclient"
 	"github.com/gravitational/teleport/api/constants"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/types"
 	apisshutils "github.com/gravitational/teleport/api/utils/sshutils"
 	"github.com/gravitational/teleport/lib"
@@ -277,7 +278,7 @@ func (a *Agent) getReverseTunnelDetails() *reverseTunnelDetails {
 	return &pd
 }
 
-func (a *Agent) connect() (conn *ssh.Client, err error) {
+func (a *Agent) connect() (conn *tracessh.Client, err error) {
 	if a.reverseTunnelDetails == nil {
 		a.reverseTunnelDetails = a.getReverseTunnelDetails()
 	}
@@ -295,7 +296,7 @@ func (a *Agent) connect() (conn *ssh.Client, err error) {
 	for _, authMethod := range a.authMethods {
 		// Create a dialer (that respects HTTP proxies) and connect to remote host.
 		dialer := proxy.DialerFromEnvironment(a.Addr.Addr, opts...)
-		pconn, err := dialer.DialTimeout(a.Addr.AddrNetwork, a.Addr.Addr, apidefaults.DefaultDialTimeout)
+		pconn, err := dialer.DialTimeout(a.Context, a.Addr.AddrNetwork, a.Addr.Addr, apidefaults.DefaultDialTimeout)
 		if err != nil {
 			a.log.WithError(err).Debugf("Dial to %v failed.", a.Addr.Addr)
 			continue
@@ -316,7 +317,7 @@ func (a *Agent) connect() (conn *ssh.Client, err error) {
 
 		// Build a new client connection. This is done to get access to incoming
 		// global requests which dialer.Dial would not provide.
-		conn, chans, reqs, err := ssh.NewClientConn(pconn, a.Addr.Addr, &ssh.ClientConfig{
+		conn, chans, reqs, err := tracessh.NewClientConn(a.Context, pconn, a.Addr.Addr, &ssh.ClientConfig{
 			User:            a.Username,
 			Auth:            []ssh.AuthMethod{authMethod},
 			HostKeyCallback: callback,
@@ -332,7 +333,7 @@ func (a *Agent) connect() (conn *ssh.Client, err error) {
 		emptyCh := make(chan *ssh.Request)
 		close(emptyCh)
 
-		client := ssh.NewClient(conn, chans, emptyCh)
+		client := tracessh.NewClient(conn, chans, emptyCh)
 
 		// Start a goroutine to process global requests from the server.
 		go a.handleGlobalRequests(a.ctx, reqs)
@@ -460,7 +461,7 @@ const ConnectedEvent = "connected"
 // processRequests is a blocking function which runs in a loop sending heartbeats
 // to the given SSH connection and processes inbound requests from the
 // remote proxy
-func (a *Agent) processRequests(conn *ssh.Client) error {
+func (a *Agent) processRequests(conn *tracessh.Client) error {
 	netConfig, err := a.AccessPoint.GetClusterNetworkingConfig(a.ctx)
 	if err != nil {
 		return trace.Wrap(err)

lib/reversetunnel/emit_conn.go

@@ -23,8 +23,8 @@ import (
 	"sync"
 
 	apievents "github.com/gravitational/teleport/api/types/events"
+	"github.com/gravitational/teleport/api/utils/sshutils"
 	"github.com/gravitational/teleport/lib/events"
-	"github.com/gravitational/teleport/lib/sshutils"
 )
 
 func min(a, b int) int {

lib/reversetunnel/localsite.go

@@ -344,7 +344,7 @@ with the cluster.`
 
 	// If no tunnel connection was found, dial to the target host.
 	dialer := proxy.DialerFromEnvironment(params.To.String())
-	conn, directErr := dialer.DialTimeout(params.To.Network(), params.To.String(), apidefaults.DefaultDialTimeout)
+	conn, directErr := dialer.DialTimeout(s.srv.Context, params.To.Network(), params.To.String(), apidefaults.DefaultDialTimeout)
 	if directErr != nil {
 		directMsg := fmt.Sprintf(errorMessageTemplate, params.ConnType, dreq.Address, "direct dial", directErr)
 		s.log.WithError(directErr).WithField("address", params.To.String()).Debug("Error occurred while dialing directly.")

lib/reversetunnel/transport.go

@@ -105,7 +105,7 @@ func (t *TunnelAuthDialer) DialContext(ctx context.Context, _, _ string) (net.Co
 	}
 
 	dialer := proxy.DialerFromEnvironment(addr.Addr, opts...)
-	sconn, err := dialer.Dial(addr.AddrNetwork, addr.Addr, t.ClientConfig)
+	sconn, err := dialer.Dial(ctx, addr.AddrNetwork, addr.Addr, t.ClientConfig)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}

lib/srv/alpnproxy/local_proxy.go

@@ -32,6 +32,7 @@ import (
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/crypto/ssh/agent"
 
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/srv/alpnproxy/common"
 	"github.com/gravitational/teleport/lib/utils"
@@ -117,7 +118,7 @@ func NewLocalProxy(cfg LocalProxyConfig) (*LocalProxy, error) {
 
 // SSHProxy is equivalent of `ssh -o 'ForwardAgent yes' -p port  %r@host -s proxy:%h:%p` but established SSH
 // connection to RemoteProxyAddr is wrapped with TLS protocol.
-func (l *LocalProxy) SSHProxy(localAgent *client.LocalKeyAgent) error {
+func (l *LocalProxy) SSHProxy(ctx context.Context, localAgent *client.LocalKeyAgent) error {
 	if l.cfg.ClientTLSConfig == nil {
 		return trace.BadParameter("client TLS config is missing")
 	}
@@ -133,7 +134,7 @@ func (l *LocalProxy) SSHProxy(localAgent *client.LocalKeyAgent) error {
 	}
 	defer upstreamConn.Close()
 
-	client, err := makeSSHClient(upstreamConn, l.cfg.RemoteProxyAddr, &ssh.ClientConfig{
+	client, err := makeSSHClient(ctx, upstreamConn, l.cfg.RemoteProxyAddr, &ssh.ClientConfig{
 		User: l.cfg.SSHUser,
 		Auth: []ssh.AuthMethod{
 			ssh.PublicKeysCallback(localAgent.Signers),
@@ -145,13 +146,13 @@ func (l *LocalProxy) SSHProxy(localAgent *client.LocalKeyAgent) error {
 	}
 	defer client.Close()
 
-	sess, err := client.NewSession()
+	sess, err := client.NewSession(ctx)
 	if err != nil {
 		return trace.Wrap(err)
 	}
 	defer sess.Close()
 
-	err = agent.ForwardToAgent(client, localAgent)
+	err = agent.ForwardToAgent(client.Client, localAgent)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -177,12 +178,12 @@ func proxySubsystemName(userHost, cluster string) string {
 	return subsystem
 }
 
-func makeSSHClient(conn *tls.Conn, addr string, cfg *ssh.ClientConfig) (*ssh.Client, error) {
-	cc, chs, reqs, err := ssh.NewClientConn(conn, addr, cfg)
+func makeSSHClient(ctx context.Context, conn *tls.Conn, addr string, cfg *ssh.ClientConfig) (*tracessh.Client, error) {
+	cc, chs, reqs, err := tracessh.NewClientConn(ctx, conn, addr, cfg)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
-	return ssh.NewClient(cc, chs, reqs), nil
+	return tracessh.NewClient(cc, chs, reqs), nil
 }
 
 func proxySession(ctx context.Context, sess *ssh.Session) error {

lib/srv/ctx.go

@@ -31,6 +31,7 @@ import (
 	"golang.org/x/crypto/ssh"
 
 	"github.com/gravitational/teleport"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
 	apiutils "github.com/gravitational/teleport/api/utils"
@@ -274,7 +275,7 @@ type ServerContext struct {
 
 	// RemoteClient holds a SSH client to a remote server. Only used by the
 	// recording proxy.
-	RemoteClient *ssh.Client
+	RemoteClient *tracessh.Client
 
 	// RemoteSession holds a SSH session to a remote server. Only used by the
 	// recording proxy.

lib/srv/forward/sshserver.go

@@ -28,9 +28,9 @@ import (
 
 	"github.com/gravitational/teleport"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
-	apisshutils "github.com/gravitational/teleport/api/utils/sshutils"
 	"github.com/gravitational/teleport/lib/auth"
 	"github.com/gravitational/teleport/lib/bpf"
 	"github.com/gravitational/teleport/lib/events"
@@ -91,7 +91,7 @@ type Server struct {
 
 	// remoteClient exposes an API to SSH functionality like shells, port
 	// forwarding, subsystems.
-	remoteClient *ssh.Client
+	remoteClient *tracessh.Client
 
 	// connectionContext is used to construct ServerContext instances
 	// and supports registration of connection-scoped resource closers.
@@ -503,7 +503,7 @@ func (s *Server) Serve() {
 
 	// Connect and authenticate to the remote node.
 	s.log.Debugf("Creating remote connection to %v@%v", sconn.User(), s.clientConn.RemoteAddr().String())
-	s.remoteClient, err = s.newRemoteClient(sconn.User())
+	s.remoteClient, err = s.newRemoteClient(ctx, sconn.User())
 	if err != nil {
 		// Reject the connection with an error so the client doesn't hang then
 		// close the connection.
@@ -569,7 +569,7 @@ func (s *Server) Close() error {
 
 // newRemoteSession will create and return a *ssh.Client and *ssh.Session
 // with a remote host.
-func (s *Server) newRemoteClient(systemLogin string) (*ssh.Client, error) {
+func (s *Server) newRemoteClient(ctx context.Context, systemLogin string) (*tracessh.Client, error) {
 	// the proxy will use the agent that has been forwarded to it as the auth
 	// method when connecting to the remote host
 	if s.userAgent == nil {
@@ -596,7 +596,7 @@ func (s *Server) newRemoteClient(systemLogin string) (*ssh.Client, error) {
 	// the correct host. It must occur in the list of principals presented by
 	// the remote server.
 	dstAddr := net.JoinHostPort(s.address, "0")
-	client, err := apisshutils.NewClientConnWithDeadline(s.targetConn, dstAddr, clientConfig)
+	client, err := tracessh.NewClientConnWithDeadline(ctx, s.targetConn, dstAddr, clientConfig)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -810,7 +810,7 @@ func (s *Server) handleSessionChannel(ctx context.Context, nch ssh.NewChannel) {
 	// create the remote session channel before accepting the local
 	// channel request; this allows us to propagate the rejection
 	// reason/message in the event the channel is rejected.
-	remoteSession, err := s.remoteClient.NewSession()
+	remoteSession, err := s.remoteClient.NewSession(ctx)
 	if err != nil {
 		s.log.Warnf("Remote session open failed: %v", err)
 		reason, msg := ssh.ConnectionFailed, fmt.Sprintf("remote session open failed: %v", err)
@@ -912,7 +912,7 @@ func (s *Server) dispatch(ctx context.Context, ch ssh.Channel, req *ssh.Request,
 			// We ignore all SSH setenv requests for join-only principals.
 			// SSH will send them anyway but it seems fine to silently drop them.
 		case sshutils.SubsystemRequest:
-			return s.handleSubsystem(ch, req, scx)
+			return s.handleSubsystem(ctx, ch, req, scx)
 		default:
 			return trace.AccessDenied("attempted %v request in join-only mode", req.Type)
 		}
@@ -932,7 +932,7 @@ func (s *Server) dispatch(ctx context.Context, ch ssh.Channel, req *ssh.Request,
 	case sshutils.EnvRequest:
 		return s.handleEnv(ch, req, scx)
 	case sshutils.SubsystemRequest:
-		return s.handleSubsystem(ch, req, scx)
+		return s.handleSubsystem(ctx, ch, req, scx)
 	case sshutils.X11ForwardRequest:
 		return s.handleX11Forward(ctx, ch, req, scx)
 	case sshutils.AgentForwardRequest:
@@ -958,7 +958,7 @@ func (s *Server) handleAgentForward(ch ssh.Channel, req *ssh.Request, ctx *srv.S
 	}
 
 	// Route authentication requests to the agent that was forwarded to the proxy.
-	err = agent.ForwardToAgent(ctx.RemoteClient, s.userAgent)
+	err = agent.ForwardToAgent(ctx.RemoteClient.Client, s.userAgent)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -1062,7 +1062,7 @@ func (s *Server) handleX11Forward(ctx context.Context, ch ssh.Channel, req *ssh.
 		return trace.AccessDenied("X11 forwarding request denied by server")
 	}
 
-	err = x11.ServeChannelRequests(ctx, s.remoteClient, s.handleX11ChannelRequest)
+	err = x11.ServeChannelRequests(ctx, s.remoteClient.Client, s.handleX11ChannelRequest)
 	if err != nil {
 		return trace.Wrap(err)
 	}
@@ -1070,16 +1070,16 @@ func (s *Server) handleX11Forward(ctx context.Context, ch ssh.Channel, req *ssh.
 	return nil
 }
 
-func (s *Server) handleSubsystem(ch ssh.Channel, req *ssh.Request, ctx *srv.ServerContext) error {
-	subsystem, err := parseSubsystemRequest(req, ctx)
+func (s *Server) handleSubsystem(ctx context.Context, ch ssh.Channel, req *ssh.Request, serverContext *srv.ServerContext) error {
+	subsystem, err := parseSubsystemRequest(req, serverContext)
 	if err != nil {
 		return trace.Wrap(err)
 	}
 
 	// start the requested subsystem, if it fails to start return result right away
-	err = subsystem.Start(ch)
+	err = subsystem.Start(ctx, ch)
 	if err != nil {
-		ctx.SendSubsystemResult(srv.SubsystemResult{
+		serverContext.SendSubsystemResult(srv.SubsystemResult{
 			Name: subsystem.subsytemName,
 			Err:  trace.Wrap(err),
 		})
@@ -1089,7 +1089,7 @@ func (s *Server) handleSubsystem(ch ssh.Channel, req *ssh.Request, ctx *srv.Serv
 	// wait for the subsystem to finish and return that result
 	go func() {
 		err := subsystem.Wait()
-		ctx.SendSubsystemResult(srv.SubsystemResult{
+		serverContext.SendSubsystemResult(srv.SubsystemResult{
 			Name: subsystem.subsytemName,
 			Err:  trace.Wrap(err),
 		})

lib/srv/forward/subsystem.go

@@ -58,8 +58,8 @@ func parseRemoteSubsystem(ctx context.Context, subsytemName string, serverContex
 	}
 }
 
-// Start will begin execution of the remote subsytem on the passed in channel.
-func (r *remoteSubsystem) Start(channel ssh.Channel) error {
+// Start will begin execution of the remote subsystem on the passed in channel.
+func (r *remoteSubsystem) Start(ctx context.Context, channel ssh.Channel) error {
 	session := r.serverContext.RemoteSession
 
 	stdout, err := session.StdoutPipe()

lib/srv/regular/proxy.go

@@ -18,6 +18,7 @@ package regular
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -30,8 +31,10 @@ import (
 
 	"github.com/gravitational/teleport"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	"github.com/gravitational/teleport/api/observability/tracing"
 	"github.com/gravitational/teleport/api/types"
 	apiutils "github.com/gravitational/teleport/api/utils"
+	apisshutils "github.com/gravitational/teleport/api/utils/sshutils"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/reversetunnel"
 	"github.com/gravitational/teleport/lib/services"
@@ -224,7 +227,7 @@ func (t *proxySubsys) String() string {
 
 // Start is called by Golang's ssh when it needs to engage this sybsystem (typically to establish
 // a mapping connection between a client & remote node we're proxying to)
-func (t *proxySubsys) Start(sconn *ssh.ServerConn, ch ssh.Channel, req *ssh.Request, ctx *srv.ServerContext) error {
+func (t *proxySubsys) Start(ctx context.Context, sconn *ssh.ServerConn, ch ssh.Channel, req *ssh.Request, serverContext *srv.ServerContext) error {
 	// once we start the connection, update logger to include component fields
 	t.log = logrus.WithFields(logrus.Fields{
 		trace.Component: teleport.ComponentSubsystemProxy,
@@ -238,12 +241,12 @@ func (t *proxySubsys) Start(sconn *ssh.ServerConn, ch ssh.Channel, req *ssh.Requ
 	var (
 		site       reversetunnel.RemoteSite
 		err        error
-		tunnel     = t.srv.tunnelWithRoles(ctx)
+		tunnel     = t.srv.tunnelWithRoles(serverContext)
 		clientAddr = sconn.RemoteAddr()
 	)
 	// did the client pass us a true client IP ahead of time via an environment variable?
 	// (usually the web client would do that)
-	trueClientIP, ok := ctx.GetEnv(sshutils.TrueClientAddrVar)
+	trueClientIP, ok := serverContext.GetEnv(sshutils.TrueClientAddrVar)
 	if ok {
 		a, err := utils.ParseAddr(trueClientIP)
 		if err == nil {
@@ -277,7 +280,7 @@ func (t *proxySubsys) Start(sconn *ssh.ServerConn, ch ssh.Channel, req *ssh.Requ
 		return t.proxyToHost(ctx, site, clientAddr, ch)
 	}
 	// connect to a site's auth server:
-	return t.proxyToSite(ctx, site, clientAddr, ch)
+	return t.proxyToSite(serverContext, site, clientAddr, ch)
 }
 
 // proxyToSite establishes a proxy connection from the connected SSH client to the
@@ -314,7 +317,7 @@ func (t *proxySubsys) proxyToSite(
 // proxyToHost establishes a proxy connection from the connected SSH client to the
 // requested remote node (t.host:t.port) via the given site
 func (t *proxySubsys) proxyToHost(
-	ctx *srv.ServerContext, site reversetunnel.RemoteSite, remoteAddr net.Addr, ch ssh.Channel) error {
+	ctx context.Context, site reversetunnel.RemoteSite, remoteAddr net.Addr, ch ssh.Channel) error {
 	//
 	// first, lets fetch a list of servers at the given site. this allows us to
 	// match the given "host name" against node configuration (their 'nodename' setting)
@@ -333,7 +336,7 @@ func (t *proxySubsys) proxyToHost(
 	if site.GetName() == localCluster.GetName() {
 		nodeWatcher = t.srv.nodeWatcher
 
-		cfg, err := t.srv.authService.GetClusterNetworkingConfig(ctx.CancelContext())
+		cfg, err := t.srv.authService.GetClusterNetworkingConfig(ctx)
 		if err != nil {
 			t.log.Warn(err)
 		} else {
@@ -352,7 +355,7 @@ func (t *proxySubsys) proxyToHost(
 				nodeWatcher = watcher
 			}
 
-			cfg, err := siteClient.GetClusterNetworkingConfig(ctx.CancelContext())
+			cfg, err := siteClient.GetClusterNetworkingConfig(ctx)
 			if err != nil {
 				t.log.Warn(err)
 			} else {
@@ -430,7 +433,7 @@ func (t *proxySubsys) proxyToHost(
 
 	// this custom SSH handshake allows SSH proxy to relay the client's IP
 	// address to the SSH server
-	t.doHandshake(remoteAddr, ch, conn)
+	t.doHandshake(ctx, remoteAddr, ch, conn)
 
 	proxiedSessions.Inc()
 	go func() {
@@ -557,8 +560,8 @@ func (t *proxySubsys) Wait() error {
 
 // doHandshake allows a proxy server to send additional information (client IP)
 // to an SSH server before establishing a bridge
-func (t *proxySubsys) doHandshake(clientAddr net.Addr, clientConn io.ReadWriter, serverConn io.ReadWriter) {
-	// on behalf of a client ask the server for it's version:
+func (t *proxySubsys) doHandshake(ctx context.Context, clientAddr net.Addr, clientConn io.ReadWriter, serverConn io.ReadWriter) {
+	// on behalf of a client ask the server for its version:
 	buff := make([]byte, sshutils.MaxVersionStringBytes)
 	n, err := serverConn.Read(buff)
 	if err != nil {
@@ -572,22 +575,23 @@ func (t *proxySubsys) doHandshake(clientAddr net.Addr, clientConn io.ReadWriter,
 	if bytes.HasPrefix(buff, []byte(sshutils.SSHVersionPrefix)) {
 		// if we're connecting to a Teleport SSH server, send our own "handshake payload"
 		// message, along with a client's IP:
-		hp := &sshutils.HandshakePayload{
-			ClientAddr: clientAddr.String(),
+		hp := &apisshutils.HandshakePayload{
+			ClientAddr:     clientAddr.String(),
+			TracingContext: tracing.PropagationContextFromContext(ctx),
 		}
 		payloadJSON, err := json.Marshal(hp)
 		if err != nil {
 			t.log.Error(err)
 		} else {
-			// send a JSON payload sandwitched between 'teleport proxy signature' and 0x00:
-			payload := fmt.Sprintf("%s%s\x00", sshutils.ProxyHelloSignature, payloadJSON)
+			// send a JSON payload sandwiched between 'teleport proxy signature' and 0x00:
+			payload := fmt.Sprintf("%s%s\x00", apisshutils.ProxyHelloSignature, payloadJSON)
 			_, err = serverConn.Write([]byte(payload))
 			if err != nil {
 				t.log.Error(err)
 			}
 		}
 	}
-	// forwrd server's response to the client:
+	// forward server's response to the client:
 	_, err = clientConn.Write(buff)
 	if err != nil {
 		t.log.Error(err)

lib/srv/regular/sites.go

@@ -17,6 +17,7 @@ limitations under the License.
 package regular
 
 import (
+	"context"
 	"encoding/json"
 
 	"golang.org/x/crypto/ssh"
@@ -50,9 +51,9 @@ func (t *proxySitesSubsys) Wait() error {
 
 // Start serves a request for "proxysites" custom SSH subsystem. It builds an array of
 // service.Site structures, and writes it serialized as JSON back to the SSH client
-func (t *proxySitesSubsys) Start(sconn *ssh.ServerConn, ch ssh.Channel, req *ssh.Request, ctx *srv.ServerContext) error {
-	log.Debugf("proxysites.start(%v)", ctx)
-	remoteSites, err := t.srv.tunnelWithRoles(ctx).GetSites()
+func (t *proxySitesSubsys) Start(ctx context.Context, sconn *ssh.ServerConn, ch ssh.Channel, req *ssh.Request, serverContext *srv.ServerContext) error {
+	log.Debugf("proxysites.start(%v)", serverContext)
+	remoteSites, err := t.srv.tunnelWithRoles(serverContext).GetSites()
 	if err != nil {
 		return trace.Wrap(err)
 	}

lib/srv/regular/sshserver.go

@@ -20,6 +20,7 @@ package regular
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"net"
@@ -30,10 +31,10 @@ import (
 	"strings"
 	"sync"
 
-	"golang.org/x/crypto/ssh"
-
 	"github.com/gravitational/teleport"
 	apidefaults "github.com/gravitational/teleport/api/defaults"
+	"github.com/gravitational/teleport/api/observability/tracing"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
 	"github.com/gravitational/teleport/lib/auth"
@@ -54,10 +55,10 @@ import (
 	"github.com/gravitational/teleport/lib/utils"
 
 	"github.com/gravitational/trace"
-
 	"github.com/jonboulle/clockwork"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/crypto/ssh"
 )
 
 var (
@@ -916,7 +917,7 @@ func (s *Server) serveAgent(ctx *srv.ServerContext) error {
 // req.Reply(false, nil).
 //
 // For more details: https://tools.ietf.org/html/rfc4254.html#page-4
-func (s *Server) HandleRequest(r *ssh.Request) {
+func (s *Server) HandleRequest(ctx context.Context, r *ssh.Request) {
 	switch r.Type {
 	case teleport.KeepAliveReqType:
 		s.handleKeepAlive(r)
@@ -1380,7 +1381,21 @@ func (s *Server) handleSessionRequests(ctx context.Context, ccx *sshutils.Connec
 				scx.Debugf("Client %v disconnected.", scx.ServerConn.RemoteAddr())
 				return
 			}
-			if err := s.dispatch(ch, req, scx); err != nil {
+
+			// handle the tracing request inline here to update the context for this session.
+			// this should only be requested once, right after the session has been created.
+			if req.Type == tracessh.TracingRequest {
+				var traceCtx tracing.PropagationContext
+				if err := json.Unmarshal(req.Payload, &traceCtx); err != nil {
+					scx.WithError(err).Error("Failed to unmarshal tracing request.")
+					continue
+				}
+
+				ctx = tracing.WithPropagationContext(ctx, traceCtx)
+				continue
+			}
+
+			if err := s.dispatch(ctx, ch, req, scx); err != nil {
 				s.replyError(ch, req, err)
 				return
 			}
@@ -1407,24 +1422,24 @@ func (s *Server) handleSessionRequests(ctx context.Context, ccx *sshutils.Connec
 	}
 }
 
-// dispatch receives an SSH request for a subsystem and disptaches the request to the
+// dispatch receives an SSH request for a subsystem and dispatches the request to the
 // appropriate subsystem implementation
-func (s *Server) dispatch(ch ssh.Channel, req *ssh.Request, ctx *srv.ServerContext) error {
-	ctx.Debugf("Handling request %v, want reply %v.", req.Type, req.WantReply)
+func (s *Server) dispatch(ctx context.Context, ch ssh.Channel, req *ssh.Request, serverContext *srv.ServerContext) error {
+	serverContext.Debugf("Handling request %v, want reply %v.", req.Type, req.WantReply)
 
 	// If this SSH server is configured to only proxy, we do not support anything
 	// other than our own custom "subsystems" and environment manipulation.
 	if s.proxyMode {
 		switch req.Type {
 		case sshutils.SubsystemRequest:
-			return s.handleSubsystem(ch, req, ctx)
+			return s.handleSubsystem(ctx, ch, req, serverContext)
 		case sshutils.EnvRequest:
 			// we currently ignore setting any environment variables via SSH for security purposes
-			return s.handleEnv(ch, req, ctx)
+			return s.handleEnv(ch, req, serverContext)
 		case sshutils.AgentForwardRequest:
 			// process agent forwarding, but we will only forward agent to proxy in
 			// recording proxy mode.
-			err := s.handleAgentForwardProxy(req, ctx)
+			err := s.handleAgentForwardProxy(req, serverContext)
 			if err != nil {
 				log.Warn(err)
 			}
@@ -1444,21 +1459,21 @@ func (s *Server) dispatch(ch ssh.Channel, req *ssh.Request, ctx *srv.ServerConte
 
 	// Certs with a join-only principal can only use a
 	// subset of all the possible request types.
-	if ctx.JoinOnly {
+	if serverContext.JoinOnly {
 		switch req.Type {
 		case sshutils.PTYRequest:
-			return s.termHandlers.HandlePTYReq(ch, req, ctx)
+			return s.termHandlers.HandlePTYReq(ch, req, serverContext)
 		case sshutils.ShellRequest:
-			return s.termHandlers.HandleShell(ch, req, ctx)
+			return s.termHandlers.HandleShell(ch, req, serverContext)
 		case sshutils.WindowChangeRequest:
-			return s.termHandlers.HandleWinChange(ch, req, ctx)
+			return s.termHandlers.HandleWinChange(ch, req, serverContext)
 		case teleport.ForceTerminateRequest:
-			return s.termHandlers.HandleForceTerminate(ch, req, ctx)
+			return s.termHandlers.HandleForceTerminate(ch, req, serverContext)
 		case sshutils.EnvRequest:
 			// We ignore all SSH setenv requests for join-only principals.
 			// SSH will send them anyway but it seems fine to silently drop them.
 		case sshutils.SubsystemRequest:
-			return s.handleSubsystem(ch, req, ctx)
+			return s.handleSubsystem(ctx, ch, req, serverContext)
 		default:
 			return trace.AccessDenied("attempted %v request in join-only mode", req.Type)
 		}
@@ -1466,23 +1481,23 @@ func (s *Server) dispatch(ch ssh.Channel, req *ssh.Request, ctx *srv.ServerConte
 
 	switch req.Type {
 	case sshutils.ExecRequest:
-		return s.termHandlers.HandleExec(ch, req, ctx)
+		return s.termHandlers.HandleExec(ch, req, serverContext)
 	case sshutils.PTYRequest:
-		return s.termHandlers.HandlePTYReq(ch, req, ctx)
+		return s.termHandlers.HandlePTYReq(ch, req, serverContext)
 	case sshutils.ShellRequest:
-		return s.termHandlers.HandleShell(ch, req, ctx)
+		return s.termHandlers.HandleShell(ch, req, serverContext)
 	case sshutils.WindowChangeRequest:
-		return s.termHandlers.HandleWinChange(ch, req, ctx)
+		return s.termHandlers.HandleWinChange(ch, req, serverContext)
 	case teleport.ForceTerminateRequest:
-		return s.termHandlers.HandleForceTerminate(ch, req, ctx)
+		return s.termHandlers.HandleForceTerminate(ch, req, serverContext)
 	case sshutils.EnvRequest:
-		return s.handleEnv(ch, req, ctx)
+		return s.handleEnv(ch, req, serverContext)
 	case sshutils.SubsystemRequest:
 		// subsystems are SSH subsystems defined in http://tools.ietf.org/html/rfc4254 6.6
 		// they are in essence SSH session extensions, allowing to implement new SSH commands
-		return s.handleSubsystem(ch, req, ctx)
+		return s.handleSubsystem(ctx, ch, req, serverContext)
 	case sshutils.X11ForwardRequest:
-		return s.handleX11Forward(ch, req, ctx)
+		return s.handleX11Forward(ch, req, serverContext)
 	case sshutils.AgentForwardRequest:
 		// This happens when SSH client has agent forwarding enabled, in this case
 		// client sends a special request, in return SSH server opens new channel
@@ -1494,7 +1509,7 @@ func (s *Server) dispatch(ch ssh.Channel, req *ssh.Request, ctx *srv.ServerConte
 		// to maintain interoperability with OpenSSH, agent forwarding requests
 		// should never fail, all errors should be logged and we should continue
 		// processing requests.
-		err := s.handleAgentForwardNode(req, ctx)
+		err := s.handleAgentForwardNode(req, serverContext)
 		if err != nil {
 			log.Warn(err)
 		}
@@ -1615,24 +1630,24 @@ func (s *Server) handleX11Forward(ch ssh.Channel, req *ssh.Request, ctx *srv.Ser
 	return nil
 }
 
-func (s *Server) handleSubsystem(ch ssh.Channel, req *ssh.Request, ctx *srv.ServerContext) error {
-	sb, err := s.parseSubsystemRequest(req, ctx)
+func (s *Server) handleSubsystem(ctx context.Context, ch ssh.Channel, req *ssh.Request, serverContext *srv.ServerContext) error {
+	sb, err := s.parseSubsystemRequest(req, serverContext)
 	if err != nil {
-		ctx.Warnf("Failed to parse subsystem request: %v: %v.", req, err)
+		serverContext.Warnf("Failed to parse subsystem request: %v: %v.", req, err)
 		return trace.Wrap(err)
 	}
-	ctx.Debugf("Subsystem request: %v.", sb)
+	serverContext.Debugf("Subsystem request: %v.", sb)
 	// starting subsystem is blocking to the client,
 	// while collecting its result and waiting is not blocking
-	if err := sb.Start(ctx.ServerConn, ch, req, ctx); err != nil {
-		ctx.Warnf("Subsystem request %v failed: %v.", sb, err)
-		ctx.SendSubsystemResult(srv.SubsystemResult{Err: trace.Wrap(err)})
+	if err := sb.Start(ctx, serverContext.ServerConn, ch, req, serverContext); err != nil {
+		serverContext.Warnf("Subsystem request %v failed: %v.", sb, err)
+		serverContext.SendSubsystemResult(srv.SubsystemResult{Err: trace.Wrap(err)})
 		return trace.Wrap(err)
 	}
 	go func() {
 		err := sb.Wait()
 		log.Debugf("Subsystem %v finished with result: %v.", sb, err)
-		ctx.SendSubsystemResult(srv.SubsystemResult{Err: trace.Wrap(err)})
+		serverContext.SendSubsystemResult(srv.SubsystemResult{Err: trace.Wrap(err)})
 	}()
 	return nil
 }
@@ -1794,7 +1809,7 @@ func (s *Server) handleProxyJump(ctx context.Context, ccx *sshutils.ConnectionCo
 		return
 	}
 
-	if err := subsys.Start(scx.ServerConn, ch, &ssh.Request{}, scx); err != nil {
+	if err := subsys.Start(ctx, scx.ServerConn, ch, &ssh.Request{}, scx); err != nil {
 		log.Errorf("Unable to start proxy subsystem: %v.", err)
 		writeStderr(ch, "Unable to start proxy subsystem.")
 		return

lib/srv/subsystem.go

@@ -17,6 +17,8 @@ limitations under the License.
 package srv
 
 import (
+	"context"
+
 	"golang.org/x/crypto/ssh"
 )
 
@@ -29,12 +31,12 @@ type SubsystemResult struct {
 	Err error
 }
 
-// Subsystem represents SSH subsytem - special command executed
+// Subsystem represents SSH subsystem - special command executed
 // in the context of the session.
 type Subsystem interface {
 	// Start starts subsystem
-	Start(*ssh.ServerConn, ssh.Channel, *ssh.Request, *ServerContext) error
+	Start(context.Context, *ssh.ServerConn, ssh.Channel, *ssh.Request, *ServerContext) error
 
-	// Wait is returned by subystem when it's completed
+	// Wait is returned by subsystem when it's completed
 	Wait() error
 }

lib/sshutils/server.go

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-// Package sshutils contains contains the implementations of the base SSH
+// Package sshutils contains the implementations of the base SSH
 // server used throughout Teleport.
 package sshutils
 
@@ -28,18 +28,17 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/gravitational/trace"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
 
 	"github.com/gravitational/teleport"
+	"github.com/gravitational/teleport/api/observability/tracing"
 	"github.com/gravitational/teleport/api/utils/sshutils"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/limiter"
 	"github.com/gravitational/teleport/lib/utils"
-
-	"github.com/gravitational/trace"
-
-	"github.com/prometheus/client_golang/prometheus"
-	log "github.com/sirupsen/logrus"
 )
 
 var proxyConnectionLimitHitCount = prometheus.NewCounter(
@@ -54,7 +53,7 @@ var proxyConnectionLimitHitCount = prometheus.NewCounter(
 type Server struct {
 	sync.RWMutex
 
-	log log.FieldLogger
+	log logrus.FieldLogger
 	// component is a name of the facility which uses this server,
 	// used for logging/debugging. typically it's "proxy" or "auth api", etc
 	component string
@@ -95,11 +94,6 @@ const (
 	// https://tools.ietf.org/html/rfc4253#page-4
 	SSHVersionPrefix = "SSH-2.0-Teleport"
 
-	// ProxyHelloSignature is a string which Teleport proxy will send
-	// right after the initial SSH "handshake/version" message if it detects
-	// talking to a Teleport server.
-	ProxyHelloSignature = "Teleport-Proxy"
-
 	// MaxVersionStringBytes is the maximum number of bytes allowed for a
 	// SSH version string
 	// https://tools.ietf.org/html/rfc4253
@@ -114,7 +108,7 @@ const (
 type ServerOption func(cfg *Server) error
 
 // SetLogger sets the logger for the server
-func SetLogger(logger log.FieldLogger) ServerOption {
+func SetLogger(logger logrus.FieldLogger) ServerOption {
 	return func(s *Server) error {
 		s.log = logger.WithField(trace.Component, "ssh:"+s.component)
 		return nil
@@ -160,7 +154,7 @@ func NewServer(
 
 	closeContext, cancel := context.WithCancel(context.TODO())
 	s := &Server{
-		log: log.WithFields(log.Fields{
+		log: logrus.WithFields(logrus.Fields{
 			trace.Component: "ssh:" + component,
 		}),
 		addr:           a,
@@ -396,13 +390,13 @@ func (s *Server) HandleConnection(conn net.Conn) {
 	// in case of error as ssh server takes care of this
 	remoteAddr, _, err := net.SplitHostPort(conn.RemoteAddr().String())
 	if err != nil {
-		log.Errorf(err.Error())
+		s.log.Errorf(err.Error())
 	}
 	if err := s.limiter.AcquireConnection(remoteAddr); err != nil {
 		if trace.IsLimitExceeded(err) {
 			proxyConnectionLimitHitCount.Inc()
 		}
-		log.Errorf(err.Error())
+		s.log.Errorf(err.Error())
 		conn.Close()
 		return
 	}
@@ -420,12 +414,15 @@ func (s *Server) HandleConnection(conn net.Conn) {
 	// create a new SSH server which handles the handshake (and pass the custom
 	// payload structure which will be populated only when/if this connection
 	// comes from another Teleport proxy):
-	sconn, chans, reqs, err := ssh.NewServerConn(wrapConnection(wconn), &s.cfg)
+	wrappedConn := wrapConnection(wconn, s.log)
+	sconn, chans, reqs, err := ssh.NewServerConn(wrappedConn, &s.cfg)
 	if err != nil {
 		conn.SetDeadline(time.Time{})
 		return
 	}
 
+	ctx := tracing.WithPropagationContext(context.Background(), wrappedConn.traceContext)
+
 	certType := "unknown"
 	if sconn.Permissions != nil {
 		certType = sconn.Permissions.Extensions[utils.ExtIntCertType]
@@ -438,7 +435,7 @@ func (s *Server) HandleConnection(conn net.Conn) {
 
 	user := sconn.User()
 	if err := s.limiter.RegisterRequest(user); err != nil {
-		log.Errorf(err.Error())
+		s.log.Errorf(err.Error())
 		sconn.Close()
 		conn.Close()
 		return
@@ -462,7 +459,7 @@ func (s *Server) HandleConnection(conn net.Conn) {
 	// closeContext field is used to trigger starvation on cancellation by halting
 	// the acceptance of new connections; it is not intended to halt in-progress
 	// connection handling, and is therefore orthogonal to the role of ConnectionContext.
-	ctx, ccx := NewConnectionContext(context.Background(), wconn, sconn)
+	ctx, ccx := NewConnectionContext(ctx, wconn, sconn)
 	defer ccx.Close()
 
 	if s.newConnHandler != nil {
@@ -501,7 +498,7 @@ func (s *Server) HandleConnection(conn net.Conn) {
 			}
 			s.log.Debugf("Received out-of-band request: %+v.", req)
 			if s.reqHandler != nil {
-				go s.reqHandler.HandleRequest(req)
+				go s.reqHandler.HandleRequest(ctx, req)
 			}
 			// handle channels:
 		case nch := <-chans:
@@ -515,7 +512,7 @@ func (s *Server) HandleConnection(conn net.Conn) {
 			const wantReply = true
 			_, _, err = sconn.SendRequest(teleport.KeepAliveReqType, wantReply, keepAlivePayload[:])
 			if err != nil {
-				log.Errorf("Failed sending keepalive request: %v", err)
+				s.log.Errorf("Failed sending keepalive request: %v", err)
 			}
 		case <-ctx.Done():
 			s.log.Debugf("Connection context canceled: %v -> %v", conn.RemoteAddr(), conn.LocalAddr())
@@ -525,7 +522,7 @@ func (s *Server) HandleConnection(conn net.Conn) {
 }
 
 type RequestHandler interface {
-	HandleRequest(r *ssh.Request)
+	HandleRequest(ctx context.Context, r *ssh.Request)
 }
 
 type NewChanHandler interface {
@@ -608,23 +605,15 @@ type (
 	PasswordFunc  func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error)
 )
 
-// HandshakePayload structure is sent as a JSON blob by the teleport
-// proxy to every SSH server who identifies itself as Teleport server
-//
-// It allows teleport proxies to communicate additional data to server
-type HandshakePayload struct {
-	// ClientAddr is the IP address of the remote client
-	ClientAddr string `json:"clientAddr,omitempty"`
-}
-
 // connectionWrapper allows the SSH server to perform custom handshake which
 // lets teleport proxy servers to relay a true remote client IP address
 // to the SSH server.
 //
 // (otherwise connection.RemoteAddr (client IP) will always point to a proxy IP
-// instead of oa true client IP)
+// instead of a true client IP)
 type connectionWrapper struct {
 	net.Conn
+	logger logrus.FieldLogger
 
 	// upstreamReader reads from the underlying (wrapped) connection
 	upstreamReader io.Reader
@@ -633,6 +622,10 @@ type connectionWrapper struct {
 	// a proxy). Keeping this address is the entire point of the
 	// connection wrapper.
 	clientAddr net.Addr
+
+	// traceContext is the tracing context that was passed across the
+	// connection, used to correlate spans.
+	traceContext tracing.PropagationContext
 }
 
 // RemoteAddr returns the behind-the-proxy client address
@@ -654,7 +647,7 @@ func (c *connectionWrapper) Read(b []byte) (int, error) {
 	if err != nil {
 		// EOF happens quite often, don't pollute the logs with it
 		if !trace.IsEOF(err) {
-			log.Error(err)
+			c.logger.Error(err)
 		}
 		return n, err
 	}
@@ -663,23 +656,23 @@ func (c *connectionWrapper) Read(b []byte) (int, error) {
 	skip := 0
 
 	// are we reading from a Teleport proxy?
-	if bytes.HasPrefix(buff, []byte(ProxyHelloSignature)) {
-		// the JSON paylaod ends with a binary zero:
+	if bytes.HasPrefix(buff, []byte(sshutils.ProxyHelloSignature)) {
+		// the JSON payload ends with a binary zero:
 		payloadBoundary := bytes.IndexByte(buff, 0x00)
 		if payloadBoundary > 0 {
-			var hp HandshakePayload
-			payload := buff[len(ProxyHelloSignature):payloadBoundary]
+			var hp sshutils.HandshakePayload
+			payload := buff[len(sshutils.ProxyHelloSignature):payloadBoundary]
 			if err = json.Unmarshal(payload, &hp); err != nil {
-				log.Error(err)
+				c.logger.Error(err)
 			} else {
 				ca, err := utils.ParseAddr(hp.ClientAddr)
-				if err != nil {
-					log.Error(err)
-				} else {
+				if err == nil {
 					// replace proxy's client addr with a real client address
 					// we just got from the custom payload:
 					c.clientAddr = ca
 				}
+
+				c.traceContext = hp.TracingContext
 			}
 			skip = payloadBoundary + 1
 		}
@@ -690,9 +683,10 @@ func (c *connectionWrapper) Read(b []byte) (int, error) {
 
 // wrapConnection takes a network connection, wraps it into connectionWrapper
 // object (which overrides Read method) and returns the wrapper.
-func wrapConnection(conn net.Conn) net.Conn {
+func wrapConnection(conn net.Conn, logger logrus.FieldLogger) *connectionWrapper {
 	return &connectionWrapper{
 		Conn:       conn,
 		clientAddr: conn.RemoteAddr(),
+		logger:     logger,
 	}
 }

lib/utils/proxy/proxy.go

@@ -13,6 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
+
 package proxy
 
 import (
@@ -21,17 +22,15 @@ import (
 	"net"
 	"time"
 
-	"github.com/gravitational/trace"
-
 	"github.com/gravitational/teleport"
 	apiclient "github.com/gravitational/teleport/api/client"
 	apiproxy "github.com/gravitational/teleport/api/client/proxy"
-	"github.com/gravitational/teleport/api/utils/sshutils"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/lib/utils"
 
-	"golang.org/x/crypto/ssh"
-
+	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/crypto/ssh"
 )
 
 var log = logrus.WithFields(logrus.Fields{
@@ -41,18 +40,22 @@ var log = logrus.WithFields(logrus.Fields{
 // dialWithDeadline works around the case when net.DialWithTimeout
 // succeeds, but key exchange hangs. Setting deadline on connection
 // prevents this case from happening
-func dialWithDeadline(network string, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
-	conn, err := net.DialTimeout(network, addr, config.Timeout)
+func dialWithDeadline(ctx context.Context, network string, addr string, config *ssh.ClientConfig) (*tracessh.Client, error) {
+	dialer := &net.Dialer{
+		Timeout: config.Timeout,
+	}
+
+	conn, err := dialer.DialContext(ctx, network, addr)
 	if err != nil {
 		return nil, err
 	}
-	return sshutils.NewClientConnWithDeadline(conn, addr, config)
+	return tracessh.NewClientConnWithDeadline(ctx, conn, addr, config)
 }
 
 // dialALPNWithDeadline allows connecting to Teleport in single-port mode. SSH protocol is wrapped into
 // TLS connection where TLS ALPN protocol is set to ProtocolReverseTunnel allowing ALPN Proxy to route the
 // incoming connection to ReverseTunnel proxy service.
-func (d directDial) dialALPNWithDeadline(network string, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
+func (d directDial) dialALPNWithDeadline(ctx context.Context, network string, addr string, config *ssh.ClientConfig) (*tracessh.Client, error) {
 	dialer := &net.Dialer{
 		Timeout: config.Timeout,
 	}
@@ -64,20 +67,26 @@ func (d directDial) dialALPNWithDeadline(network string, addr string, config *ss
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
-	tlsConn, err := tls.DialWithDialer(dialer, network, addr, conf)
+
+	tlsDialer := tls.Dialer{
+		NetDialer: dialer,
+		Config:    conf,
+	}
+
+	tlsConn, err := tlsDialer.DialContext(ctx, network, addr)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
-	return sshutils.NewClientConnWithDeadline(tlsConn, addr, config)
+	return tracessh.NewClientConnWithDeadline(ctx, tlsConn, addr, config)
 }
 
 // A Dialer is a means for a client to establish a SSH connection.
 type Dialer interface {
 	// Dial establishes a client connection to a SSH server.
-	Dial(network string, addr string, config *ssh.ClientConfig) (*ssh.Client, error)
+	Dial(ctx context.Context, network string, addr string, config *ssh.ClientConfig) (*tracessh.Client, error)
 
 	// DialTimeout acts like Dial but takes a timeout.
-	DialTimeout(network, address string, timeout time.Duration) (net.Conn, error)
+	DialTimeout(ctx context.Context, network, address string, timeout time.Duration) (net.Conn, error)
 }
 
 type directDial struct {
@@ -101,15 +110,15 @@ func (d directDial) getTLSConfig(addr *utils.NetAddr) (*tls.Config, error) {
 }
 
 // Dial calls ssh.Dial directly.
-func (d directDial) Dial(network string, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
+func (d directDial) Dial(ctx context.Context, network string, addr string, config *ssh.ClientConfig) (*tracessh.Client, error) {
 	if d.tlsRoutingEnabled {
-		client, err := d.dialALPNWithDeadline(network, addr, config)
+		client, err := d.dialALPNWithDeadline(ctx, network, addr, config)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
 		return client, nil
 	}
-	client, err := dialWithDeadline(network, addr, config)
+	client, err := dialWithDeadline(ctx, network, addr, config)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -117,7 +126,11 @@ func (d directDial) Dial(network string, addr string, config *ssh.ClientConfig)
 }
 
 // DialTimeout acts like Dial but takes a timeout.
-func (d directDial) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
+func (d directDial) DialTimeout(ctx context.Context, network, address string, timeout time.Duration) (net.Conn, error) {
+	dialer := &net.Dialer{
+		Timeout: timeout,
+	}
+
 	if d.tlsRoutingEnabled {
 		addr, err := utils.ParseAddr(address)
 		if err != nil {
@@ -127,15 +140,19 @@ func (d directDial) DialTimeout(network, address string, timeout time.Duration)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
-		tlsConn, err := tls.DialWithDialer(&net.Dialer{
-			Timeout: timeout,
-		}, "tcp", address, conf)
+
+		tlsDialer := tls.Dialer{
+			NetDialer: dialer,
+			Config:    conf,
+		}
+
+		tlsConn, err := tlsDialer.DialContext(ctx, "tcp", address)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
 		return tlsConn, nil
 	}
-	conn, err := net.DialTimeout(network, address, timeout)
+	conn, err := dialer.DialContext(ctx, network, address)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
@@ -165,9 +182,8 @@ func (d proxyDial) getTLSConfig(addr *utils.NetAddr) (*tls.Config, error) {
 }
 
 // DialTimeout acts like Dial but takes a timeout.
-func (d proxyDial) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
+func (d proxyDial) DialTimeout(ctx context.Context, network, address string, timeout time.Duration) (net.Conn, error) {
 	// Build a proxy connection first.
-	ctx := context.Background()
 	if timeout > 0 {
 		timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
 		defer cancel()
@@ -198,14 +214,16 @@ func (d proxyDial) DialTimeout(network, address string, timeout time.Duration) (
 
 // Dial first connects to a proxy, then uses the connection to establish a new
 // SSH connection.
-func (d proxyDial) Dial(network string, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
+func (d proxyDial) Dial(ctx context.Context, network string, addr string, config *ssh.ClientConfig) (*tracessh.Client, error) {
 	// Build a proxy connection first.
-	pconn, err := apiclient.DialProxy(context.Background(), d.proxyHost, addr)
+	pconn, err := apiclient.DialProxy(ctx, d.proxyHost, addr)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 	if config.Timeout > 0 {
-		pconn.SetReadDeadline(time.Now().Add(config.Timeout))
+		if err := pconn.SetReadDeadline(time.Now().Add(config.Timeout)); err != nil {
+			return nil, trace.Wrap(err)
+		}
 	}
 	if d.tlsRoutingEnabled {
 		address, err := utils.ParseAddr(addr)
@@ -220,14 +238,16 @@ func (d proxyDial) Dial(network string, addr string, config *ssh.ClientConfig) (
 	}
 
 	// Do the same as ssh.Dial but pass in proxy connection.
-	c, chans, reqs, err := ssh.NewClientConn(pconn, addr, config)
+	c, chans, reqs, err := tracessh.NewClientConn(ctx, pconn, addr, config)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 	if config.Timeout > 0 {
-		pconn.SetReadDeadline(time.Time{})
+		if err := pconn.SetReadDeadline(time.Time{}); err != nil {
+			return nil, trace.Wrap(err)
+		}
 	}
-	return ssh.NewClient(c, chans, reqs), nil
+	return tracessh.NewClient(c, chans, reqs), nil
 }
 
 type dialerOptions struct {

lib/web/terminal.go

@@ -28,7 +28,6 @@ import (
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/gorilla/websocket"
-	"github.com/gravitational/teleport/lib/services"
 	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh"
@@ -37,11 +36,13 @@ import (
 
 	"github.com/gravitational/teleport"
 	authproto "github.com/gravitational/teleport/api/client/proto"
+	tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
 	"github.com/gravitational/teleport/api/types"
 	wanlib "github.com/gravitational/teleport/lib/auth/webauthn"
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/events"
+	"github.com/gravitational/teleport/lib/services"
 	"github.com/gravitational/teleport/lib/session"
 	"github.com/gravitational/teleport/lib/sshutils"
 	"github.com/gravitational/teleport/lib/utils"
@@ -328,7 +329,7 @@ func (t *TerminalHandler) makeClient(ws *websocket.Conn, r *http.Request) (*clie
 	// Save the *ssh.Session after the shell has been created. The session is
 	// used to update all other parties window size to that of the web client and
 	// to allow future window changes.
-	tc.OnShellCreated = func(s *ssh.Session, c *ssh.Client, _ io.ReadWriteCloser) (bool, error) {
+	tc.OnShellCreated = func(s *ssh.Session, c *tracessh.Client, _ io.ReadWriteCloser) (bool, error) {
 		t.sshSession = s
 		t.windowChange(&t.params.Term)
 

tool/tsh/proxy.go

@@ -108,7 +108,7 @@ func sshProxyWithTLSRouting(cf *CLIConf, tc *libclient.TeleportClient, targetHos
 		return trace.Wrap(err)
 	}
 	defer lp.Close()
-	if err := lp.SSHProxy(tc.LocalAgent()); err != nil {
+	if err := lp.SSHProxy(cf.Context, tc.LocalAgent()); err != nil {
 		return trace.Wrap(err)
 	}
 	return nil