diff --git a/client/internal/peer/handshaker_pq_test.go b/client/internal/peer/handshaker_pq_test.go new file mode 100644 index 000000000..ec28954aa --- /dev/null +++ b/client/internal/peer/handshaker_pq_test.go @@ -0,0 +1,74 @@ +package peer + +import ( + "sync" + "testing" + + log "github.com/sirupsen/logrus" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.zx2c4.com/wireguard/wgctrl/wgtypes" + + relayClient "github.com/netbirdio/netbird/shared/relay/client" + signal "github.com/netbirdio/netbird/shared/signal/client" + sProto "github.com/netbirdio/netbird/shared/signal/proto" +) + +// bootstrapPQ is a fakePQ whose controller decision is set by the test. +type bootstrapPQ struct { + fakePQ + bootstrap bool +} + +func (b bootstrapPQ) ShouldSendBootstrapOffer(string) bool { return b.bootstrap } + +// newControllerHandshaker builds a handshaker for the controller side of a peer pair and +// records the type of every signalling message it sends. +func newControllerHandshaker(t *testing.T, pq PQHandshaker) (*Handshaker, func() []sProto.Body_Type) { + t.Helper() + var mu sync.Mutex + var sent []sProto.Body_Type + client := &signal.MockClient{ + ReadyFunc: func() bool { return true }, + SendFunc: func(msg *sProto.Message) error { + mu.Lock() + defer mu.Unlock() + sent = append(sent, msg.GetBody().GetType()) + return nil + }, + } + key, err := wgtypes.GeneratePrivateKey() + require.NoError(t, err) + config := ConnConfig{Key: "aaaa", LocalKey: "bbbb", PQ: pq} // LocalKey > Key: controller + relay := &WorkerRelay{relayManager: &relayClient.Manager{}} + h := NewHandshaker(log.WithField("test", t.Name()), config, NewSignaler(client, key), nil, relay, nil) + return h, func() []sProto.Body_Type { + mu.Lock() + defer mu.Unlock() + return append([]sProto.Body_Type(nil), sent...) + } +} + +// TestHandshaker_ControllerReoffersOnResponderOffer drives the responder-originated +// reconnect through the handshaker: when the KEM manager asks for a fresh bootstrap, the +// controller replies to the responder's KEM-less offer with its own offer and sends no +// answer, so the reconnect re-derives one PSK on both sides. +func TestHandshaker_ControllerReoffersOnResponderOffer(t *testing.T) { + h, sent := newControllerHandshaker(t, bootstrapPQ{bootstrap: true}) + + h.handleRemoteOffer(OfferAnswer{}) + + assert.Equal(t, []sProto.Body_Type{sProto.Body_OFFER}, sent(), + "the controller must reply with its own offer, not an answer") +} + +// TestHandshaker_ControllerAnswersWhileExchangeInFlight keeps the glare protection: with +// the controller's own exchange in flight, the responder's offer gets a normal answer. +func TestHandshaker_ControllerAnswersWhileExchangeInFlight(t *testing.T) { + h, sent := newControllerHandshaker(t, bootstrapPQ{bootstrap: false}) + + h.handleRemoteOffer(OfferAnswer{}) + + assert.Equal(t, []sProto.Body_Type{sProto.Body_ANSWER}, sent(), + "with an exchange in flight the controller must answer normally") +} diff --git a/client/internal/pqkem/capability_test.go b/client/internal/pqkem/capability_test.go index aabf2edd3..88ab7b2b9 100644 --- a/client/internal/pqkem/capability_test.go +++ b/client/internal/pqkem/capability_test.go @@ -45,13 +45,17 @@ func TestManager_MarkNonCapableCancelsInFlight(t *testing.T) { var cancelled atomic.Bool d.mu.Lock() ex := d.exchanges["aaaa"] - require.NotNil(t, ex, "precondition: the bootstrap exchange is in flight") - loopCancel := ex.cancel - ex.cancel = func() { - cancelled.Store(true) - loopCancel() + if ex != nil { + loopCancel := ex.cancel + ex.cancel = func() { + cancelled.Store(true) + loopCancel() + } } d.mu.Unlock() + // Asserted after unlocking: a failing require stops the test, and the deferred Stop + // would then block on the lock. + require.NotNil(t, ex, "precondition: the bootstrap exchange is in flight") // Now we learn the peer is non-PQ: the exchange must be dropped. d.MarkNonCapable("aaaa")