coder/peerbroker/proxy.go

package peerbroker

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"

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

	"cdr.dev/slog"
	"github.com/coder/coder/coderd/database"
	"github.com/coder/coder/peerbroker/proto"
)

var (
	// Each NegotiateConnection() function call spawns a new stream.
	streamIDLength = len(uuid.NewString())
	// We shouldn't PubSub anything larger than this!
	maxPayloadSizeBytes = 8192
)

// ProxyOptions provides values to configure a proxy.
type ProxyOptions struct {
	ChannelID string
	Logger    slog.Logger
	Pubsub    database.Pubsub
}

// ProxyDial writes client negotiation streams over PubSub.
//
// PubSub is used to geodistribute WebRTC handshakes. All negotiation
// messages are small in size (<=8KB), and we don't require delivery
// guarantees because connections can always be renegotiated.
//
//	┌────────────────────┐   ┌─────────────────────────────┐
//	│      coderd        │   │          coderd             │
//
// ┌─────────────────────┐   │/<agent-id>/connect │   │    /<agent-id>/listen       │
// │       client        │   │                    │   │                             │   ┌─────┐
// │                     ├──►│Creates a stream ID │◄─►│Subscribe() to the <agent-id>│◄──┤agent│
// │NegotiateConnection()│   │and Publish() to the│   │channel. Parse the stream ID │   └─────┘
// └─────────────────────┘   │<agent-id> channel: │   │from payloads to create new  │
//
//	│                    │   │NegotiateConnection() streams│
//	│<stream-id><payload>│   │or write to existing ones.   │
//	└────────────────────┘   └─────────────────────────────┘
func ProxyDial(client proto.DRPCPeerBrokerClient, options ProxyOptions) (io.Closer, error) {
	proxyDial := &proxyDial{
		channelID:  options.ChannelID,
		logger:     options.Logger,
		pubsub:     options.Pubsub,
		connection: client,
		streams:    make(map[string]proto.DRPCPeerBroker_NegotiateConnectionClient),
	}
	return proxyDial, proxyDial.listen()
}

// ProxyListen accepts client negotiation streams over PubSub and writes them to the listener
// as new NegotiateConnection() streams.
func ProxyListen(ctx context.Context, connListener net.Listener, options ProxyOptions) error {
	mux := drpcmux.New()
	err := proto.DRPCRegisterPeerBroker(mux, &proxyListen{
		channelID: options.ChannelID,
		pubsub:    options.Pubsub,
		logger:    options.Logger,
	})
	if err != nil {
		return xerrors.Errorf("register peer broker: %w", err)
	}
	server := drpcserver.New(mux)
	err = server.Serve(ctx, connListener)
	if err != nil {
		if errors.Is(err, yamux.ErrSessionShutdown) {
			return nil
		}
		return xerrors.Errorf("serve: %w", err)
	}
	return nil
}

type proxyListen struct {
	channelID string
	pubsub    database.Pubsub
	logger    slog.Logger
}

func (p *proxyListen) NegotiateConnection(stream proto.DRPCPeerBroker_NegotiateConnectionStream) error {
	streamID := uuid.NewString()
	var err error
	closeSubscribe, err := p.pubsub.Subscribe(proxyInID(p.channelID), func(ctx context.Context, message []byte) {
		err := p.onServerToClientMessage(streamID, stream, message)
		if err != nil {
			p.logger.Debug(ctx, "failed to accept server message", slog.Error(err))
		}
	})
	if err != nil {
		return xerrors.Errorf("subscribe: %w", err)
	}
	defer closeSubscribe()
	for {
		clientToServerMessage, err := stream.Recv()
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			return xerrors.Errorf("recv: %w", err)
		}
		data, err := protobuf.Marshal(clientToServerMessage)
		if err != nil {
			return xerrors.Errorf("marshal: %w", err)
		}
		if len(data) > maxPayloadSizeBytes {
			return xerrors.Errorf("maximum payload size %d exceeded", maxPayloadSizeBytes)
		}
		data = append([]byte(streamID), data...)
		err = p.pubsub.Publish(proxyOutID(p.channelID), marshal(data))
		if err != nil {
			return xerrors.Errorf("publish: %w", err)
		}
	}
	return nil
}

func (*proxyListen) onServerToClientMessage(streamID string, stream proto.DRPCPeerBroker_NegotiateConnectionStream, message []byte) error {
	var err error
	message, err = unmarshal(message)
	if err != nil {
		return xerrors.Errorf("decode: %w", err)
	}
	if len(message) < streamIDLength {
		return xerrors.Errorf("got message length %d < %d", len(message), streamIDLength)
	}
	serverStreamID := string(message[0:streamIDLength])
	if serverStreamID != streamID {
		// It's not trying to communicate with this stream!
		return nil
	}
	var msg proto.Exchange
	err = protobuf.Unmarshal(message[streamIDLength:], &msg)
	if err != nil {
		return xerrors.Errorf("unmarshal message: %w", err)
	}
	err = stream.Send(&msg)
	if err != nil {
		return xerrors.Errorf("send message: %w", err)
	}
	return nil
}

type proxyDial struct {
	channelID string
	pubsub    database.Pubsub
	logger    slog.Logger

	connection     proto.DRPCPeerBrokerClient
	closeSubscribe func()
	streamMutex    sync.Mutex
	streams        map[string]proto.DRPCPeerBroker_NegotiateConnectionClient
}

func (p *proxyDial) listen() error {
	var err error
	p.closeSubscribe, err = p.pubsub.Subscribe(proxyOutID(p.channelID), func(ctx context.Context, message []byte) {
		err := p.onClientToServerMessage(ctx, message)
		if err != nil {
			p.logger.Debug(ctx, "failed to accept client message", slog.Error(err))
		}
	})
	if err != nil {
		return err
	}
	return nil
}

func (p *proxyDial) onClientToServerMessage(ctx context.Context, message []byte) error {
	var err error
	message, err = unmarshal(message)
	if err != nil {
		return xerrors.Errorf("decode: %w", err)
	}
	if len(message) < streamIDLength {
		return xerrors.Errorf("got message length %d < %d", len(message), streamIDLength)
	}
	streamID := string(message[0:streamIDLength])
	p.streamMutex.Lock()
	stream, ok := p.streams[streamID]
	if !ok {
		stream, err = p.connection.NegotiateConnection(ctx)
		if err != nil {
			p.streamMutex.Unlock()
			return xerrors.Errorf("negotiate connection: %w", err)
		}
		p.streams[streamID] = stream
		go func() {
			defer stream.Close()

			err := p.onServerToClientMessage(streamID, stream)
			if err != nil {
				p.logger.Debug(ctx, "failed to accept server message", slog.Error(err))
			}
		}()
		go func() {
			<-stream.Context().Done()
			p.streamMutex.Lock()
			delete(p.streams, streamID)
			p.streamMutex.Unlock()
		}()
	}
	p.streamMutex.Unlock()

	var msg proto.Exchange
	err = protobuf.Unmarshal(message[streamIDLength:], &msg)
	if err != nil {
		return xerrors.Errorf("unmarshal message: %w", err)
	}
	err = stream.Send(&msg)
	if err != nil {
		return xerrors.Errorf("write message: %w", err)
	}
	return nil
}

func (p *proxyDial) onServerToClientMessage(streamID string, stream proto.DRPCPeerBroker_NegotiateConnectionClient) error {
	for {
		serverToClientMessage, err := stream.Recv()
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			if errors.Is(err, context.Canceled) {
				break
			}
			return xerrors.Errorf("recv: %w", err)
		}
		data, err := protobuf.Marshal(serverToClientMessage)
		if err != nil {
			return xerrors.Errorf("marshal: %w", err)
		}
		if len(data) > maxPayloadSizeBytes {
			return xerrors.Errorf("maximum payload size %d exceeded", maxPayloadSizeBytes)
		}
		data = append([]byte(streamID), data...)
		err = p.pubsub.Publish(proxyInID(p.channelID), marshal(data))
		if err != nil {
			return xerrors.Errorf("publish: %w", err)
		}
	}
	return nil
}

func (p *proxyDial) Close() error {
	p.streamMutex.Lock()
	defer p.streamMutex.Unlock()
	p.closeSubscribe()
	return nil
}

// base64 needs to be used here to keep the pubsub messages in UTF-8 range.
// PostgreSQL cannot handle non UTF-8 messages over pubsub.
func marshal(data []byte) []byte {
	return []byte(base64.StdEncoding.EncodeToString(data))
}

func unmarshal(data []byte) ([]byte, error) {
	return base64.StdEncoding.DecodeString(string(data))
}

func proxyOutID(channelID string) string {
	return fmt.Sprintf("%s-out", channelID)
}

func proxyInID(channelID string) string {
	return fmt.Sprintf("%s-in", channelID)
}