mirror of
https://github.com/netbirdio/netbird.git
synced 2026-10-21 12:59:08 +02:00
98 lines
4.0 KiB
Go
98 lines
4.0 KiB
Go
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")
|
|
}
|