sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-09-08 16:01:29 +02:00
Code Issues Packages Projects Releases Wiki Activity

API-key normalization and model-discovery snapshot validation

Browse Source
This commit is contained in:
Brandon Hopkins
2026-07-26 18:58:39 -07:00
parent 0463b30152
commit fea131101d
4 changed files with 61 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
proxy/model_discovery_sync_test.go
View File

@@ -94,6 +94,17 @@ func (s *modelDiscoverySyncStream) Context() context.Context {
return s.ctx
}
func completeModelDiscoveryInitialSync(t *testing.T, stream *modelDiscoverySyncStream) {
t.Helper()
stream.recv <- &proto.SyncMappingsResponse{InitialSyncComplete: true}
select {
case sent := <-stream.sent:
require.NotNil(t, sent.GetAck())
case <-time.After(time.Second):
t.Fatal("initial snapshot was not acknowledged")
}
}
func TestProxyCapabilitiesAdvertiseModelDiscovery(t *testing.T) {
t.Parallel()
@@ -178,10 +189,11 @@ func TestHandleSyncMappingsStreamRunsDiscoveryOutOfBand(t *testing.T) {
stream.sendWait = 10 * time.Millisecond
done := make(chan error, 1)
initialSyncDone := true
initialSyncDone := false
go func() {
done <- server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
}()
completeModelDiscoveryInitialSync(t, stream)
stream.recv <- &proto.SyncMappingsResponse{
ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{
@@ -242,10 +254,11 @@ func TestHandleSyncMappingsStreamBoundsConcurrentDiscovery(t *testing.T) {
stream := newModelDiscoverySyncStream(ctx)
done := make(chan error, 1)
initialSyncDone := true
initialSyncDone := false
go func() {
done <- server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
}()
completeModelDiscoveryInitialSync(t, stream)
for i := range 5 {
stream.recv <- &proto.SyncMappingsResponse{
@@ -298,7 +311,8 @@ func TestHandleSyncMappingsStreamRejectsMixedDiscoveryMessage(t *testing.T) {
}
stream := newModelDiscoverySyncStream(ctx)
stream.recv <- &proto.SyncMappingsResponse{
Mapping: []*proto.ProxyMapping{{Id: "mapping-1"}},
Mapping: []*proto.ProxyMapping{{Id: "mapping-1"}},
InitialSyncComplete: true,
ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{
RequestId: "request-1",
},
@@ -307,5 +321,29 @@ func TestHandleSyncMappingsStreamRejectsMixedDiscoveryMessage(t *testing.T) {
initialSyncDone := true
err := server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
require.EqualError(t, err, "model discovery message must not include mapping data")
require.EqualError(t, err, "model discovery message must not include mapping data or set initial_sync_complete")
}
func TestHandleSyncMappingsStreamRejectsDiscoveryBeforeInitialSnapshot(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
server := &Server{
Logger: log.New(),
routerReady: closedChan(),
modelDiscoverer: &stubModelDiscoverer{},
}
stream := newModelDiscoverySyncStream(ctx)
stream.recv <- &proto.SyncMappingsResponse{
ModelDiscoveryRequest: &proto.ModelDiscoveryRequest{
RequestId: "request-1",
},
}
close(stream.recv)
initialSyncDone := true
err := server.handleSyncMappingsStream(ctx, stream, &initialSyncDone, time.Time{})
require.EqualError(t, err, "model discovery request received before initial sync completed")
}

12
proxy/server.go
View File

@@ -1372,6 +1372,7 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox
}
tracker := s.newSnapshotTracker(initialSyncDone, connectTime)
initialSnapshotComplete := false
discoverer := s.modelDiscoverer
if discoverer == nil {
discoverer = modeldiscovery.New(s.Logger)
@@ -1411,7 +1412,10 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox
if discovery := msg.GetModelDiscoveryRequest(); discovery != nil {
if len(msg.GetMapping()) != 0 || msg.GetInitialSyncComplete() {
return errors.New("model discovery message must not include mapping data")
return errors.New("model discovery message must not include mapping data or set initial_sync_complete")
}
if !initialSnapshotComplete {
return errors.New("model discovery request received before initial sync completed")
}
select {
@@ -1455,7 +1459,11 @@ func (s *Server) handleSyncMappingsStream(ctx context.Context, stream proto.Prox
return err
}
s.Logger.Debug("Processing mapping update completed")
tracker.recordBatch(ctx, s, msg.GetMapping(), msg.GetInitialSyncComplete(), batchStart)
syncComplete := msg.GetInitialSyncComplete()
tracker.recordBatch(ctx, s, msg.GetMapping(), syncComplete, batchStart)
if syncComplete {
initialSnapshotComplete = true
}
if err := send(&proto.SyncMappingsRequest{
Msg: &proto.SyncMappingsRequest_Ack{Ack: &proto.SyncMappingsAck{}},