peer/peerstates_test.go

@@ -1,143 +1,12 @@
 package peer
 
 import (
-	"net/netip"
 	"testing"
-	"time"
 	"vppn/m"
-
-	"git.crumpington.com/lib/go/ratelimiter"
 )
 
 // ----------------------------------------------------------------------------
 
-type PeerStateControlMsg struct {
-	Peer   remotePeer
-	Packet any
-}
-
-type PeerStateTestHarness struct {
-	State     peerState
-	Published remotePeer
-	Sent      []PeerStateControlMsg
-}
-
-func NewPeerStateTestHarness() *PeerStateTestHarness {
-	h := &PeerStateTestHarness{}
-
-	keys := generateKeys()
-
-	state := &peerData{
-		publish: func(rp remotePeer) {
-			h.Published = rp
-		},
-		sendControlPacket: func(rp remotePeer, pkt marshaller) {
-			h.Sent = append(h.Sent, PeerStateControlMsg{rp, pkt})
-		},
-		pingTimer: time.NewTicker(pingInterval),
-		localIP:   2,
-		remoteIP:  3,
-		privKey:   keys.PrivKey,
-		pubAddrs:  newPubAddrStore(netip.AddrPort{}),
-		limiter: ratelimiter.New(ratelimiter.Config{
-			FillPeriod:   20 * time.Millisecond,
-			MaxWaitCount: 1,
-		}),
-	}
-
-	h.State = enterStateDisconnected(state)
-	return h
-}
-
-func (h *PeerStateTestHarness) PeerUpdate(p *m.Peer) {
-	h.State = h.State.OnMsg(peerUpdateMsg{p})
-}
-
-func (h *PeerStateTestHarness) OnInit(msg controlMsg[packetInit]) {
-	h.State = h.State.OnMsg(msg)
-}
-
-func (h *PeerStateTestHarness) OnSyn(msg controlMsg[packetSyn]) {
-	h.State = h.State.OnMsg(msg)
-}
-
-func (h *PeerStateTestHarness) OnProbe(msg controlMsg[packetProbe]) {
-	h.State = h.State.OnMsg(msg)
-}
-
-func (h *PeerStateTestHarness) OnPingTimer() {
-	h.State = h.State.OnMsg(pingTimerMsg{})
-}
-
-func (h *PeerStateTestHarness) ConfigServer_Public(t *testing.T) *stateServer {
-	keys := generateKeys()
-
-	state := h.State.(*stateDisconnected)
-	state.localAddr = addrPort4(1, 1, 1, 2, 200)
-
-	peer := &m.Peer{
-		PeerIP:     3,
-		PublicIP:   []byte{1, 1, 1, 3},
-		Port:       456,
-		PubKey:     keys.PubKey,
-		PubSignKey: keys.PubSignKey,
-	}
-
-	h.PeerUpdate(peer)
-	assertEqual(t, h.Published.Up, false)
-	return assertType[*stateServer](t, h.State)
-}
-
-func (h *PeerStateTestHarness) ConfigServer_Relayed(t *testing.T) *stateServer {
-	keys := generateKeys()
-	peer := &m.Peer{
-		PeerIP:     3,
-		Port:       456,
-		PubKey:     keys.PubKey,
-		PubSignKey: keys.PubSignKey,
-	}
-
-	h.PeerUpdate(peer)
-	assertEqual(t, h.Published.Up, false)
-	return assertType[*stateServer](t, h.State)
-}
-
-func (h *PeerStateTestHarness) ConfigClientDirect(t *testing.T) *stateClientDirect {
-	keys := generateKeys()
-	peer := &m.Peer{
-		PeerIP:     3,
-		PublicIP:   []byte{1, 2, 3, 4},
-		Port:       456,
-		PubKey:     keys.PubKey,
-		PubSignKey: keys.PubSignKey,
-	}
-
-	h.PeerUpdate(peer)
-	assertEqual(t, h.Published.Up, false)
-
-	return assertType[*stateClientDirect](t, h.State)
-}
-
-func (h *PeerStateTestHarness) ConfigClientRelayed(t *testing.T) *stateClientRelayed {
-	keys := generateKeys()
-
-	state := h.State.(*stateDisconnected)
-	state.remoteIP = 1
-
-	peer := &m.Peer{
-		PeerIP:     3,
-		Port:       456,
-		PubKey:     keys.PubKey,
-		PubSignKey: keys.PubSignKey,
-	}
-
-	h.PeerUpdate(peer)
-	assertEqual(t, h.Published.Up, false)
-	return assertType[*stateClientRelayed](t, h.State)
-}
-
-// ----------------------------------------------------------------------------
-
 func TestPeerState_OnPeerUpdate_nilPeer(t *testing.T) {
 	h := NewPeerStateTestHarness()
 	h.PeerUpdate(nil)
@@ -163,6 +32,7 @@ func TestPeerState_OnPeerUpdate_publicLocalIsServer(t *testing.T) {
 	assertType[*stateServer](t, h.State)
 }
 
+/*
 func TestPeerState_OnPeerUpdate_serverDirect(t *testing.T) {
 	h := NewPeerStateTestHarness()
 	h.ConfigServer_Public(t)
@@ -178,11 +48,13 @@ func TestPeerState_OnPeerUpdate_clientDirect(t *testing.T) {
 	h.ConfigClientDirect(t)
 }
 
+/*
 func TestPeerState_OnPeerUpdate_clientRelayed(t *testing.T) {
 	h := NewPeerStateTestHarness()
 	h.ConfigClientRelayed(t)
 }
 
+/*
 func TestStateServer_directSyn(t *testing.T) {
 	h := NewPeerStateTestHarness()
 	h.ConfigServer_Relayed(t)
@@ -505,3 +377,4 @@ func TestStateClientRelayed_OnProbe_upgradeDirect(t *testing.T) {
 
 	assertType[*stateClientDirect](t, h.State)
 }
+*/

peer/state-clientinit_test.go

@@ -0,0 +1,83 @@
+package peer
+
+import (
+	"testing"
+	"time"
+)
+
+func TestPeerState_ClientInit_initWithIncorrectTraceID(t *testing.T) {
+	h := NewPeerStateTestHarness()
+	h.ConfigClientInit(t)
+
+	// Should have sent the first init packet.
+	assertEqual(t, len(h.Sent), 1)
+	init := assertType[packetInit](t, h.Sent[0].Packet)
+
+	init.TraceID = newTraceID()
+	h.OnInit(controlMsg[packetInit]{Packet: init})
+
+	assertType[*stateClientInit](t, h.State)
+}
+
+func TestPeerState_ClientInit_init(t *testing.T) {
+	h := NewPeerStateTestHarness()
+	h.ConfigClientInit(t)
+
+	// Should have sent the first init packet.
+	assertEqual(t, len(h.Sent), 1)
+	init := assertType[packetInit](t, h.Sent[0].Packet)
+	h.OnInit(controlMsg[packetInit]{Packet: init})
+
+	assertType[*stateClient](t, h.State)
+}
+
+func TestPeerState_ClientInit_onPing(t *testing.T) {
+	h := NewPeerStateTestHarness()
+	h.ConfigClientInit(t)
+
+	// Should have sent the first init packet.
+	assertEqual(t, len(h.Sent), 1)
+	h.Sent = h.Sent[:0]
+
+	for range 3 {
+		h.OnPingTimer()
+	}
+
+	assertEqual(t, len(h.Sent), 3)
+
+	for i := range h.Sent {
+		assertType[packetInit](t, h.Sent[i].Packet)
+	}
+}
+
+func TestPeerState_ClientInit_onPingTimeout(t *testing.T) {
+	h := NewPeerStateTestHarness()
+	h.ConfigClientInit(t)
+
+	state := assertType[*stateClientInit](t, h.State)
+	state.startedAt = time.Now().Add(-2 * timeoutInterval)
+
+	h.OnPingTimer()
+
+	// Should have moved into the client state due to timeout.
+	assertType[*stateClient](t, h.State)
+}
+
+func TestPeerState_ClientInit_onPeerUpdate(t *testing.T) {
+	h := NewPeerStateTestHarness()
+	h.ConfigClientInit(t)
+
+	h.PeerUpdate(nil)
+
+	// Should have moved into the client state due to timeout.
+	assertType[*stateDisconnected](t, h.State)
+}
+
+func TestPeerState_ClientInit_ignoreMessage(t *testing.T) {
+	h := NewPeerStateTestHarness()
+	h.ConfigClientInit(t)
+	h.OnProbe(controlMsg[packetProbe]{})
+
+	// Shouldn't do anything.
+	assertType[*stateClientInit](t, h.State)
+}

peer/state-disconnected.go

@@ -2,11 +2,11 @@ package peer
 
 import "net/netip"
 
-type stateDisconnected2 struct {
+type stateDisconnected struct {
 	*peerData
 }
 
-func enterStateDisconnected2(data *peerData) peerState {
+func enterStateDisconnected(data *peerData) peerState {
 	data.staged.Up = false
 	data.staged.Relay = false
 	data.staged.Direct = false
@@ -19,10 +19,10 @@ func enterStateDisconnected2(data *peerData) peerState {
 
 	data.pingTimer.Stop()
 
-	return &stateDisconnected2{data}
+	return &stateDisconnected{data}
 }
 
-func (s *stateDisconnected2) OnMsg(raw any) peerState {
+func (s *stateDisconnected) OnMsg(raw any) peerState {
 	switch msg := raw.(type) {
 	case peerUpdateMsg:
 		return initPeerState(s.peerData, msg.Peer)

peer/state-util_test.go

@@ -0,0 +1,146 @@
+package peer
+
+import (
+	"net/netip"
+	"testing"
+	"time"
+	"vppn/m"
+
+	"git.crumpington.com/lib/go/ratelimiter"
+)
+
+type PeerStateControlMsg struct {
+	Peer   remotePeer
+	Packet any
+}
+
+type PeerStateTestHarness struct {
+	State     peerState
+	Published remotePeer
+	Sent      []PeerStateControlMsg
+}
+
+func NewPeerStateTestHarness() *PeerStateTestHarness {
+	h := &PeerStateTestHarness{}
+
+	keys := generateKeys()
+
+	state := &peerData{
+		publish: func(rp remotePeer) {
+			h.Published = rp
+		},
+		sendControlPacket: func(rp remotePeer, pkt marshaller) {
+			h.Sent = append(h.Sent, PeerStateControlMsg{rp, pkt})
+		},
+		pingTimer: time.NewTicker(pingInterval),
+		localIP:   2,
+		remoteIP:  3,
+		privKey:   keys.PrivKey,
+		pubAddrs:  newPubAddrStore(netip.AddrPort{}),
+		limiter: ratelimiter.New(ratelimiter.Config{
+			FillPeriod:   20 * time.Millisecond,
+			MaxWaitCount: 1,
+		}),
+	}
+
+	h.State = enterStateDisconnected(state)
+	return h
+}
+
+func (h *PeerStateTestHarness) PeerUpdate(p *m.Peer) {
+	h.State = h.State.OnMsg(peerUpdateMsg{p})
+}
+
+func (h *PeerStateTestHarness) OnInit(msg controlMsg[packetInit]) {
+	h.State = h.State.OnMsg(msg)
+}
+
+func (h *PeerStateTestHarness) OnSyn(msg controlMsg[packetSyn]) {
+	h.State = h.State.OnMsg(msg)
+}
+
+func (h *PeerStateTestHarness) OnProbe(msg controlMsg[packetProbe]) {
+	h.State = h.State.OnMsg(msg)
+}
+
+func (h *PeerStateTestHarness) OnPingTimer() {
+	h.State = h.State.OnMsg(pingTimerMsg{})
+}
+
+func (h *PeerStateTestHarness) ConfigServer_Public(t *testing.T) *stateServer {
+	keys := generateKeys()
+
+	state := h.State.(*stateDisconnected)
+	state.localAddr = addrPort4(1, 1, 1, 2, 200)
+
+	peer := &m.Peer{
+		PeerIP:     3,
+		PublicIP:   []byte{1, 1, 1, 3},
+		Port:       456,
+		PubKey:     keys.PubKey,
+		PubSignKey: keys.PubSignKey,
+	}
+
+	h.PeerUpdate(peer)
+	assertEqual(t, h.Published.Up, false)
+	return assertType[*stateServer](t, h.State)
+}
+
+func (h *PeerStateTestHarness) ConfigServer_Relayed(t *testing.T) *stateServer {
+	keys := generateKeys()
+	peer := &m.Peer{
+		PeerIP:     3,
+		Port:       456,
+		PubKey:     keys.PubKey,
+		PubSignKey: keys.PubSignKey,
+	}
+
+	h.PeerUpdate(peer)
+	assertEqual(t, h.Published.Up, false)
+	return assertType[*stateServer](t, h.State)
+}
+
+func (h *PeerStateTestHarness) ConfigClientInit(t *testing.T) *stateClientInit {
+	keys := generateKeys()
+	peer := &m.Peer{
+		PeerIP:     3,
+		PublicIP:   []byte{1, 2, 3, 4},
+		Port:       456,
+		PubKey:     keys.PubKey,
+		PubSignKey: keys.PubSignKey,
+	}
+
+	h.PeerUpdate(peer)
+	assertEqual(t, h.Published.Up, false)
+	return assertType[*stateClientInit](t, h.State)
+}
+
+func (h *PeerStateTestHarness) ConfigClientDirect(t *testing.T) *stateClient {
+	h.ConfigClientInit(t)
+	init := assertType[packetInit](t, h.Sent[0].Packet)
+	h.OnInit(controlMsg[packetInit]{
+		Packet: init,
+	})
+
+	return assertType[*stateClient](t, h.State)
+}
+
+func (h *PeerStateTestHarness) ConfigClientRelayed(t *testing.T) *stateClient {
+	keys := generateKeys()
+
+	state := h.State.(*stateDisconnected)
+	state.remoteIP = 1
+
+	peer := &m.Peer{
+		PeerIP:     3,
+		Port:       456,
+		PubKey:     keys.PubKey,
+		PubSignKey: keys.PubSignKey,
+	}
+
+	// TODO: Fix me.
+
+	h.PeerUpdate(peer)
+	assertEqual(t, h.Published.Up, false)
+	return assertType[*stateClient](t, h.State)
+}

peer/statedata.go

@@ -92,7 +92,7 @@ func initPeerState(data *peerData, peer *m.Peer) peerState {
 	data.peer = peer
 
 	if peer == nil {
-		return enterStateDisconnected2(data)
+		return enterStateDisconnected(data)
 	}
 
 	if _, isValid := netip.AddrFromSlice(peer.PublicIP); isValid {