coder/tailnet/service.go

package tailnet

import (
	"context"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
	"github.com/hashicorp/yamux"
	"golang.org/x/xerrors"
	"storj.io/drpc/drpcmux"
	"storj.io/drpc/drpcserver"
	"tailscale.com/tailcfg"

	"cdr.dev/slog/v3"
	"github.com/coder/coder/v2/apiversion"
	"github.com/coder/coder/v2/codersdk/drpcsdk"
	"github.com/coder/coder/v2/tailnet/proto"
	"github.com/coder/quartz"
)

var ErrUnsupportedVersion = xerrors.New("unsupported version")

type streamIDContextKey struct{}

// StreamID identifies the caller of the CoordinateTailnet RPC.  We store this
// on the context, since the information is extracted at the HTTP layer for
// remote clients of the API, or set outside tailnet for local clients (e.g.
// Coderd's single_tailnet)
type StreamID struct {
	Name string
	ID   uuid.UUID
	Auth CoordinateeAuth
}

func WithStreamID(ctx context.Context, streamID StreamID) context.Context {
	return context.WithValue(ctx, streamIDContextKey{}, streamID)
}

type WorkspaceUpdatesProvider interface {
	io.Closer
	Subscribe(ctx context.Context, userID uuid.UUID) (Subscription, error)
}

type Subscription interface {
	io.Closer
	Updates() <-chan *proto.WorkspaceUpdate
}

type TunnelAuthorizer interface {
	AuthorizeTunnel(ctx context.Context, agentID uuid.UUID) error
}

type ClientServiceOptions struct {
	Logger                   slog.Logger
	CoordPtr                 *atomic.Pointer[Coordinator]
	DERPMapUpdateFrequency   time.Duration
	DERPMapFn                func() *tailcfg.DERPMap
	NetworkTelemetryHandler  func(batch []*proto.TelemetryEvent)
	ResumeTokenProvider      ResumeTokenProvider
	WorkspaceUpdatesProvider WorkspaceUpdatesProvider
}

// ClientService is a tailnet coordination service that accepts a connection and version from a
// tailnet client, and support versions 2.x of the Tailnet API protocol.
type ClientService struct {
	Logger   slog.Logger
	CoordPtr *atomic.Pointer[Coordinator]
	drpc     *drpcserver.Server
}

// NewClientService returns a ClientService based on the given Coordinator pointer.  The pointer is
// loaded on each processed connection.
func NewClientService(options ClientServiceOptions) (
	*ClientService, error,
) {
	s := &ClientService{Logger: options.Logger, CoordPtr: options.CoordPtr}
	mux := drpcmux.New()
	drpcService := &DRPCService{
		CoordPtr:                 options.CoordPtr,
		Logger:                   options.Logger,
		DerpMapUpdateFrequency:   options.DERPMapUpdateFrequency,
		DerpMapFn:                options.DERPMapFn,
		NetworkTelemetryHandler:  options.NetworkTelemetryHandler,
		ResumeTokenProvider:      options.ResumeTokenProvider,
		WorkspaceUpdatesProvider: options.WorkspaceUpdatesProvider,
	}
	err := proto.DRPCRegisterTailnet(mux, drpcService)
	if err != nil {
		return nil, xerrors.Errorf("register DRPC service: %w", err)
	}
	server := drpcserver.NewWithOptions(mux, drpcserver.Options{
		Manager: drpcsdk.DefaultDRPCOptions(nil),
		Log: func(err error) {
			if xerrors.Is(err, io.EOF) ||
				xerrors.Is(err, context.Canceled) ||
				xerrors.Is(err, context.DeadlineExceeded) {
				return
			}
			options.Logger.Debug(context.Background(), "drpc server error", slog.Error(err))
		},
	})
	s.drpc = server
	return s, nil
}

func (s *ClientService) ServeClient(ctx context.Context, version string, conn net.Conn, streamID StreamID) error {
	major, _, err := apiversion.Parse(version)
	if err != nil {
		s.Logger.Warn(ctx, "serve client called with unparsable version", slog.Error(err))
		return err
	}
	switch major {
	case 2:
		return s.ServeConnV2(ctx, conn, streamID)
	default:
		s.Logger.Warn(ctx, "serve client called with unsupported version", slog.F("version", version))
		return ErrUnsupportedVersion
	}
}

func (s ClientService) ServeConnV2(ctx context.Context, conn net.Conn, streamID StreamID) error {
	config := yamux.DefaultConfig()
	config.LogOutput = io.Discard
	session, err := yamux.Server(conn, config)
	if err != nil {
		return xerrors.Errorf("yamux init failed: %w", err)
	}
	ctx = WithStreamID(ctx, streamID)
	s.Logger.Debug(ctx, "serving dRPC tailnet v2 API session",
		slog.F("peer_id", streamID.ID.String()))
	return s.drpc.Serve(ctx, session)
}

// DRPCService is the dRPC-based, version 2.x of the tailnet API and implements proto.DRPCClientServer
type DRPCService struct {
	CoordPtr                 *atomic.Pointer[Coordinator]
	Logger                   slog.Logger
	DerpMapUpdateFrequency   time.Duration
	DerpMapFn                func() *tailcfg.DERPMap
	NetworkTelemetryHandler  func(batch []*proto.TelemetryEvent)
	ResumeTokenProvider      ResumeTokenProvider
	WorkspaceUpdatesProvider WorkspaceUpdatesProvider
}

func (s *DRPCService) PostTelemetry(_ context.Context, req *proto.TelemetryRequest) (*proto.TelemetryResponse, error) {
	if s.NetworkTelemetryHandler != nil {
		s.NetworkTelemetryHandler(req.Events)
	}
	return &proto.TelemetryResponse{}, nil
}

func (s *DRPCService) StreamDERPMaps(_ *proto.StreamDERPMapsRequest, stream proto.DRPCTailnet_StreamDERPMapsStream) error {
	defer stream.Close()

	ticker := time.NewTicker(s.DerpMapUpdateFrequency)
	defer ticker.Stop()

	var lastDERPMap *tailcfg.DERPMap
	for {
		derpMap := s.DerpMapFn()
		if derpMap == nil {
			// in testing, we send nil to close the stream.
			return io.EOF
		}
		if lastDERPMap == nil || !CompareDERPMaps(lastDERPMap, derpMap) {
			protoDERPMap := DERPMapToProto(derpMap)
			err := stream.Send(protoDERPMap)
			if err != nil {
				return xerrors.Errorf("send derp map: %w", err)
			}
			lastDERPMap = derpMap
		}

		ticker.Reset(s.DerpMapUpdateFrequency)
		select {
		case <-stream.Context().Done():
			return nil
		case <-ticker.C:
		}
	}
}

func (s *DRPCService) RefreshResumeToken(ctx context.Context, _ *proto.RefreshResumeTokenRequest) (*proto.RefreshResumeTokenResponse, error) {
	streamID, ok := ctx.Value(streamIDContextKey{}).(StreamID)
	if !ok {
		return nil, xerrors.New("no Stream ID")
	}

	res, err := s.ResumeTokenProvider.GenerateResumeToken(ctx, streamID.ID)
	if err != nil {
		return nil, xerrors.Errorf("generate resume token: %w", err)
	}
	return res, nil
}

func (s *DRPCService) Coordinate(stream proto.DRPCTailnet_CoordinateStream) error {
	ctx := stream.Context()
	streamID, ok := ctx.Value(streamIDContextKey{}).(StreamID)
	if !ok {
		_ = stream.Close()
		return xerrors.New("no Stream ID")
	}
	logger := s.Logger.With(slog.F("peer_id", streamID), slog.F("name", streamID.Name))
	logger.Debug(ctx, "starting tailnet Coordinate")
	coord := *(s.CoordPtr.Load())
	reqs, resps := coord.Coordinate(ctx, streamID.ID, streamID.Name, streamID.Auth)
	c := communicator{
		logger: logger,
		stream: stream,
		reqs:   reqs,
		resps:  resps,
	}
	c.communicate()
	return nil
}

func (s *DRPCService) WorkspaceUpdates(req *proto.WorkspaceUpdatesRequest, stream proto.DRPCTailnet_WorkspaceUpdatesStream) error {
	defer stream.Close()

	ctx := stream.Context()

	ownerID, err := uuid.FromBytes(req.WorkspaceOwnerId)
	if err != nil {
		return xerrors.Errorf("parse workspace owner ID: %w", err)
	}

	sub, err := s.WorkspaceUpdatesProvider.Subscribe(ctx, ownerID)
	if err != nil {
		return xerrors.Errorf("subscribe to workspace updates: %w", err)
	}
	defer sub.Close()

	for {
		select {
		case updates, ok := <-sub.Updates():
			if !ok {
				return nil
			}
			err := stream.Send(updates)
			if err != nil {
				return xerrors.Errorf("send workspace update: %w", err)
			}
		case <-stream.Context().Done():
			return nil
		}
	}
}

type communicator struct {
	logger slog.Logger
	stream proto.DRPCTailnet_CoordinateStream
	reqs   chan<- *proto.CoordinateRequest
	resps  <-chan *proto.CoordinateResponse
}

func (c communicator) communicate() {
	go c.loopReq()
	c.loopResp()
}

func (c communicator) loopReq() {
	ctx := c.stream.Context()
	defer close(c.reqs)
	for {
		req, err := c.stream.Recv()
		if err != nil {
			c.logger.Debug(ctx, "error receiving requests from DRPC stream", slog.Error(err))
			return
		}
		err = SendCtx(ctx, c.reqs, req)
		if err != nil {
			c.logger.Debug(ctx, "context done while sending coordinate request", slog.Error(ctx.Err()))
			return
		}
	}
}

func (c communicator) loopResp() {
	ctx := c.stream.Context()
	defer func() {
		err := c.stream.Close()
		if err != nil {
			c.logger.Debug(ctx, "loopResp hit error closing stream", slog.Error(err))
		}
	}()
	for {
		resp, err := RecvCtx(ctx, c.resps)
		if err != nil {
			c.logger.Debug(ctx, "loopResp failed to get response", slog.Error(err))
			return
		}
		err = c.stream.Send(resp)
		if err != nil {
			c.logger.Debug(ctx, "loopResp failed to send response to DRPC stream", slog.Error(err))
			return
		}
	}
}

type NetworkTelemetryBatcher struct {
	clock     quartz.Clock
	frequency time.Duration
	maxSize   int
	batchFn   func(batch []*proto.TelemetryEvent)

	mu      sync.Mutex
	closed  chan struct{}
	done    chan struct{}
	ticker  *quartz.Ticker
	pending []*proto.TelemetryEvent
}

func NewNetworkTelemetryBatcher(clk quartz.Clock, frequency time.Duration, maxSize int, batchFn func(batch []*proto.TelemetryEvent)) *NetworkTelemetryBatcher {
	b := &NetworkTelemetryBatcher{
		clock:     clk,
		frequency: frequency,
		maxSize:   maxSize,
		batchFn:   batchFn,
		closed:    make(chan struct{}),
		done:      make(chan struct{}),
	}
	if b.batchFn == nil {
		b.batchFn = func(_ []*proto.TelemetryEvent) {}
	}
	b.start()
	return b
}

func (b *NetworkTelemetryBatcher) Close() error {
	close(b.closed)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	select {
	case <-ctx.Done():
		return xerrors.New("timed out waiting for batcher to close")
	case <-b.done:
	}
	return nil
}

func (b *NetworkTelemetryBatcher) sendTelemetryBatch() {
	b.mu.Lock()
	defer b.mu.Unlock()
	events := b.pending
	if len(events) == 0 {
		return
	}
	b.pending = []*proto.TelemetryEvent{}
	b.batchFn(events)
}

func (b *NetworkTelemetryBatcher) start() {
	b.ticker = b.clock.NewTicker(b.frequency)

	go func() {
		defer func() {
			// The lock prevents Handler from racing with Close.
			b.mu.Lock()
			defer b.mu.Unlock()
			close(b.done)
			b.ticker.Stop()
		}()

		for {
			select {
			case <-b.ticker.C:
				b.sendTelemetryBatch()
				b.ticker.Reset(b.frequency)
			case <-b.closed:
				// Send any remaining telemetry events before exiting.
				b.sendTelemetryBatch()
				return
			}
		}
	}()
}

func (b *NetworkTelemetryBatcher) Handler(events []*proto.TelemetryEvent) {
	b.mu.Lock()
	defer b.mu.Unlock()
	select {
	case <-b.closed:
		return
	default:
	}

	for _, event := range events {
		b.pending = append(b.pending, event)

		if len(b.pending) >= b.maxSize {
			// This can't call sendTelemetryBatch directly because we already
			// hold the lock.
			events := b.pending
			b.pending = []*proto.TelemetryEvent{}
			// Resetting the ticker is best effort. We don't care if the ticker
			// has already fired or has a pending message, because the only risk
			// is that we send two telemetry events in short succession (which
			// is totally fine).
			b.ticker.Reset(b.frequency)
			// Perform the send in a goroutine to avoid blocking the DRPC call.
			go b.batchFn(events)
		}
	}
}