81 lines
1.6 KiB
Go
81 lines
1.6 KiB
Go
package peer
|
|
|
|
import (
|
|
"bytes"
|
|
"net"
|
|
"net/netip"
|
|
"sync"
|
|
)
|
|
|
|
type TestPacket struct {
|
|
Addr netip.AddrPort
|
|
Data []byte
|
|
}
|
|
|
|
type TestNetwork struct {
|
|
lock sync.Mutex
|
|
packets map[netip.AddrPort]chan TestPacket
|
|
}
|
|
|
|
func NewTestNetwork() *TestNetwork {
|
|
return &TestNetwork{packets: map[netip.AddrPort]chan TestPacket{}}
|
|
}
|
|
|
|
func (n *TestNetwork) NewUDPConn(localAddr netip.AddrPort) *TestUDPConn {
|
|
n.lock.Lock()
|
|
defer n.lock.Unlock()
|
|
if _, ok := n.packets[localAddr]; !ok {
|
|
n.packets[localAddr] = make(chan TestPacket, 1024)
|
|
}
|
|
return &TestUDPConn{
|
|
addr: localAddr,
|
|
n: n,
|
|
packets: n.packets[localAddr],
|
|
}
|
|
}
|
|
|
|
func (n *TestNetwork) write(b []byte, from, to netip.AddrPort) {
|
|
n.lock.Lock()
|
|
defer n.lock.Unlock()
|
|
if _, ok := n.packets[to]; !ok {
|
|
n.packets[to] = make(chan TestPacket, 1024)
|
|
}
|
|
n.packets[to] <- TestPacket{
|
|
Addr: from,
|
|
Data: bytes.Clone(b),
|
|
}
|
|
}
|
|
|
|
type TestUDPConn struct {
|
|
addr netip.AddrPort
|
|
n *TestNetwork
|
|
packets chan TestPacket
|
|
}
|
|
|
|
func (c *TestUDPConn) WriteToUDPAddrPort(b []byte, addr netip.AddrPort) (int, error) {
|
|
c.n.write(b, c.addr, addr)
|
|
return len(b), nil
|
|
}
|
|
|
|
func (c *TestUDPConn) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
|
|
return c.WriteToUDPAddrPort(b, addr.AddrPort())
|
|
}
|
|
|
|
func (c *TestUDPConn) ReadFromUDPAddrPort(b []byte) (n int, addr netip.AddrPort, err error) {
|
|
pkt := <-c.packets
|
|
b = b[:len(pkt.Data)]
|
|
copy(b, pkt.Data)
|
|
return len(b), pkt.Addr, nil
|
|
}
|
|
|
|
func (c *TestUDPConn) Packets() (out []TestPacket) {
|
|
for {
|
|
select {
|
|
case pkt := <-c.packets:
|
|
out = append(out, pkt)
|
|
default:
|
|
return
|
|
}
|
|
}
|
|
}
|