diff --git a/relay/metrics/realy.go b/relay/metrics/realy.go index 49a357557..1af7a1ca2 100644 --- a/relay/metrics/realy.go +++ b/relay/metrics/realy.go @@ -8,6 +8,8 @@ import ( log "github.com/sirupsen/logrus" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/metric" + + "github.com/netbirdio/netbird/shared/relay/messages" ) const ( @@ -17,16 +19,17 @@ const ( type Metrics struct { metric.Meter - TransferBytesSent metric.Int64Counter - TransferBytesRecv metric.Int64Counter - AuthenticationTime metric.Float64Histogram - PeerStoreTime metric.Float64Histogram - peerReconnections metric.Int64Counter - peers metric.Int64UpDownCounter - peerActivityChan chan string - peerLastActive map[string]time.Time - mutexActivity sync.Mutex - ctx context.Context + TransferBytesSent metric.Int64Counter + TransferBytesRecv metric.Int64Counter + AuthenticationTime metric.Float64Histogram + PeerStoreTime metric.Float64Histogram + peerReconnections metric.Int64Counter + peerProtocolVersions metric.Int64Counter + peers metric.Int64UpDownCounter + peerActivityChan chan string + peerLastActive map[string]time.Time + mutexActivity sync.Mutex + ctx context.Context } func NewMetrics(ctx context.Context, meter metric.Meter) (*Metrics, error) { @@ -88,14 +91,22 @@ func NewMetrics(ctx context.Context, meter metric.Meter) (*Metrics, error) { return nil, err } + peerProtocolVersions, err := meter.Int64Counter("relay_peer_protocol_versions_total", + metric.WithDescription("Total number of peer connections by advertised and selected protocol version"), + ) + if err != nil { + return nil, err + } + m := &Metrics{ - Meter: meter, - TransferBytesSent: bytesSent, - TransferBytesRecv: bytesRecv, - AuthenticationTime: authTime, - PeerStoreTime: peerStoreTime, - peers: peers, - peerReconnections: peerReconnections, + Meter: meter, + TransferBytesSent: bytesSent, + TransferBytesRecv: bytesRecv, + AuthenticationTime: authTime, + PeerStoreTime: peerStoreTime, + peers: peers, + peerReconnections: peerReconnections, + peerProtocolVersions: peerProtocolVersions, ctx: ctx, peerActivityChan: make(chan string, 10), @@ -151,6 +162,15 @@ func (m *Metrics) RecordPeerReconnection() { m.peerReconnections.Add(m.ctx, 1) } +// RecordProtocolVersion counts a new peer connection by the protocol version the client advertised +// and the version the server selected for it +func (m *Metrics) RecordProtocolVersion(advertised, selected messages.ProtocolVersion) { + m.peerProtocolVersions.Add(m.ctx, 1, metric.WithAttributes( + attribute.Int("advertised", int(advertised)), + attribute.Int("selected", int(selected)), + )) +} + // PeerActivity increases the active connections func (m *Metrics) PeerActivity(peerID string) { select { diff --git a/relay/server/handshake.go b/relay/server/handshake.go index f064b3501..c5ec39283 100644 --- a/relay/server/handshake.go +++ b/relay/server/handshake.go @@ -19,20 +19,9 @@ type Validator interface { Validate(any) error } -// preparedMsg contains the marshalled success response message +// preparedMsg contains the marshalled success response message for every supported protocol version type preparedMsg struct { - responseAuthMsg []byte -} - -func newPreparedMsg(instanceURL string) (*preparedMsg, error) { - ram, err := messages.MarshalAuthResponse(instanceURL) - if err != nil { - return nil, fmt.Errorf("failed to marshal auth response msg: %w", err) - } - - return &preparedMsg{ - responseAuthMsg: ram, - }, nil + responseAuthMsg map[messages.ProtocolVersion][]byte } type handshake struct { @@ -40,7 +29,35 @@ type handshake struct { validator Validator preparedMsg *preparedMsg - peerID *messages.PeerID + peerID *messages.PeerID + advertisedVersion messages.ProtocolVersion + selectedVersion messages.ProtocolVersion +} + +func newPreparedMsg(instanceURL string) (*preparedMsg, error) { + responses := make(map[messages.ProtocolVersion][]byte, messages.CurrentProtocolVersion) + for version := messages.ProtocolV1; ; version++ { + ram, err := messages.NewEncoder(version).AuthResponse(instanceURL) + if err != nil { + return nil, fmt.Errorf("marshal auth response for version %d: %w", version, err) + } + responses[version] = ram + if version == messages.CurrentProtocolVersion { + break + } + } + + return &preparedMsg{ + responseAuthMsg: responses, + }, nil +} + +func (p *preparedMsg) authResponse(version messages.ProtocolVersion) ([]byte, error) { + msg, ok := p.responseAuthMsg[version] + if !ok { + return nil, fmt.Errorf("no prepared auth response for protocol version %d", version) + } + return msg, nil } func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, error) { @@ -52,9 +69,9 @@ func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, err buf = buf[:n] - _, err = messages.ValidateVersion(buf) + advertised, err := messages.ReadVersion(buf) if err != nil { - return nil, fmt.Errorf("validate version from %s: %w", h.conn.RemoteAddr(), err) + return nil, fmt.Errorf("read version from %s: %w", h.conn.RemoteAddr(), err) } msgType, err := messages.DetermineClientMessageType(buf) @@ -70,12 +87,20 @@ func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, err if err != nil { return peerID, err } + h.peerID = peerID + h.advertisedVersion = advertised + h.selectedVersion = messages.SelectVersion(advertised, messages.CurrentProtocolVersion) return peerID, nil } func (h *handshake) handshakeResponse(ctx context.Context) error { - if _, err := h.conn.Write(ctx, h.preparedMsg.responseAuthMsg); err != nil { + msg, err := h.preparedMsg.authResponse(h.selectedVersion) + if err != nil { + return err + } + + if _, err := h.conn.Write(ctx, msg); err != nil { return fmt.Errorf("handshake response write to %s (%s): %w", h.peerID, h.conn.RemoteAddr(), err) } diff --git a/relay/server/handshake_test.go b/relay/server/handshake_test.go new file mode 100644 index 000000000..5cdcaa8ea --- /dev/null +++ b/relay/server/handshake_test.go @@ -0,0 +1,378 @@ +package server + +import ( + "context" + "encoding/binary" + "errors" + "net" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/shared/relay/auth/allow" + "github.com/netbirdio/netbird/shared/relay/messages" +) + +const ( + testWaitTimeout = 2 * time.Second + + // authTokenLenOffset is where the envelope's token length starts: header(2) + magic(4) + peer ID(36). + authTokenLenOffset = 2 + 4 + 36 +) + +// scriptedConn is a listener.Conn whose inbound frames are queued by the test and whose +// outbound frames are recorded. +type scriptedConn struct { + reads chan []byte + closed chan struct{} + closeOnce sync.Once + onWrite func([]byte) + + mu sync.Mutex + writes [][]byte +} + +type recordingValidator struct { + err error + + mu sync.Mutex + payloads [][]byte +} + +func newScriptedConn() *scriptedConn { + return &scriptedConn{ + reads: make(chan []byte, 16), + closed: make(chan struct{}), + } +} + +func (c *scriptedConn) Read(ctx context.Context, b []byte) (int, error) { + select { + case msg := <-c.reads: + return copy(b, msg), nil + case <-c.closed: + return 0, net.ErrClosed + case <-ctx.Done(): + return 0, ctx.Err() + } +} + +func (c *scriptedConn) Write(_ context.Context, b []byte) (int, error) { + select { + case <-c.closed: + return 0, net.ErrClosed + default: + } + + frame := append([]byte(nil), b...) + c.mu.Lock() + c.writes = append(c.writes, frame) + c.mu.Unlock() + + if c.onWrite != nil { + c.onWrite(frame) + } + return len(b), nil +} + +func (c *scriptedConn) RemoteAddr() net.Addr { + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1} +} + +func (c *scriptedConn) Close() error { + c.closeOnce.Do(func() { close(c.closed) }) + return nil +} + +func (c *scriptedConn) Protocol() string { + return "test" +} + +func (c *scriptedConn) deliver(t *testing.T, frame []byte) { + t.Helper() + select { + case c.reads <- frame: + case <-time.After(testWaitTimeout): + t.Fatal("scripted connection read queue is full") + } +} + +func (c *scriptedConn) isClosed() bool { + select { + case <-c.closed: + return true + default: + return false + } +} + +func (c *scriptedConn) frames() [][]byte { + c.mu.Lock() + defer c.mu.Unlock() + return append([][]byte(nil), c.writes...) +} + +func (c *scriptedConn) waitFrames(t *testing.T, n int) [][]byte { + t.Helper() + deadline := time.Now().Add(testWaitTimeout) + for { + frames := c.frames() + if len(frames) >= n { + return frames + } + if time.Now().After(deadline) { + t.Fatalf("expected %d written frames, got %d", n, len(frames)) + } + time.Sleep(5 * time.Millisecond) + } +} + +func (c *scriptedConn) waitClosed(t *testing.T) { + t.Helper() + select { + case <-c.closed: + case <-time.After(testWaitTimeout): + t.Fatal("connection was not closed") + } +} + +func (v *recordingValidator) Validate(credentials any) error { + v.mu.Lock() + defer v.mu.Unlock() + + payload, _ := credentials.([]byte) + v.payloads = append(v.payloads, append([]byte(nil), payload...)) + return v.err +} + +func (v *recordingValidator) calls() [][]byte { + v.mu.Lock() + defer v.mu.Unlock() + return append([][]byte(nil), v.payloads...) +} + +func TestRelay_AuthV1ResponseIsUnchanged(t *testing.T) { + validator := &recordingValidator{} + relay := newTestRelay(t, validator) + peerID := messages.HashID("v1-client") + token := []byte("v1-token") + auth, err := messages.MarshalAuthMsg(peerID, token) + require.NoError(t, err) + + conn := newScriptedConn() + conn.deliver(t, auth) + relay.Accept(conn) + + frames := conn.waitFrames(t, 1) + want, err := messages.MarshalAuthResponse(relay.instanceURL.String()) + require.NoError(t, err) + assert.Equal(t, want, frames[0], "a v1 client must receive the same response bytes as today") + assert.Equal(t, [][]byte{token}, validator.calls(), "the validator must receive the v1 token bytes") + assert.False(t, conn.isClosed(), "an authenticated v1 connection must stay open") + + item, ok := relay.store.Peer(peerID) + require.True(t, ok, "the peer must be stored after the response") + assert.Equal(t, messages.ProtocolV1, item.(*Peer).Version(), "a v1 client selects v1") +} + +func TestRelay_AuthEnvelopeSelectsV1(t *testing.T) { + cases := []struct { + name string + version messages.ProtocolVersion + ext []byte + }{ + {"v2 without extension", 2, nil}, + {"v2 with json extension", 2, []byte(`{"cred2":"value"}`)}, + {"v255 with opaque extension", 255, []byte("opaque-extension-bytes")}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + validator := &recordingValidator{} + relay := newTestRelay(t, validator) + peerID := messages.HashID("future-client") + token := []byte("future-token") + auth, err := messages.MarshalAuthEnvelope(tc.version, peerID, token, tc.ext) + require.NoError(t, err) + + conn := newScriptedConn() + conn.deliver(t, auth) + relay.Accept(conn) + + frames := conn.waitFrames(t, 1) + want, err := messages.MarshalAuthResponse(relay.instanceURL.String()) + require.NoError(t, err) + assert.Equal(t, want, frames[0], "the response must carry the selected v1 header") + assert.Equal(t, [][]byte{token}, validator.calls(), "the validator must receive the token without the extension") + assert.False(t, conn.isClosed(), "a higher advertised version must not be rejected") + + item, ok := relay.store.Peer(peerID) + require.True(t, ok, "the peer must be stored after the response") + assert.Equal(t, messages.ProtocolV1, item.(*Peer).Version(), "the connection must use the selected v1") + }) + } +} + +func TestRelay_AuthRejectsInvalidFirstFrame(t *testing.T) { + peerID := messages.HashID("client") + validEnvelope, err := messages.MarshalAuthEnvelope(2, peerID, []byte("token"), []byte("ext")) + require.NoError(t, err) + validV1, err := messages.MarshalAuthMsg(peerID, []byte("token")) + require.NoError(t, err) + transportFirst, err := messages.MarshalTransportMsg(peerID, []byte("payload")) + require.NoError(t, err) + + versionZero := append([]byte(nil), validV1...) + versionZero[0] = 0 + + badMagic := append([]byte(nil), validEnvelope...) + badMagic[2] ^= 0xFF + + zeroTokenLen := append([]byte(nil), validEnvelope...) + binary.BigEndian.PutUint16(zeroTokenLen[authTokenLenOffset:], 0) + + tokenLenPastEnd := append([]byte(nil), validEnvelope...) + binary.BigEndian.PutUint16(tokenLenPastEnd[authTokenLenOffset:], uint16(len(validEnvelope))) + + cases := []struct { + name string + frame []byte + }{ + {"version zero", versionZero}, + {"reserved hello type", []byte{byte(messages.ProtocolV1), byte(messages.MsgTypeHello)}}, + {"unknown type before auth", []byte{byte(messages.ProtocolV1), 200}}, + {"transport before auth", transportFirst}, + {"bad magic in envelope", badMagic}, + {"zero token length", zeroTokenLen}, + {"token length past the frame", tokenLenPastEnd}, + {"envelope truncated at the length field", validEnvelope[:authTokenLenOffset+1]}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + validator := &recordingValidator{} + relay := newTestRelay(t, validator) + + conn := newScriptedConn() + conn.deliver(t, tc.frame) + relay.Accept(conn) + + assert.True(t, conn.isClosed(), "the connection must be closed") + assert.Empty(t, conn.frames(), "no response may be written") + assert.Empty(t, validator.calls(), "the validator must not run") + _, ok := relay.store.Peer(peerID) + assert.False(t, ok, "the peer must not be stored") + }) + } +} + +func TestRelay_AuthRejectsInvalidToken(t *testing.T) { + peerID := messages.HashID("client") + v1Auth, err := messages.MarshalAuthMsg(peerID, []byte("token")) + require.NoError(t, err) + envelopeAuth, err := messages.MarshalAuthEnvelope(2, peerID, []byte("token"), []byte(`{"ok":true}`)) + require.NoError(t, err) + + cases := []struct { + name string + frame []byte + }{ + {"v1 layout", v1Auth}, + {"envelope with well-formed extension", envelopeAuth}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + validator := &recordingValidator{err: errors.New("invalid token")} + relay := newTestRelay(t, validator) + + conn := newScriptedConn() + conn.deliver(t, tc.frame) + relay.Accept(conn) + + assert.True(t, conn.isClosed(), "the connection must be closed") + assert.Empty(t, conn.frames(), "no response may be written") + assert.Len(t, validator.calls(), 1, "the validator must run exactly once") + _, ok := relay.store.Peer(peerID) + assert.False(t, ok, "the peer must not be stored") + }) + } +} + +func TestRelay_AuthResponsePrecedesPublish(t *testing.T) { + relay := newTestRelay(t, &allow.Auth{}) + peerID := messages.HashID("client") + + onlineEvents := make(chan []messages.PeerID, 1) + watcher := relay.notifier.NewListener(func(ids []messages.PeerID) { onlineEvents <- ids }, func([]messages.PeerID) {}) + t.Cleanup(func() { relay.notifier.RemoveListener(watcher) }) + watcher.AddInterestedPeers([]messages.PeerID{peerID}) + + var storedAtResponse, notifiedAtResponse bool + conn := newScriptedConn() + conn.onWrite = func(frame []byte) { + if frame[1] != byte(messages.MsgTypeAuthResponse) { + return + } + _, storedAtResponse = relay.store.Peer(peerID) + notifiedAtResponse = len(onlineEvents) > 0 + } + + auth, err := messages.MarshalAuthMsg(peerID, []byte("token")) + require.NoError(t, err) + conn.deliver(t, auth) + relay.Accept(conn) + + conn.waitFrames(t, 1) + assert.False(t, storedAtResponse, "the peer must not be in the store when the response is written") + assert.False(t, notifiedAtResponse, "no online notification may precede the response") + + select { + case ids := <-onlineEvents: + assert.Equal(t, []messages.PeerID{peerID}, ids, "the peer must be published after the response") + case <-time.After(testWaitTimeout): + t.Fatal("the peer was not published after the response") + } +} + +func TestRelay_ReconnectClosesPreviousConnection(t *testing.T) { + relay := newTestRelay(t, &allow.Auth{}) + first, peerID := acceptPeer(t, relay, "client") + + auth, err := messages.MarshalAuthMsg(peerID, []byte("token")) + require.NoError(t, err) + second := newScriptedConn() + second.deliver(t, auth) + relay.Accept(second) + + second.waitFrames(t, 1) + first.waitClosed(t) + assert.False(t, second.isClosed(), "the reconnecting client must keep its new connection") + item, ok := relay.store.Peer(peerID) + require.True(t, ok, "the peer must stay stored after a reconnect") + assert.Same(t, second, item.(*Peer).conn, "the store must hold the new connection") +} + +func newTestRelay(t *testing.T, validator Validator) *Relay { + t.Helper() + relay, err := NewRelay(Config{ + ExposedAddress: "rel://127.0.0.1:1234", + AuthValidator: validator, + }) + require.NoError(t, err) + t.Cleanup(func() { relay.Shutdown(context.Background()) }) + return relay +} + +func acceptPeer(t *testing.T, relay *Relay, name string) (*scriptedConn, messages.PeerID) { + t.Helper() + peerID := messages.HashID(name) + auth, err := messages.MarshalAuthMsg(peerID, []byte("token")) + require.NoError(t, err) + + conn := newScriptedConn() + conn.deliver(t, auth) + relay.Accept(conn) + conn.waitFrames(t, 1) + return conn, peerID +} diff --git a/relay/server/peer.go b/relay/server/peer.go index 8376cdfa7..887695d94 100644 --- a/relay/server/peer.go +++ b/relay/server/peer.go @@ -27,6 +27,8 @@ type Peer struct { metrics *metrics.Metrics log *log.Entry id messages.PeerID + version messages.ProtocolVersion + enc messages.Encoder conn listener.Conn connMu sync.RWMutex store *store.Store @@ -42,12 +44,14 @@ type Peer struct { } // NewPeer creates a new Peer instance and prepare custom logging -func NewPeer(metrics *metrics.Metrics, id messages.PeerID, conn listener.Conn, store *store.Store, notifier *store.PeerNotifier) *Peer { +func NewPeer(metrics *metrics.Metrics, id messages.PeerID, version messages.ProtocolVersion, conn listener.Conn, store *store.Store, notifier *store.PeerNotifier) *Peer { ctx, cancel := context.WithCancel(context.Background()) p := &Peer{ metrics: metrics, log: log.WithField("peer_id", id.String()), id: id, + version: version, + enc: messages.NewEncoder(version), conn: conn, store: store, notifier: notifier, @@ -95,16 +99,15 @@ func (p *Peer) Work() { msg := buf[:n] - _, err = messages.ValidateVersion(msg) - if err != nil { - p.log.Warnf("failed to validate protocol version: %s", err) - return + if err := messages.ValidateConnectionVersion(msg, p.version); err != nil { + p.log.Tracef("drop message: %s", err) + continue } msgType, err := messages.DetermineClientMessageType(msg) if err != nil { - p.log.Errorf("failed to determine message type: %s", err) - return + p.log.Tracef("drop message: %s", err) + continue } p.handleMsgType(ctx, msgType, hc, n, msg) @@ -115,6 +118,11 @@ func (p *Peer) ID() messages.PeerID { return p.id } +// Version returns the protocol version selected for this connection +func (p *Peer) Version() messages.ProtocolVersion { + return p.version +} + func (p *Peer) handleMsgType(ctx context.Context, msgType messages.MsgType, hc *healthcheck.Sender, n int, msg []byte) { switch msgType { case messages.MsgTypeHealthCheck: @@ -133,7 +141,7 @@ func (p *Peer) handleMsgType(ctx context.Context, msgType messages.MsgType, hc * case messages.MsgTypeUnsubscribePeerState: p.handleUnsubscribePeerState(msg) default: - p.log.Warnf("received unexpected message type: %s", msgType) + p.log.Tracef("drop %s message: not allowed after authentication", msgType) } } @@ -149,7 +157,7 @@ func (p *Peer) Write(ctx context.Context, b []byte) (int, error) { func (p *Peer) CloseGracefully(ctx context.Context) { p.connMu.Lock() defer p.connMu.Unlock() - err := p.writeWithTimeout(ctx, messages.MarshalCloseMsg()) + err := p.writeWithTimeout(ctx, p.enc.CloseMsg()) if err != nil { p.log.Errorf("failed to send close message to peer: %s", p.String()) } @@ -187,7 +195,7 @@ func (p *Peer) handleHealthcheckEvents(ctx context.Context, hc *healthcheck.Send for { select { case <-hc.HealthCheck: - _, err := p.Write(ctx, messages.MarshalHealthcheck()) + _, err := p.Write(ctx, p.enc.Healthcheck()) if err != nil { p.log.Errorf("failed to send healthcheck message: %s", err) return @@ -220,7 +228,7 @@ func (p *Peer) handleTransportMsg(msg []byte) { } dp := item.(*Peer) - err = messages.UpdateTransportMsg(msg, p.id) + err = messages.UpdateTransportMsg(msg, dp.version, p.id) if err != nil { p.log.Errorf("failed to update transport message: %s", err) return @@ -267,7 +275,7 @@ func (p *Peer) handleUnsubscribePeerState(msg []byte) { } func (p *Peer) sendPeersOnline(peers []messages.PeerID) { - msgs, err := messages.MarshalPeersOnline(peers) + msgs, err := p.enc.PeersOnline(peers) if err != nil { p.log.Errorf("failed to marshal peer location message: %s", err) return @@ -284,7 +292,7 @@ func (p *Peer) sendPeersWentOffline(peers []messages.PeerID) { p.notificationMutex.Lock() defer p.notificationMutex.Unlock() - msgs, err := messages.MarshalPeersWentOffline(peers) + msgs, err := p.enc.PeersWentOffline(peers) if err != nil { p.log.Errorf("failed to marshal peer location message: %s", err) return diff --git a/relay/server/peer_test.go b/relay/server/peer_test.go new file mode 100644 index 000000000..3b6a8de41 --- /dev/null +++ b/relay/server/peer_test.go @@ -0,0 +1,99 @@ +package server + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/shared/relay/auth/allow" + "github.com/netbirdio/netbird/shared/relay/messages" +) + +func TestPeer_DropsFramesItCannotHandle(t *testing.T) { + secondAuth, err := messages.MarshalAuthMsg(messages.HashID("alice"), []byte("token")) + require.NoError(t, err) + + cases := []struct { + name string + frame []byte + }{ + {"unknown message type", []byte{byte(messages.ProtocolV1), 200, 0xDE, 0xAD}}, + {"version above the selected one", []byte{2, byte(messages.MsgTypeHealthCheck)}}, + {"version zero", []byte{0, byte(messages.MsgTypeHealthCheck)}}, + {"header shorter than two bytes", []byte{byte(messages.ProtocolV1)}}, + {"reserved hello type", []byte{byte(messages.ProtocolV1), byte(messages.MsgTypeHello)}}, + {"server to client peers online type", []byte{byte(messages.ProtocolV1), byte(messages.MsgTypePeersOnline)}}, + {"auth after authentication", secondAuth}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + relay := newTestRelay(t, &allow.Auth{}) + alice, _ := acceptPeer(t, relay, "alice") + _, bobID := acceptPeer(t, relay, "bob") + + alice.deliver(t, tc.frame) + subscribe, err := messages.MarshalSubPeerStateMsg([]messages.PeerID{bobID}) + require.NoError(t, err) + alice.deliver(t, subscribe[0]) + + // The subscription answer proves the loop survived the dropped frame before it. + frames := alice.waitFrames(t, 2) + assert.False(t, alice.isClosed(), "a dropped frame must not close the connection") + online := frames[1] + assert.Equal(t, byte(messages.ProtocolV1), online[0], "notifications must carry the connection version") + assert.Equal(t, byte(messages.MsgTypePeersOnline), online[1], "the subscription must be answered") + ids, err := messages.UnmarshalPeersOnlineMsg(online) + require.NoError(t, err) + assert.Equal(t, []messages.PeerID{bobID}, ids, "the online peer must be reported") + }) + } +} + +func TestPeer_ClosesOnClientCloseMessage(t *testing.T) { + relay := newTestRelay(t, &allow.Auth{}) + conn, peerID := acceptPeer(t, relay, "alice") + + conn.deliver(t, messages.MarshalCloseMsg()) + + conn.waitClosed(t) + assert.Len(t, conn.frames(), 1, "only the auth response may be written") + assert.Eventually(t, func() bool { + _, ok := relay.store.Peer(peerID) + return !ok + }, testWaitTimeout, 5*time.Millisecond, "the peer must be removed from the store") +} + +func TestPeer_ForwardsTransportWithSenderAndDestinationVersion(t *testing.T) { + relay := newTestRelay(t, &allow.Auth{}) + alice, aliceID := acceptPeer(t, relay, "alice") + bob, bobID := acceptPeer(t, relay, "bob") + + payload := []byte("hello bob") + transport, err := messages.MarshalTransportMsg(bobID, payload) + require.NoError(t, err) + alice.deliver(t, transport) + + frames := bob.waitFrames(t, 2) + forwarded := frames[1] + assert.Equal(t, byte(messages.ProtocolV1), forwarded[0], "the forwarded frame must carry the destination's version") + assert.Equal(t, byte(messages.MsgTypeTransport), forwarded[1], "the forwarded frame must stay a transport message") + senderID, gotPayload, err := messages.UnmarshalTransportMsg(forwarded) + require.NoError(t, err) + assert.Equal(t, aliceID, *senderID, "the peer ID must be rewritten to the sender") + assert.Equal(t, payload, gotPayload, "the payload must be preserved") + assert.False(t, alice.isClosed(), "the sender must stay connected") +} + +func TestPeer_GracefulCloseCarriesConnectionVersion(t *testing.T) { + relay := newTestRelay(t, &allow.Auth{}) + conn, _ := acceptPeer(t, relay, "alice") + + relay.Shutdown(context.Background()) + + frames := conn.waitFrames(t, 2) + assert.Equal(t, messages.NewEncoder(messages.ProtocolV1).CloseMsg(), frames[1], "the close message must carry the connection version") + assert.True(t, conn.isClosed(), "the connection must be closed after the close message") +} diff --git a/relay/server/relay.go b/relay/server/relay.go index 58f9db840..bd580efd0 100644 --- a/relay/server/relay.go +++ b/relay/server/relay.go @@ -150,8 +150,17 @@ func (r *Relay) Accept(conn listener.Conn) { return } - peer := NewPeer(r.metrics, *peerID, conn, r.store, r.notifier) - peer.log.Infof("peer connected from: %s", conn.RemoteAddr()) + if err := h.handshakeResponse(hsCtx); err != nil { + log.Errorf("failed to send handshake response, close connection: %s", err) + if cErr := conn.Close(); cErr != nil { + log.Errorf("failed to close connection, %s: %s", conn.RemoteAddr(), cErr) + } + return + } + r.metrics.RecordAuthenticationTime(time.Since(acceptTime)) + + peer := NewPeer(r.metrics, *peerID, h.selectedVersion, conn, r.store, r.notifier) + peer.log.Infof("peer connected from: %s, advertised protocol version: %d, selected: %d", conn.RemoteAddr(), h.advertisedVersion, h.selectedVersion) storeTime := time.Now() if isReconnection := r.store.AddPeer(peer); isReconnection { r.metrics.RecordPeerReconnection() @@ -161,6 +170,7 @@ func (r *Relay) Accept(conn listener.Conn) { transport := conn.Protocol() r.metrics.RecordPeerStoreTime(time.Since(storeTime)) r.metrics.PeerConnected(peer.String(), transport) + r.metrics.RecordProtocolVersion(h.advertisedVersion, h.selectedVersion) go func() { peer.Work() if deleted := r.store.DeletePeer(peer); deleted { @@ -169,12 +179,6 @@ func (r *Relay) Accept(conn listener.Conn) { peer.log.Debugf("relay connection closed") r.metrics.PeerDisconnected(peer.String(), transport) }() - - if err := h.handshakeResponse(hsCtx); err != nil { - log.Errorf("failed to send handshake response, close peer: %s", err) - peer.Close() - } - r.metrics.RecordAuthenticationTime(time.Since(acceptTime)) } // Shutdown closes the relay server diff --git a/relay/server/ws_version_test.go b/relay/server/ws_version_test.go new file mode 100644 index 000000000..bd5fa7303 --- /dev/null +++ b/relay/server/ws_version_test.go @@ -0,0 +1,97 @@ +package server + +import ( + "context" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/shared/relay" + "github.com/netbirdio/netbird/shared/relay/messages" +) + +func TestServer_FutureClientOverWebSocket(t *testing.T) { + addr := freeAddress(t) + srv := newTestServer(t, addr) + errChan := make(chan error, 1) + go func() { + errChan <- srv.Listen(ListenerConfig{Address: addr}) + }() + waitForListeners(t, srv, errChan) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + instanceURL := srv.InstanceURL() + want, err := messages.MarshalAuthResponse(instanceURL.String()) + require.NoError(t, err) + + // A synthetic client advertising v2 with an extension area that a v1 server does not read. + future := dialRelayWebSocket(t, ctx, addr) + futureID := messages.HashID("future-client") + futureAuth, err := messages.MarshalAuthEnvelope(2, futureID, []byte("future-token"), []byte(`{"cred2":"value"}`)) + require.NoError(t, err) + require.NoError(t, future.Write(ctx, websocket.MessageBinary, futureAuth)) + assert.Equal(t, want, readFrame(t, ctx, future), "a v2 client must receive the v1 response") + + legacy := dialRelayWebSocket(t, ctx, addr) + legacyID := messages.HashID("legacy-client") + legacyAuth, err := messages.MarshalAuthMsg(legacyID, []byte("legacy-token")) + require.NoError(t, err) + require.NoError(t, legacy.Write(ctx, websocket.MessageBinary, legacyAuth)) + assert.Equal(t, want, readFrame(t, ctx, legacy), "a v1 client must receive the response it gets today") + + toLegacy, err := messages.MarshalTransportMsg(legacyID, []byte("from future")) + require.NoError(t, err) + require.NoError(t, future.Write(ctx, websocket.MessageBinary, toLegacy)) + assertTransportFrame(t, readFrame(t, ctx, legacy), futureID, "from future") + + toFuture, err := messages.MarshalTransportMsg(futureID, []byte("from legacy")) + require.NoError(t, err) + require.NoError(t, legacy.Write(ctx, websocket.MessageBinary, toFuture)) + assertTransportFrame(t, readFrame(t, ctx, future), legacyID, "from legacy") + + // After the response the connection is v1. A frame with another version is dropped and the + // connection stays usable, which the transport message right after it proves. + require.NoError(t, future.Write(ctx, websocket.MessageBinary, []byte{2, byte(messages.MsgTypeHealthCheck)})) + toLegacyAgain, err := messages.MarshalTransportMsg(legacyID, []byte("after the dropped frame")) + require.NoError(t, err) + require.NoError(t, future.Write(ctx, websocket.MessageBinary, toLegacyAgain)) + assertTransportFrame(t, readFrame(t, ctx, legacy), futureID, "after the dropped frame") + + require.NoError(t, future.Write(ctx, websocket.MessageBinary, messages.MarshalCloseMsg())) + require.NoError(t, legacy.Write(ctx, websocket.MessageBinary, messages.MarshalCloseMsg())) +} + +func dialRelayWebSocket(t *testing.T, ctx context.Context, addr string) *websocket.Conn { + t.Helper() + conn, resp, err := websocket.Dial(ctx, "ws://"+addr+relay.WebSocketURLPath, nil) //nolint:bodyclose + require.NoError(t, err) + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + t.Cleanup(func() { _ = conn.CloseNow() }) + return conn +} + +func readFrame(t *testing.T, ctx context.Context, conn *websocket.Conn) []byte { + t.Helper() + msgType, frame, err := conn.Read(ctx) + require.NoError(t, err) + require.Equal(t, websocket.MessageBinary, msgType, "relay frames are binary messages") + return frame +} + +func assertTransportFrame(t *testing.T, frame []byte, sender messages.PeerID, payload string) { + t.Helper() + require.GreaterOrEqual(t, len(frame), 2, "a frame must carry the protocol header") + assert.Equal(t, byte(messages.ProtocolV1), frame[0], "the forwarded frame must carry the destination's version") + assert.Equal(t, byte(messages.MsgTypeTransport), frame[1], "the forwarded frame must be a transport message") + id, got, err := messages.UnmarshalTransportMsg(frame) + require.NoError(t, err) + assert.Equal(t, sender, *id, "the peer ID must be rewritten to the sender") + assert.Equal(t, payload, string(got), "the payload must be preserved") +} diff --git a/shared/relay/messages/message.go b/shared/relay/messages/message.go index fdbf2a6e7..e25571edd 100644 --- a/shared/relay/messages/message.go +++ b/shared/relay/messages/message.go @@ -2,16 +2,27 @@ package messages import ( "bytes" + "encoding/binary" "errors" "fmt" + "math" ) const ( - MaxHandshakeSize = 212 + // TODO: MaxHandshakeSize is the buffer size of the first read on every connection before + // authentication. Weigh it against memory use under many unauthenticated connections; a + // smaller limit may still fit the extension area. + MaxHandshakeSize = 8192 MaxHandshakeRespSize = 8192 MaxMessageSize = 8820 - CurrentProtocolVersion = 1 + ProtocolV1 ProtocolVersion = 1 + // CurrentProtocolVersion is the highest protocol version this build supports. + CurrentProtocolVersion = ProtocolV1 + + // authEnvelopeMinVersion is the first advertised version whose Auth message carries a + // length-prefixed token followed by an extension area. + authEnvelopeMinVersion ProtocolVersion = 2 MsgTypeUnknown MsgType = 0 // MsgTypeHello and MsgTypeHelloResponse are the removed legacy handshake @@ -38,10 +49,13 @@ const ( // auth message sizeOfMagicByte = 4 + sizeOfTokenLen = 2 headerSizeAuth = sizeOfMagicByte + peerIDSize offsetMagicByte = sizeOfProtoHeader offsetAuthPeerID = sizeOfProtoHeader + sizeOfMagicByte headerTotalSizeAuth = sizeOfProtoHeader + headerSizeAuth + offsetAuthTokenLen = headerTotalSizeAuth + offsetAuthToken = offsetAuthTokenLen + sizeOfTokenLen // transport headerSizeTransport = peerIDSize @@ -52,6 +66,10 @@ const ( var ( ErrInvalidMessageLength = errors.New("invalid message length") ErrUnsupportedVersion = errors.New("unsupported version") + ErrInvalidTokenLength = errors.New("invalid token length") + ErrReservedMsgType = errors.New("reserved message type") + ErrMsgTypeDirection = errors.New("message type not valid from client") + ErrUnknownMsgType = errors.New("unknown message type") magicHeader = []byte{0x21, 0x12, 0xA4, 0x42} @@ -60,6 +78,19 @@ var ( type MsgType byte +// ProtocolVersion is the relay wire protocol version carried in the first byte of every frame. +type ProtocolVersion uint8 + +// Encoder marshals server-originated messages with the protocol version selected for one connection. +type Encoder struct { + version ProtocolVersion +} + +// NewEncoder returns an Encoder that stamps every message with the given protocol version. +func NewEncoder(version ProtocolVersion) Encoder { + return Encoder{version: version} +} + func (m MsgType) String() string { switch m { case MsgTypeHello: @@ -89,18 +120,99 @@ func (m MsgType) String() string { } } -// ValidateVersion checks if the given version is supported by the protocol -func ValidateVersion(msg []byte) (int, error) { +// Version returns the protocol version the Encoder stamps on its messages. +func (e Encoder) Version() ProtocolVersion { + return e.version +} + +// AuthResponse creates the response to a successful Auth message. The payload is the instance URL +// of the server, which the clients use to choose a common Relay server. +func (e Encoder) AuthResponse(address string) ([]byte, error) { + ab := []byte(address) + msg := make([]byte, sizeOfProtoHeader, sizeOfProtoHeader+len(ab)) + + msg[0] = byte(e.version) + msg[1] = byte(MsgTypeAuthResponse) + + msg = append(msg, ab...) + + if len(msg) > MaxHandshakeRespSize { + return nil, fmt.Errorf("invalid message length: %d", len(msg)) + } + + return msg, nil +} + +// CloseMsg creates a close message for a graceful disconnect. +func (e Encoder) CloseMsg() []byte { + return []byte{ + byte(e.version), + byte(MsgTypeClose), + } +} + +// Healthcheck creates a health check message. +func (e Encoder) Healthcheck() []byte { + return []byte{ + byte(e.version), + byte(MsgTypeHealthCheck), + } +} + +// PeersOnline creates the peers online notification messages. +func (e Encoder) PeersOnline(ids []PeerID) ([][]byte, error) { + return marshalPeerIDs(e.version, ids, byte(MsgTypePeersOnline)) +} + +// PeersWentOffline creates the peers went offline notification messages. +func (e Encoder) PeersWentOffline(ids []PeerID) ([][]byte, error) { + return marshalPeerIDs(e.version, ids, byte(MsgTypePeersWentOffline)) +} + +// ReadVersion returns the protocol version carried in the frame header. Zero is invalid. +func ReadVersion(msg []byte) (ProtocolVersion, error) { if len(msg) < sizeOfProtoHeader { return 0, ErrInvalidMessageLength } - version := int(msg[0]) - if version != CurrentProtocolVersion { + version := ProtocolVersion(msg[0]) + if version == 0 { return 0, fmt.Errorf("%d: %w", version, ErrUnsupportedVersion) } return version, nil } +// SelectVersion returns the highest protocol version both sides support. +func SelectVersion(advertised, supported ProtocolVersion) ProtocolVersion { + if advertised < supported { + return advertised + } + return supported +} + +// ValidateVersion checks that the frame carries the protocol version of this build. +func ValidateVersion(msg []byte) (int, error) { + version, err := ReadVersion(msg) + if err != nil { + return 0, err + } + if version != CurrentProtocolVersion { + return 0, fmt.Errorf("%d: %w", version, ErrUnsupportedVersion) + } + return int(version), nil +} + +// ValidateConnectionVersion checks that the frame carries the version selected for the connection. +func ValidateConnectionVersion(msg []byte, selected ProtocolVersion) error { + version, err := ReadVersion(msg) + if err != nil { + return err + } + if version != selected { + return fmt.Errorf("%d instead of selected %d: %w", version, selected, ErrUnsupportedVersion) + } + return nil +} + // DetermineClientMessageType determines the message type from the first the message func DetermineClientMessageType(msg []byte) (MsgType, error) { if len(msg) < sizeOfProtoHeader { @@ -117,8 +229,12 @@ func DetermineClientMessageType(msg []byte) (MsgType, error) { MsgTypeSubscribePeerState, MsgTypeUnsubscribePeerState: return msgType, nil + case MsgTypeHello, MsgTypeHelloResponse: + return MsgTypeUnknown, fmt.Errorf("%d: %w", msgType, ErrReservedMsgType) + case MsgTypeAuthResponse, MsgTypePeersOnline, MsgTypePeersWentOffline: + return MsgTypeUnknown, fmt.Errorf("%d: %w", msgType, ErrMsgTypeDirection) default: - return MsgTypeUnknown, fmt.Errorf("invalid msg type %d", msgType) + return MsgTypeUnknown, fmt.Errorf("%d: %w", msgType, ErrUnknownMsgType) } } @@ -162,8 +278,41 @@ func MarshalAuthMsg(peerID PeerID, authPayload []byte) ([]byte, error) { return msg, nil } -// UnmarshalAuthMsg extracts peerID and the auth payload from the message +// MarshalAuthEnvelope creates the Auth message layout used by clients advertising protocol +// version 2 or higher: the token is length-prefixed and followed by an extension area that +// servers without support for it ignore. +func MarshalAuthEnvelope(version ProtocolVersion, peerID PeerID, token, ext []byte) ([]byte, error) { + if version < authEnvelopeMinVersion { + return nil, fmt.Errorf("protocol version %d does not use the auth envelope", version) + } + if len(token) == 0 || len(token) > math.MaxUint16 { + return nil, fmt.Errorf("%d: %w", len(token), ErrInvalidTokenLength) + } + + size := offsetAuthToken + len(token) + len(ext) + if size > MaxHandshakeSize { + return nil, fmt.Errorf("too large auth message: %d", size) + } + + msg := make([]byte, size) + msg[0] = byte(version) + msg[1] = byte(MsgTypeAuth) + copy(msg[offsetMagicByte:], magicHeader) + copy(msg[offsetAuthPeerID:], peerID[:]) + binary.BigEndian.PutUint16(msg[offsetAuthTokenLen:], uint16(len(token))) + copy(msg[offsetAuthToken:], token) + copy(msg[offsetAuthToken+len(token):], ext) + return msg, nil +} + +// UnmarshalAuthMsg extracts the peer ID and the auth token from an Auth message of any +// protocol version. The extension area of the envelope layout is not returned. func UnmarshalAuthMsg(msg []byte) (*PeerID, []byte, error) { + version, err := ReadVersion(msg) + if err != nil { + return nil, nil, err + } + if len(msg) < headerTotalSizeAuth { return nil, nil, ErrInvalidMessageLength } @@ -174,7 +323,15 @@ func UnmarshalAuthMsg(msg []byte) (*PeerID, []byte, error) { } peerID := PeerID(msg[offsetAuthPeerID:headerTotalSizeAuth]) - return &peerID, msg[headerTotalSizeAuth:], nil + if version < authEnvelopeMinVersion { + return &peerID, msg[headerTotalSizeAuth:], nil + } + + token, err := unmarshalAuthEnvelopeToken(msg) + if err != nil { + return nil, nil, err + } + return &peerID, token, nil } // MarshalAuthResponse creates a response message to the auth. @@ -182,19 +339,7 @@ func UnmarshalAuthMsg(msg []byte) (*PeerID, []byte, error) { // instance URL. This URL will be used by choose the common Relay server in case if the peers are in different Relay // servers. func MarshalAuthResponse(address string) ([]byte, error) { - ab := []byte(address) - msg := make([]byte, sizeOfProtoHeader, sizeOfProtoHeader+len(ab)) - - msg[0] = byte(CurrentProtocolVersion) - msg[1] = byte(MsgTypeAuthResponse) - - msg = append(msg, ab...) - - if len(msg) > MaxHandshakeRespSize { - return nil, fmt.Errorf("invalid message length: %d", len(msg)) - } - - return msg, nil + return NewEncoder(CurrentProtocolVersion).AuthResponse(address) } // UnmarshalAuthResponse it is a confirmation message to auth success @@ -209,10 +354,7 @@ func UnmarshalAuthResponse(msg []byte) (string, error) { // The close message is used to close the connection gracefully between the client and the server. The server and the // client can send this message. After receiving this message, the server or client will close the connection. func MarshalCloseMsg() []byte { - return []byte{ - byte(CurrentProtocolVersion), - byte(MsgTypeClose), - } + return NewEncoder(CurrentProtocolVersion).CloseMsg() } // MarshalTransportMsg creates a transport message. @@ -252,13 +394,14 @@ func UnmarshalTransportID(buf []byte) (*PeerID, error) { return &id, nil } -// UpdateTransportMsg updates the peerID in the transport message. -// With this function the server can reuse the given byte slice to update the peerID in the transport message. So do -// need to allocate a new byte slice. -func UpdateTransportMsg(msg []byte, peerID PeerID) error { +// UpdateTransportMsg rewrites the transport message in place for forwarding: the header carries +// the protocol version of the destination connection and the peer ID field the sender's ID. +// The payload bytes are left untouched, so no new byte slice is allocated. +func UpdateTransportMsg(msg []byte, version ProtocolVersion, peerID PeerID) error { if len(msg) < offsetTransportID+peerIDSize { return ErrInvalidMessageLength } + msg[0] = byte(version) copy(msg[offsetTransportID:], peerID[:]) return nil } @@ -269,3 +412,20 @@ func UpdateTransportMsg(msg []byte, peerID PeerID) error { func MarshalHealthcheck() []byte { return healthCheckMsg } + +func unmarshalAuthEnvelopeToken(msg []byte) ([]byte, error) { + if len(msg) < offsetAuthToken { + return nil, ErrInvalidMessageLength + } + + tokenLen := int(binary.BigEndian.Uint16(msg[offsetAuthTokenLen:offsetAuthToken])) + if tokenLen == 0 { + return nil, fmt.Errorf("%d: %w", tokenLen, ErrInvalidTokenLength) + } + + end := offsetAuthToken + tokenLen + if end > len(msg) { + return nil, fmt.Errorf("%d exceeds the message: %w", tokenLen, ErrInvalidTokenLength) + } + return msg[offsetAuthToken:end], nil +} diff --git a/shared/relay/messages/message_test.go b/shared/relay/messages/message_test.go index 26a504063..b3727c86a 100644 --- a/shared/relay/messages/message_test.go +++ b/shared/relay/messages/message_test.go @@ -1,7 +1,11 @@ package messages import ( + "encoding/binary" "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestDetermineClientMessageTypeRejectsHello(t *testing.T) { @@ -119,3 +123,245 @@ func TestMarshalHealthcheck(t *testing.T) { t.Errorf("expected %d, got %d", MsgTypeHealthCheck, msgType) } } + +func TestReadVersion(t *testing.T) { + _, err := ReadVersion([]byte{byte(ProtocolV1)}) + assert.ErrorIs(t, err, ErrInvalidMessageLength, "a header needs two bytes") + + _, err = ReadVersion([]byte{0, byte(MsgTypeAuth)}) + assert.ErrorIs(t, err, ErrUnsupportedVersion, "version zero is invalid") + + for _, version := range []ProtocolVersion{1, 2, 128, 255} { + got, err := ReadVersion([]byte{byte(version), byte(MsgTypeAuth)}) + require.NoError(t, err) + assert.Equal(t, version, got, "the advertised version must be returned unchanged") + } +} + +func TestSelectVersion(t *testing.T) { + assert.Equal(t, ProtocolV1, SelectVersion(1, 1), "equal versions select themselves") + assert.Equal(t, ProtocolV1, SelectVersion(2, 1), "a higher advertised version selects the server maximum") + assert.Equal(t, ProtocolV1, SelectVersion(255, 1), "the highest advertised version selects the server maximum") + assert.Equal(t, ProtocolVersion(2), SelectVersion(2, 3), "a lower advertised version selects itself") + assert.Equal(t, ProtocolVersion(2), SelectVersion(3, 2), "a higher advertised version selects the server maximum") +} + +func TestValidateVersionAcceptsOnlyCurrent(t *testing.T) { + _, err := ValidateVersion([]byte{byte(CurrentProtocolVersion), byte(MsgTypeHealthCheck)}) + require.NoError(t, err) + + _, err = ValidateVersion([]byte{byte(CurrentProtocolVersion) + 1, byte(MsgTypeHealthCheck)}) + assert.ErrorIs(t, err, ErrUnsupportedVersion, "the client build accepts only its own version") + + _, err = ValidateVersion([]byte{0, byte(MsgTypeHealthCheck)}) + assert.ErrorIs(t, err, ErrUnsupportedVersion, "version zero is invalid") +} + +func TestValidateConnectionVersion(t *testing.T) { + require.NoError(t, ValidateConnectionVersion([]byte{1, byte(MsgTypeHealthCheck)}, ProtocolV1)) + assert.ErrorIs(t, ValidateConnectionVersion([]byte{2, byte(MsgTypeHealthCheck)}, ProtocolV1), ErrUnsupportedVersion, "a frame must carry the selected version") + assert.ErrorIs(t, ValidateConnectionVersion([]byte{1, byte(MsgTypeHealthCheck)}, 2), ErrUnsupportedVersion, "a frame below the selected version is rejected too") + assert.ErrorIs(t, ValidateConnectionVersion([]byte{0, byte(MsgTypeHealthCheck)}, ProtocolV1), ErrUnsupportedVersion, "version zero is invalid") + assert.ErrorIs(t, ValidateConnectionVersion([]byte{1}, ProtocolV1), ErrInvalidMessageLength, "a header needs two bytes") +} + +func TestDetermineClientMessageTypeClassification(t *testing.T) { + rejected := []struct { + name string + msgType MsgType + wantErr error + }{ + {"hello is reserved", MsgTypeHello, ErrReservedMsgType}, + {"hello response is reserved", MsgTypeHelloResponse, ErrReservedMsgType}, + {"auth response is server to client only", MsgTypeAuthResponse, ErrMsgTypeDirection}, + {"peers online is server to client only", MsgTypePeersOnline, ErrMsgTypeDirection}, + {"peers went offline is server to client only", MsgTypePeersWentOffline, ErrMsgTypeDirection}, + {"zero is unknown", MsgTypeUnknown, ErrUnknownMsgType}, + {"a future type is unknown", MsgType(200), ErrUnknownMsgType}, + } + for _, tc := range rejected { + t.Run(tc.name, func(t *testing.T) { + got, err := DetermineClientMessageType([]byte{byte(ProtocolV1), byte(tc.msgType)}) + assert.ErrorIs(t, err, tc.wantErr) + assert.Equal(t, MsgTypeUnknown, got, "a rejected type must not be returned") + }) + } + + accepted := []MsgType{MsgTypeAuth, MsgTypeTransport, MsgTypeClose, MsgTypeHealthCheck, MsgTypeSubscribePeerState, MsgTypeUnsubscribePeerState} + for _, msgType := range accepted { + got, err := DetermineClientMessageType([]byte{byte(ProtocolV1), byte(msgType)}) + require.NoError(t, err) + assert.Equal(t, msgType, got, "a client message type must be returned") + } +} + +func TestAuthEnvelopeRoundTrip(t *testing.T) { + peerID := HashID("future-client") + token := []byte("0123456789abcdef") + cases := []struct { + name string + version ProtocolVersion + ext []byte + }{ + {"v2 without extension", 2, nil}, + {"v2 with json extension", 2, []byte(`{"cred2":"value"}`)}, + {"v255 with opaque extension", 255, []byte("opaque-extension-bytes")}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + msg, err := MarshalAuthEnvelope(tc.version, peerID, token, tc.ext) + require.NoError(t, err) + + version, err := ReadVersion(msg) + require.NoError(t, err) + assert.Equal(t, tc.version, version, "the header must carry the advertised version") + + msgType, err := DetermineClientMessageType(msg) + require.NoError(t, err) + assert.Equal(t, MsgType(MsgTypeAuth), msgType, "the envelope is an auth message") + + gotID, gotToken, err := UnmarshalAuthMsg(msg) + require.NoError(t, err) + assert.Equal(t, peerID, *gotID, "the peer ID must be preserved") + assert.Equal(t, token, gotToken, "the token must be returned without the extension bytes") + }) + } +} + +func TestMarshalAuthEnvelopeRejectsInvalidInput(t *testing.T) { + peerID := HashID("future-client") + + _, err := MarshalAuthEnvelope(ProtocolV1, peerID, []byte("token"), nil) + assert.Error(t, err, "v1 does not use the envelope layout") + + _, err = MarshalAuthEnvelope(2, peerID, nil, nil) + assert.ErrorIs(t, err, ErrInvalidTokenLength, "the token is mandatory") + + _, err = MarshalAuthEnvelope(2, peerID, []byte("token"), make([]byte, MaxHandshakeSize)) + assert.Error(t, err, "the envelope must fit the handshake size limit") +} + +func TestUnmarshalAuthMsgV1LayoutUnchanged(t *testing.T) { + peerID := HashID("legacy-client") + // In the v1 layout the token runs to the end of the frame, so bytes that would be a + // length prefix in the envelope belong to the token here. + payload := []byte{0x00, 0x05, 'a', 'b', 'c', 'd', 'e', 'f', 'g'} + msg, err := MarshalAuthMsg(peerID, payload) + require.NoError(t, err) + + gotID, gotToken, err := UnmarshalAuthMsg(msg) + require.NoError(t, err) + assert.Equal(t, peerID, *gotID, "the peer ID must be preserved") + assert.Equal(t, payload, gotToken, "the v1 token must be returned whole") +} + +func TestUnmarshalAuthMsgRejectsMalformedFrames(t *testing.T) { + peerID := HashID("future-client") + valid, err := MarshalAuthEnvelope(2, peerID, []byte("token"), []byte("ext")) + require.NoError(t, err) + + t.Run("zero token length", func(t *testing.T) { + msg := append([]byte(nil), valid...) + binary.BigEndian.PutUint16(msg[offsetAuthTokenLen:], 0) + _, _, err := UnmarshalAuthMsg(msg) + assert.ErrorIs(t, err, ErrInvalidTokenLength) + }) + + t.Run("token length past the end", func(t *testing.T) { + msg := append([]byte(nil), valid...) + binary.BigEndian.PutUint16(msg[offsetAuthTokenLen:], uint16(len(msg))) + _, _, err := UnmarshalAuthMsg(msg) + assert.ErrorIs(t, err, ErrInvalidTokenLength) + }) + + t.Run("truncated before the length field", func(t *testing.T) { + _, _, err := UnmarshalAuthMsg(valid[:headerTotalSizeAuth+1]) + assert.ErrorIs(t, err, ErrInvalidMessageLength) + }) + + t.Run("version zero", func(t *testing.T) { + msg := append([]byte(nil), valid...) + msg[0] = 0 + _, _, err := UnmarshalAuthMsg(msg) + assert.ErrorIs(t, err, ErrUnsupportedVersion) + }) + + t.Run("invalid magic", func(t *testing.T) { + msg := append([]byte(nil), valid...) + msg[offsetMagicByte] ^= 0xFF + _, _, err := UnmarshalAuthMsg(msg) + assert.Error(t, err) + }) + + t.Run("shorter than the fixed header", func(t *testing.T) { + _, _, err := UnmarshalAuthMsg(valid[:headerTotalSizeAuth-1]) + assert.ErrorIs(t, err, ErrInvalidMessageLength) + }) +} + +func TestEncoderStampsVersion(t *testing.T) { + for _, version := range []ProtocolVersion{1, 2, 7} { + enc := NewEncoder(version) + assert.Equal(t, version, enc.Version()) + assert.Equal(t, []byte{byte(version), byte(MsgTypeHealthCheck)}, enc.Healthcheck(), "healthcheck must carry the version") + assert.Equal(t, []byte{byte(version), byte(MsgTypeClose)}, enc.CloseMsg(), "close must carry the version") + + resp, err := enc.AuthResponse("rel://relay.example") + require.NoError(t, err) + assert.Equal(t, byte(version), resp[0], "auth response must carry the version") + assert.Equal(t, byte(MsgTypeAuthResponse), resp[1]) + addr, err := UnmarshalAuthResponse(resp) + require.NoError(t, err) + assert.Equal(t, "rel://relay.example", addr, "the auth response payload must be the address") + + online, err := enc.PeersOnline([]PeerID{HashID("a")}) + require.NoError(t, err) + require.Len(t, online, 1) + assert.Equal(t, byte(version), online[0][0], "peers online must carry the version") + assert.Equal(t, byte(MsgTypePeersOnline), online[0][1]) + + offline, err := enc.PeersWentOffline([]PeerID{HashID("a")}) + require.NoError(t, err) + require.Len(t, offline, 1) + assert.Equal(t, byte(version), offline[0][0], "peers went offline must carry the version") + assert.Equal(t, byte(MsgTypePeersWentOffline), offline[0][1]) + } +} + +func TestPackageEncodersUseCurrentVersion(t *testing.T) { + current := NewEncoder(CurrentProtocolVersion) + assert.Equal(t, current.CloseMsg(), MarshalCloseMsg()) + assert.Equal(t, current.Healthcheck(), MarshalHealthcheck()) + + want, err := current.AuthResponse("addr") + require.NoError(t, err) + got, err := MarshalAuthResponse("addr") + require.NoError(t, err) + assert.Equal(t, want, got) + + ids := []PeerID{HashID("a"), HashID("b")} + wantOnline, err := current.PeersOnline(ids) + require.NoError(t, err) + gotOnline, err := MarshalPeersOnline(ids) + require.NoError(t, err) + assert.Equal(t, wantOnline, gotOnline) +} + +func TestUpdateTransportMsgRewritesVersionAndSender(t *testing.T) { + dst := HashID("dst") + sender := HashID("sender") + payload := []byte("payload") + msg, err := MarshalTransportMsg(dst, payload) + require.NoError(t, err) + + require.NoError(t, UpdateTransportMsg(msg, ProtocolVersion(2), sender)) + assert.Equal(t, byte(2), msg[0], "the header must carry the destination version") + assert.Equal(t, byte(MsgTypeTransport), msg[1], "the message type must be preserved") + + gotID, gotPayload, err := UnmarshalTransportMsg(msg) + require.NoError(t, err) + assert.Equal(t, sender, *gotID, "the peer ID must be rewritten to the sender") + assert.Equal(t, payload, gotPayload, "the payload must be preserved") + + assert.ErrorIs(t, UpdateTransportMsg(msg[:headerTotalSizeTransport-1], ProtocolV1, sender), ErrInvalidMessageLength) +} diff --git a/shared/relay/messages/peer_state.go b/shared/relay/messages/peer_state.go index f10bc7bdf..42b1762b5 100644 --- a/shared/relay/messages/peer_state.go +++ b/shared/relay/messages/peer_state.go @@ -5,7 +5,7 @@ import ( ) func MarshalSubPeerStateMsg(ids []PeerID) ([][]byte, error) { - return marshalPeerIDs(ids, byte(MsgTypeSubscribePeerState)) + return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypeSubscribePeerState)) } func UnmarshalSubPeerStateMsg(buf []byte) ([]PeerID, error) { @@ -13,7 +13,7 @@ func UnmarshalSubPeerStateMsg(buf []byte) ([]PeerID, error) { } func MarshalUnsubPeerStateMsg(ids []PeerID) ([][]byte, error) { - return marshalPeerIDs(ids, byte(MsgTypeUnsubscribePeerState)) + return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypeUnsubscribePeerState)) } func UnmarshalUnsubPeerStateMsg(buf []byte) ([]PeerID, error) { @@ -21,7 +21,7 @@ func UnmarshalUnsubPeerStateMsg(buf []byte) ([]PeerID, error) { } func MarshalPeersOnline(ids []PeerID) ([][]byte, error) { - return marshalPeerIDs(ids, byte(MsgTypePeersOnline)) + return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypePeersOnline)) } func UnmarshalPeersOnlineMsg(buf []byte) ([]PeerID, error) { @@ -29,7 +29,7 @@ func UnmarshalPeersOnlineMsg(buf []byte) ([]PeerID, error) { } func MarshalPeersWentOffline(ids []PeerID) ([][]byte, error) { - return marshalPeerIDs(ids, byte(MsgTypePeersWentOffline)) + return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypePeersWentOffline)) } func UnMarshalPeersWentOffline(buf []byte) ([]PeerID, error) { @@ -37,7 +37,7 @@ func UnMarshalPeersWentOffline(buf []byte) ([]PeerID, error) { } // marshalPeerIDs is a generic function to marshal peer IDs with a specific message type -func marshalPeerIDs(ids []PeerID, msgType byte) ([][]byte, error) { +func marshalPeerIDs(version ProtocolVersion, ids []PeerID, msgType byte) ([][]byte, error) { if len(ids) == 0 { return nil, fmt.Errorf("no list of peer ids provided") } @@ -54,7 +54,7 @@ func marshalPeerIDs(ids []PeerID, msgType byte) ([][]byte, error) { totalSize := sizeOfProtoHeader + len(chunk)*peerIDSize buf := make([]byte, totalSize) - buf[0] = byte(CurrentProtocolVersion) + buf[0] = byte(version) buf[1] = msgType offset := sizeOfProtoHeader