Files
go-libp2p/p2p/net/swarm/testing/testing.go

100 lines
2.5 KiB
Go

package testing
import (
"context"
"testing"
"github.com/libp2p/go-libp2p-core/metrics"
"github.com/libp2p/go-libp2p-core/network"
"github.com/libp2p/go-libp2p-core/peerstore"
"github.com/libp2p/go-libp2p-testing/net"
"github.com/libp2p/go-tcp-transport"
csms "github.com/libp2p/go-conn-security-multistream"
pstoremem "github.com/libp2p/go-libp2p-peerstore/pstoremem"
secio "github.com/libp2p/go-libp2p-secio"
tptu "github.com/libp2p/go-libp2p-transport-upgrader"
yamux "github.com/libp2p/go-libp2p-yamux"
msmux "github.com/libp2p/go-stream-muxer-multistream"
swarm "github.com/libp2p/go-libp2p-swarm"
)
type config struct {
disableReuseport bool
dialOnly bool
}
// Option is an option that can be passed when constructing a test swarm.
type Option func(*testing.T, *config)
// OptDisableReuseport disables reuseport in this test swarm.
var OptDisableReuseport Option = func(_ *testing.T, c *config) {
c.disableReuseport = true
}
// OptDialOnly prevents the test swarm from listening.
var OptDialOnly Option = func(_ *testing.T, c *config) {
c.dialOnly = true
}
// GenUpgrader creates a new connection upgrader for use with this swarm.
func GenUpgrader(n *swarm.Swarm) *tptu.Upgrader {
id := n.LocalPeer()
pk := n.Peerstore().PrivKey(id)
secMuxer := new(csms.SSMuxer)
secMuxer.AddTransport(secio.ID, &secio.Transport{
LocalID: id,
PrivateKey: pk,
})
stMuxer := msmux.NewBlankTransport()
stMuxer.AddTransport("/yamux/1.0.0", yamux.DefaultTransport)
return &tptu.Upgrader{
Secure: secMuxer,
Muxer: stMuxer,
Filters: n.Filters,
}
}
// GenSwarm generates a new test swarm.
func GenSwarm(t *testing.T, ctx context.Context, opts ...Option) *swarm.Swarm {
var cfg config
for _, o := range opts {
o(t, &cfg)
}
p := tnet.RandPeerNetParamsOrFatal(t)
ps := pstoremem.NewPeerstore()
ps.AddPubKey(p.ID, p.PubKey)
ps.AddPrivKey(p.ID, p.PrivKey)
s := swarm.NewSwarm(ctx, p.ID, ps, metrics.NewBandwidthCounter())
tcpTransport := tcp.NewTCPTransport(GenUpgrader(s))
tcpTransport.DisableReuseport = cfg.disableReuseport
if err := s.AddTransport(tcpTransport); err != nil {
t.Fatal(err)
}
if !cfg.dialOnly {
if err := s.Listen(p.Addr); err != nil {
t.Fatal(err)
}
s.Peerstore().AddAddrs(p.ID, s.ListenAddresses(), peerstore.PermanentAddrTTL)
}
return s
}
// DivulgeAddresses adds swarm a's addresses to swarm b's peerstore.
func DivulgeAddresses(a, b network.Network) {
id := a.LocalPeer()
addrs := a.Peerstore().Addrs(id)
b.Peerstore().AddAddrs(id, addrs, peerstore.PermanentAddrTTL)
}