vppn/peer/state-util_test.go
2025-02-26 12:45:43 +01:00

155 lines
3.4 KiB
Go

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) OnAck(msg controlMsg[packetAck]) {
h.State = h.State.OnMsg(msg)
}
func (h *PeerStateTestHarness) OnProbe(msg controlMsg[packetProbe]) {
h.State = h.State.OnMsg(msg)
}
func (h *PeerStateTestHarness) OnLocalDiscovery(msg controlMsg[packetLocalDiscovery]) {
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)
}