mirror of
https://github.com/coder/coder.git
synced 2025-07-12 00:14:10 +00:00
feat: Add peerbroker proxy for agent connections (#349)
* feat: Add peerbroker proxy for agent connections Agents will connect using this proxy. Eventually we'll intercept some of these messages for validation, but that's not necessary right now. * Add ASCII chart
This commit is contained in:
Kyle Carberry
GitHub
260
peerbroker/proxy.go
Normal file
260
peerbroker/proxy.go
Normal file
@ -0,0 +1,260 @@
|
||||
package peerbroker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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/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), data)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("publish: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (*proxyListen) onServerToClientMessage(streamID string, stream proto.DRPCPeerBroker_NegotiateConnectionStream, message []byte) error {
|
||||
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.NegotiateConnection_ServerToClient
|
||||
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 {
|
||||
if len(message) < streamIDLength {
|
||||
return xerrors.Errorf("got message length %d < %d", len(message), streamIDLength)
|
||||
}
|
||||
var err error
|
||||
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.NegotiateConnection_ClientToServer
|
||||
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), 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
|
||||
}
|
||||
|
||||
func proxyOutID(channelID string) string {
|
||||
return fmt.Sprintf("%s-out", channelID)
|
||||
}
|
||||
|
||||
func proxyInID(channelID string) string {
|
||||
return fmt.Sprintf("%s-in", channelID)
|
||||
}
|
||||
81
peerbroker/proxy_test.go
Normal file
81
peerbroker/proxy_test.go
Normal file
@ -0,0 +1,81 @@
|
||||
package peerbroker_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/pion/webrtc/v3"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"cdr.dev/slog"
|
||||
"cdr.dev/slog/sloggers/slogtest"
|
||||
"github.com/coder/coder/database"
|
||||
"github.com/coder/coder/peer"
|
||||
"github.com/coder/coder/peerbroker"
|
||||
"github.com/coder/coder/peerbroker/proto"
|
||||
"github.com/coder/coder/provisionersdk"
|
||||
)
|
||||
|
||||
func TestProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := context.Background()
|
||||
channelID := "hello"
|
||||
pubsub := database.NewPubsubInMemory()
|
||||
dialerClient, dialerServer := provisionersdk.TransportPipe()
|
||||
defer dialerClient.Close()
|
||||
defer dialerServer.Close()
|
||||
listenerClient, listenerServer := provisionersdk.TransportPipe()
|
||||
defer listenerClient.Close()
|
||||
defer listenerServer.Close()
|
||||
|
||||
listener, err := peerbroker.Listen(listenerServer, &peer.ConnOptions{
|
||||
Logger: slogtest.Make(t, nil).Named("server").Leveled(slog.LevelDebug),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
proxyCloser, err := peerbroker.ProxyDial(proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(listenerClient)), peerbroker.ProxyOptions{
|
||||
ChannelID: channelID,
|
||||
Logger: slogtest.Make(t, nil).Named("proxy-listen").Leveled(slog.LevelDebug),
|
||||
Pubsub: pubsub,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() {
|
||||
_ = proxyCloser.Close()
|
||||
})
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
err = peerbroker.ProxyListen(ctx, dialerServer, peerbroker.ProxyOptions{
|
||||
ChannelID: channelID,
|
||||
Logger: slogtest.Make(t, nil).Named("proxy-dial").Leveled(slog.LevelDebug),
|
||||
Pubsub: pubsub,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}()
|
||||
|
||||
api := proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(dialerClient))
|
||||
stream, err := api.NegotiateConnection(ctx)
|
||||
require.NoError(t, err)
|
||||
clientConn, err := peerbroker.Dial(stream, []webrtc.ICEServer{{
|
||||
URLs: []string{"stun:stun.l.google.com:19302"},
|
||||
}}, &peer.ConnOptions{
|
||||
Logger: slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer clientConn.Close()
|
||||
|
||||
serverConn, err := listener.Accept()
|
||||
require.NoError(t, err)
|
||||
defer serverConn.Close()
|
||||
_, err = serverConn.Ping()
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = clientConn.Ping()
|
||||
require.NoError(t, err)
|
||||
|
||||
_ = dialerServer.Close()
|
||||
wg.Wait()
|
||||
}
|
||||
Reference in New Issue
Block a user