github.com/gravitational/teleport/api@v0.0.0-20240507183017-3110591cbafc/observability/tracing/ssh/channel.go (about)

     1  // Copyright 2022 Gravitational, Inc
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package ssh
    16  
    17  import (
    18  	"context"
    19  	"encoding/json"
    20  	"fmt"
    21  
    22  	semconv "go.opentelemetry.io/otel/semconv/v1.10.0"
    23  	oteltrace "go.opentelemetry.io/otel/trace"
    24  	"golang.org/x/crypto/ssh"
    25  
    26  	"github.com/gravitational/teleport/api/observability/tracing"
    27  )
    28  
    29  // Channel is a wrapper around ssh.Channel that adds tracing support.
    30  type Channel struct {
    31  	ssh.Channel
    32  	tracingSupported tracingCapability
    33  	opts             []tracing.Option
    34  }
    35  
    36  // NewTraceChannel creates a new Channel.
    37  func NewTraceChannel(ch ssh.Channel, opts ...tracing.Option) *Channel {
    38  	return &Channel{
    39  		Channel: ch,
    40  		opts:    opts,
    41  	}
    42  }
    43  
    44  // SendRequest sends a global request, and returns the
    45  // reply. If tracing is enabled, the provided payload
    46  // is wrapped in an Envelope to forward any tracing context.
    47  func (c *Channel) SendRequest(ctx context.Context, name string, wantReply bool, payload []byte) (_ bool, err error) {
    48  	config := tracing.NewConfig(c.opts)
    49  	tracer := config.TracerProvider.Tracer(instrumentationName)
    50  
    51  	ctx, span := tracer.Start(
    52  		ctx,
    53  		fmt.Sprintf("ssh.ChannelRequest/%s", name),
    54  		oteltrace.WithSpanKind(oteltrace.SpanKindClient),
    55  		oteltrace.WithAttributes(
    56  			semconv.RPCServiceKey.String("ssh.Channel"),
    57  			semconv.RPCMethodKey.String("SendRequest"),
    58  			semconv.RPCSystemKey.String("ssh"),
    59  		),
    60  	)
    61  	defer func() { tracing.EndSpan(span, err) }()
    62  
    63  	return c.Channel.SendRequest(
    64  		name, wantReply, wrapPayload(ctx, c.tracingSupported, config.TextMapPropagator, payload),
    65  	)
    66  }
    67  
    68  // NewChannel is a wrapper around ssh.NewChannel that allows an
    69  // Envelope to be provided to new channels.
    70  type NewChannel struct {
    71  	ssh.NewChannel
    72  	Envelope Envelope
    73  }
    74  
    75  // NewTraceNewChannel wraps the ssh.NewChannel in a new NewChannel
    76  //
    77  // The provided ssh.NewChannel will have any Envelope provided
    78  // via ExtraData extracted so that the original payload can be
    79  // provided to callers of NewCh.ExtraData.
    80  func NewTraceNewChannel(nch ssh.NewChannel) *NewChannel {
    81  	ch := &NewChannel{
    82  		NewChannel: nch,
    83  	}
    84  
    85  	data := nch.ExtraData()
    86  
    87  	var envelope Envelope
    88  	if err := json.Unmarshal(data, &envelope); err == nil {
    89  		ch.Envelope = envelope
    90  	} else {
    91  		ch.Envelope.Payload = data
    92  	}
    93  
    94  	return ch
    95  }
    96  
    97  // ExtraData returns the arbitrary payload for this channel, as supplied
    98  // by the client. This data is specific to the channel type.
    99  func (n NewChannel) ExtraData() []byte {
   100  	return n.Envelope.Payload
   101  }