mirror of
https://github.com/netbirdio/netbird.git
synced 2026-05-19 15:19:55 +00:00
511 lines
14 KiB
Go
511 lines
14 KiB
Go
package proxy
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
grpcstatus "google.golang.org/grpc/status"
|
|
|
|
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
|
"github.com/netbirdio/netbird/shared/management/proto"
|
|
)
|
|
|
|
func TestIntegration_SyncMappings_HappyPath(t *testing.T) {
|
|
setup := setupIntegrationTest(t)
|
|
defer setup.cleanup()
|
|
|
|
conn, err := grpc.NewClient(setup.grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
require.NoError(t, err)
|
|
defer conn.Close()
|
|
|
|
client := proto.NewProxyServiceClient(conn)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
stream, err := client.SyncMappings(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Send init.
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Init{
|
|
Init: &proto.SyncMappingsInit{
|
|
ProxyId: "sync-proxy-1",
|
|
Version: "test-v1",
|
|
Address: "test.proxy.io",
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
mappingsByID := make(map[string]*proto.ProxyMapping)
|
|
for {
|
|
msg, err := stream.Recv()
|
|
require.NoError(t, err)
|
|
for _, m := range msg.GetMapping() {
|
|
mappingsByID[m.GetId()] = m
|
|
}
|
|
|
|
// Ack every batch.
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
if msg.GetInitialSyncComplete() {
|
|
break
|
|
}
|
|
}
|
|
|
|
assert.Len(t, mappingsByID, 2, "Should receive 2 mappings")
|
|
|
|
rp1 := mappingsByID["rp-1"]
|
|
require.NotNil(t, rp1)
|
|
assert.Equal(t, "app1.test.proxy.io", rp1.GetDomain())
|
|
assert.Equal(t, "test-account-1", rp1.GetAccountId())
|
|
assert.Equal(t, proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED, rp1.GetType())
|
|
assert.NotEmpty(t, rp1.GetAuthToken(), "Should have auth token")
|
|
|
|
rp2 := mappingsByID["rp-2"]
|
|
require.NotNil(t, rp2)
|
|
assert.Equal(t, "app2.test.proxy.io", rp2.GetDomain())
|
|
}
|
|
|
|
func TestIntegration_SyncMappings_BackPressure(t *testing.T) {
|
|
setup := setupIntegrationTest(t)
|
|
defer setup.cleanup()
|
|
|
|
// Add more services so we get multiple batches.
|
|
addServicesToStore(t, setup, 20, "test.proxy.io")
|
|
|
|
conn, err := grpc.NewClient(setup.grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
require.NoError(t, err)
|
|
defer conn.Close()
|
|
|
|
client := proto.NewProxyServiceClient(conn)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
stream, err := client.SyncMappings(ctx)
|
|
require.NoError(t, err)
|
|
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Init{
|
|
Init: &proto.SyncMappingsInit{
|
|
ProxyId: "sync-proxy-backpressure",
|
|
Version: "test-v1",
|
|
Address: "test.proxy.io",
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Record the ordering of events to verify back-pressure.
|
|
var mu sync.Mutex
|
|
var events []string
|
|
var totalMappings int
|
|
|
|
for {
|
|
msg, err := stream.Recv()
|
|
require.NoError(t, err)
|
|
|
|
mu.Lock()
|
|
events = append(events, "recv")
|
|
totalMappings += len(msg.GetMapping())
|
|
mu.Unlock()
|
|
|
|
// Simulate processing delay.
|
|
time.Sleep(50 * time.Millisecond)
|
|
|
|
mu.Lock()
|
|
events = append(events, "ack")
|
|
mu.Unlock()
|
|
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
if msg.GetInitialSyncComplete() {
|
|
break
|
|
}
|
|
}
|
|
|
|
// 2 original + 20 added = 22 services total.
|
|
assert.Equal(t, 22, totalMappings, "should receive all 22 mappings")
|
|
|
|
// Events should alternate recv/ack — no two recvs in a row
|
|
// (management waits for ack before sending next).
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
for i := 0; i < len(events)-1; i += 2 {
|
|
assert.Equal(t, "recv", events[i], "event %d should be recv", i)
|
|
if i+1 < len(events) {
|
|
assert.Equal(t, "ack", events[i+1], "event %d should be ack", i+1)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestIntegration_SyncMappings_IncrementalUpdate(t *testing.T) {
|
|
setup := setupIntegrationTest(t)
|
|
defer setup.cleanup()
|
|
|
|
conn, err := grpc.NewClient(setup.grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
require.NoError(t, err)
|
|
defer conn.Close()
|
|
|
|
client := proto.NewProxyServiceClient(conn)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
stream, err := client.SyncMappings(ctx)
|
|
require.NoError(t, err)
|
|
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Init{
|
|
Init: &proto.SyncMappingsInit{
|
|
ProxyId: "sync-proxy-incremental",
|
|
Version: "test-v1",
|
|
Address: "test.proxy.io",
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Drain initial snapshot.
|
|
for {
|
|
msg, err := stream.Recv()
|
|
require.NoError(t, err)
|
|
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
if msg.GetInitialSyncComplete() {
|
|
break
|
|
}
|
|
}
|
|
|
|
// Now send an incremental update via the management server.
|
|
setup.proxyService.SendServiceUpdate(&proto.GetMappingUpdateResponse{
|
|
Mapping: []*proto.ProxyMapping{{
|
|
Type: proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED,
|
|
Id: "rp-1",
|
|
AccountId: "test-account-1",
|
|
Domain: "app1.test.proxy.io",
|
|
}},
|
|
})
|
|
|
|
// Receive the incremental update on the sync stream.
|
|
msg, err := stream.Recv()
|
|
require.NoError(t, err)
|
|
require.NotEmpty(t, msg.GetMapping())
|
|
assert.Equal(t, "rp-1", msg.GetMapping()[0].GetId())
|
|
assert.Equal(t, proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED, msg.GetMapping()[0].GetType())
|
|
}
|
|
|
|
func TestIntegration_SyncMappings_MixedProxyVersions(t *testing.T) {
|
|
setup := setupIntegrationTest(t)
|
|
defer setup.cleanup()
|
|
|
|
conn, err := grpc.NewClient(setup.grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
require.NoError(t, err)
|
|
defer conn.Close()
|
|
|
|
client := proto.NewProxyServiceClient(conn)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
// Old proxy uses GetMappingUpdate.
|
|
legacyStream, err := client.GetMappingUpdate(ctx, &proto.GetMappingUpdateRequest{
|
|
ProxyId: "legacy-proxy",
|
|
Version: "old-v1",
|
|
Address: "test.proxy.io",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
var legacyMappings []*proto.ProxyMapping
|
|
for {
|
|
msg, err := legacyStream.Recv()
|
|
require.NoError(t, err)
|
|
legacyMappings = append(legacyMappings, msg.GetMapping()...)
|
|
if msg.GetInitialSyncComplete() {
|
|
break
|
|
}
|
|
}
|
|
|
|
// New proxy uses SyncMappings.
|
|
syncStream, err := client.SyncMappings(ctx)
|
|
require.NoError(t, err)
|
|
|
|
err = syncStream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Init{
|
|
Init: &proto.SyncMappingsInit{
|
|
ProxyId: "new-proxy",
|
|
Version: "new-v2",
|
|
Address: "test.proxy.io",
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
var syncMappings []*proto.ProxyMapping
|
|
for {
|
|
msg, err := syncStream.Recv()
|
|
require.NoError(t, err)
|
|
syncMappings = append(syncMappings, msg.GetMapping()...)
|
|
|
|
err = syncStream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
if msg.GetInitialSyncComplete() {
|
|
break
|
|
}
|
|
}
|
|
|
|
// Both should receive the same set of mappings.
|
|
assert.Equal(t, len(legacyMappings), len(syncMappings),
|
|
"legacy and sync proxies should receive the same number of mappings")
|
|
|
|
legacyIDs := make(map[string]bool)
|
|
for _, m := range legacyMappings {
|
|
legacyIDs[m.GetId()] = true
|
|
}
|
|
for _, m := range syncMappings {
|
|
assert.True(t, legacyIDs[m.GetId()],
|
|
"mapping %s should be present in both streams", m.GetId())
|
|
}
|
|
|
|
// Both proxies should be connected.
|
|
proxies := setup.proxyService.GetConnectedProxies()
|
|
assert.Contains(t, proxies, "legacy-proxy")
|
|
assert.Contains(t, proxies, "new-proxy")
|
|
|
|
// Both should receive incremental updates.
|
|
setup.proxyService.SendServiceUpdate(&proto.GetMappingUpdateResponse{
|
|
Mapping: []*proto.ProxyMapping{{
|
|
Type: proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED,
|
|
Id: "rp-1",
|
|
AccountId: "test-account-1",
|
|
Domain: "app1.test.proxy.io",
|
|
}},
|
|
})
|
|
|
|
// Legacy proxy receives via GetMappingUpdateResponse.
|
|
legacyMsg, err := legacyStream.Recv()
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "rp-1", legacyMsg.GetMapping()[0].GetId())
|
|
|
|
// Sync proxy receives via SyncMappingsResponse.
|
|
syncMsg, err := syncStream.Recv()
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "rp-1", syncMsg.GetMapping()[0].GetId())
|
|
}
|
|
|
|
func TestIntegration_SyncMappings_Reconnect(t *testing.T) {
|
|
setup := setupIntegrationTest(t)
|
|
defer setup.cleanup()
|
|
|
|
conn, err := grpc.NewClient(setup.grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
require.NoError(t, err)
|
|
defer conn.Close()
|
|
|
|
client := proto.NewProxyServiceClient(conn)
|
|
proxyID := "sync-proxy-reconnect"
|
|
|
|
receiveMappings := func() []*proto.ProxyMapping {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
stream, err := client.SyncMappings(ctx)
|
|
require.NoError(t, err)
|
|
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Init{
|
|
Init: &proto.SyncMappingsInit{
|
|
ProxyId: proxyID,
|
|
Version: "test-v1",
|
|
Address: "test.proxy.io",
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
var mappings []*proto.ProxyMapping
|
|
for {
|
|
msg, err := stream.Recv()
|
|
require.NoError(t, err)
|
|
mappings = append(mappings, msg.GetMapping()...)
|
|
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
if msg.GetInitialSyncComplete() {
|
|
break
|
|
}
|
|
}
|
|
return mappings
|
|
}
|
|
|
|
first := receiveMappings()
|
|
time.Sleep(100 * time.Millisecond)
|
|
second := receiveMappings()
|
|
|
|
assert.Equal(t, len(first), len(second),
|
|
"should receive same mappings on reconnect")
|
|
|
|
firstIDs := make(map[string]bool)
|
|
for _, m := range first {
|
|
firstIDs[m.GetId()] = true
|
|
}
|
|
for _, m := range second {
|
|
assert.True(t, firstIDs[m.GetId()],
|
|
"mapping %s should be present in both connections", m.GetId())
|
|
}
|
|
}
|
|
|
|
// --- Fallback tests: old management returns Unimplemented ---
|
|
|
|
// unimplementedProxyServer embeds UnimplementedProxyServiceServer so
|
|
// SyncMappings returns codes.Unimplemented while GetMappingUpdate works.
|
|
type unimplementedSyncServer struct {
|
|
proto.UnimplementedProxyServiceServer
|
|
getMappingCalls atomic.Int32
|
|
}
|
|
|
|
func (s *unimplementedSyncServer) GetMappingUpdate(_ *proto.GetMappingUpdateRequest, stream proto.ProxyService_GetMappingUpdateServer) error {
|
|
s.getMappingCalls.Add(1)
|
|
return stream.Send(&proto.GetMappingUpdateResponse{
|
|
Mapping: []*proto.ProxyMapping{{Id: "svc-1", AccountId: "acct-1", Domain: "example.com"}},
|
|
InitialSyncComplete: true,
|
|
})
|
|
}
|
|
|
|
func TestIntegration_FallbackToGetMappingUpdate(t *testing.T) {
|
|
// Start a gRPC server that does NOT implement SyncMappings.
|
|
lis, err := net.Listen("tcp", "127.0.0.1:0")
|
|
require.NoError(t, err)
|
|
|
|
srv := &unimplementedSyncServer{}
|
|
grpcServer := grpc.NewServer()
|
|
proto.RegisterProxyServiceServer(grpcServer, srv)
|
|
go func() { _ = grpcServer.Serve(lis) }()
|
|
defer grpcServer.GracefulStop()
|
|
|
|
conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
require.NoError(t, err)
|
|
defer conn.Close()
|
|
|
|
client := proto.NewProxyServiceClient(conn)
|
|
|
|
// Try SyncMappings — should get Unimplemented.
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
|
|
stream, err := client.SyncMappings(ctx)
|
|
require.NoError(t, err)
|
|
|
|
err = stream.Send(&proto.SyncMappingsRequest{
|
|
Msg: &proto.SyncMappingsRequest_Init{
|
|
Init: &proto.SyncMappingsInit{
|
|
ProxyId: "test-proxy",
|
|
Address: "test.example.com",
|
|
},
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
_, err = stream.Recv()
|
|
require.Error(t, err)
|
|
st, ok := grpcstatus.FromError(err)
|
|
require.True(t, ok)
|
|
assert.Equal(t, codes.Unimplemented, st.Code(),
|
|
"unimplemented SyncMappings should return Unimplemented code")
|
|
|
|
// isSyncUnimplemented should detect this.
|
|
assert.True(t, isSyncUnimplemented(err))
|
|
|
|
// The actual fallback: GetMappingUpdate should work.
|
|
legacyStream, err := client.GetMappingUpdate(ctx, &proto.GetMappingUpdateRequest{
|
|
ProxyId: "test-proxy",
|
|
Address: "test.example.com",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
msg, err := legacyStream.Recv()
|
|
require.NoError(t, err)
|
|
assert.True(t, msg.GetInitialSyncComplete())
|
|
assert.Len(t, msg.GetMapping(), 1)
|
|
assert.Equal(t, int32(1), srv.getMappingCalls.Load())
|
|
}
|
|
|
|
func TestIsSyncUnimplemented(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
err error
|
|
want bool
|
|
}{
|
|
{"nil error", nil, false},
|
|
{"non-grpc error", errors.New("random"), false},
|
|
{"grpc internal", grpcstatus.Error(codes.Internal, "fail"), false},
|
|
{"grpc unavailable", grpcstatus.Error(codes.Unavailable, "fail"), false},
|
|
{"grpc unimplemented", grpcstatus.Error(codes.Unimplemented, "method not found"), true},
|
|
{
|
|
"wrapped unimplemented",
|
|
fmt.Errorf("create sync stream: %w", grpcstatus.Error(codes.Unimplemented, "nope")),
|
|
// grpc/status.FromError unwraps in recent versions of grpc-go.
|
|
true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
assert.Equal(t, tt.want, isSyncUnimplemented(tt.err))
|
|
})
|
|
}
|
|
}
|
|
|
|
// addServicesToStore adds n extra services to the test store for the given cluster.
|
|
func addServicesToStore(t *testing.T, setup *integrationTestSetup, n int, cluster string) {
|
|
t.Helper()
|
|
ctx := context.Background()
|
|
for i := 0; i < n; i++ {
|
|
svc := &service.Service{
|
|
ID: fmt.Sprintf("extra-svc-%d", i),
|
|
AccountID: "test-account-1",
|
|
Name: fmt.Sprintf("Extra Service %d", i),
|
|
Domain: fmt.Sprintf("extra-%d.test.proxy.io", i),
|
|
ProxyCluster: cluster,
|
|
Enabled: true,
|
|
Targets: []*service.Target{{
|
|
Path: strPtr("/"),
|
|
Host: fmt.Sprintf("10.0.1.%d", i%256),
|
|
Port: 8080,
|
|
Protocol: "http",
|
|
TargetId: fmt.Sprintf("peer-extra-%d", i),
|
|
TargetType: "peer",
|
|
Enabled: true,
|
|
}},
|
|
}
|
|
require.NoError(t, setup.store.CreateService(ctx, svc))
|
|
}
|
|
}
|