ceremonyclient/go-libp2p/p2p/transport/quicreuse/reuse.go
2023-08-20 23:07:43 -05:00

359 lines
9.1 KiB
Go

package quicreuse
import (
"context"
"crypto/tls"
"net"
"sync"
"time"
"github.com/google/gopacket/routing"
"github.com/libp2p/go-netroute"
"github.com/quic-go/quic-go"
)
type refCountedQuicTransport interface {
LocalAddr() net.Addr
// Used to send packets directly around QUIC. Useful for hole punching.
WriteTo([]byte, net.Addr) (int, error)
Close() error
// count transport reference
DecreaseCount()
IncreaseCount()
Dial(ctx context.Context, addr net.Addr, tlsConf *tls.Config, conf *quic.Config) (quic.Connection, error)
Listen(tlsConf *tls.Config, conf *quic.Config) (*quic.Listener, error)
}
type singleOwnerTransport struct {
quic.Transport
// Used to write packets directly around QUIC.
packetConn net.PacketConn
}
func (c *singleOwnerTransport) IncreaseCount() {}
func (c *singleOwnerTransport) DecreaseCount() {
c.Transport.Close()
}
func (c *singleOwnerTransport) LocalAddr() net.Addr {
return c.Transport.Conn.LocalAddr()
}
func (c *singleOwnerTransport) Close() error {
// TODO(when we drop support for go 1.19) use errors.Join
c.Transport.Close()
return c.packetConn.Close()
}
func (c *singleOwnerTransport) WriteTo(b []byte, addr net.Addr) (int, error) {
// Safe because we called quic.OptimizeConn ourselves.
return c.packetConn.WriteTo(b, addr)
}
// Constant. Defined as variables to simplify testing.
var (
garbageCollectInterval = 30 * time.Second
maxUnusedDuration = 10 * time.Second
)
type refcountedTransport struct {
quic.Transport
// Used to write packets directly around QUIC.
packetConn net.PacketConn
mutex sync.Mutex
refCount int
unusedSince time.Time
}
func (c *refcountedTransport) IncreaseCount() {
c.mutex.Lock()
c.refCount++
c.unusedSince = time.Time{}
c.mutex.Unlock()
}
func (c *refcountedTransport) Close() error {
// TODO(when we drop support for go 1.19) use errors.Join
c.Transport.Close()
return c.packetConn.Close()
}
func (c *refcountedTransport) WriteTo(b []byte, addr net.Addr) (int, error) {
// Safe because we called quic.OptimizeConn ourselves.
return c.packetConn.WriteTo(b, addr)
}
func (c *refcountedTransport) LocalAddr() net.Addr {
return c.Transport.Conn.LocalAddr()
}
func (c *refcountedTransport) DecreaseCount() {
c.mutex.Lock()
c.refCount--
if c.refCount == 0 {
c.unusedSince = time.Now()
}
c.mutex.Unlock()
}
func (c *refcountedTransport) ShouldGarbageCollect(now time.Time) bool {
c.mutex.Lock()
defer c.mutex.Unlock()
return !c.unusedSince.IsZero() && c.unusedSince.Add(maxUnusedDuration).Before(now)
}
type reuse struct {
mutex sync.Mutex
closeChan chan struct{}
gcStopChan chan struct{}
routes routing.Router
unicast map[string] /* IP.String() */ map[int] /* port */ *refcountedTransport
// globalListeners contains transports that are listening on 0.0.0.0 / ::
globalListeners map[int]*refcountedTransport
// globalDialers contains transports that we've dialed out from. These transports are listening on 0.0.0.0 / ::
// On Dial, transports are reused from this map if no transport is available in the globalListeners
// On Listen, transports are reused from this map if the requested port is 0, and then moved to globalListeners
globalDialers map[int]*refcountedTransport
statelessResetKey *quic.StatelessResetKey
metricsTracer *metricsTracer
}
func newReuse(srk *quic.StatelessResetKey, mt *metricsTracer) *reuse {
r := &reuse{
unicast: make(map[string]map[int]*refcountedTransport),
globalListeners: make(map[int]*refcountedTransport),
globalDialers: make(map[int]*refcountedTransport),
closeChan: make(chan struct{}),
gcStopChan: make(chan struct{}),
statelessResetKey: srk,
metricsTracer: mt,
}
go r.gc()
return r
}
func (r *reuse) gc() {
defer func() {
r.mutex.Lock()
for _, tr := range r.globalListeners {
tr.Close()
}
for _, tr := range r.globalDialers {
tr.Close()
}
for _, trs := range r.unicast {
for _, tr := range trs {
tr.Close()
}
}
r.mutex.Unlock()
close(r.gcStopChan)
}()
ticker := time.NewTicker(garbageCollectInterval)
defer ticker.Stop()
for {
select {
case <-r.closeChan:
return
case <-ticker.C:
now := time.Now()
r.mutex.Lock()
for key, tr := range r.globalListeners {
if tr.ShouldGarbageCollect(now) {
tr.Close()
delete(r.globalListeners, key)
}
}
for key, tr := range r.globalDialers {
if tr.ShouldGarbageCollect(now) {
tr.Close()
delete(r.globalDialers, key)
}
}
for ukey, trs := range r.unicast {
for key, tr := range trs {
if tr.ShouldGarbageCollect(now) {
tr.Close()
delete(trs, key)
}
}
if len(trs) == 0 {
delete(r.unicast, ukey)
// If we've dropped all transports with a unicast binding,
// assume our routes may have changed.
if len(r.unicast) == 0 {
r.routes = nil
} else {
// Ignore the error, there's nothing we can do about
// it.
r.routes, _ = netroute.New()
}
}
}
r.mutex.Unlock()
}
}
}
func (r *reuse) TransportForDial(network string, raddr *net.UDPAddr) (*refcountedTransport, error) {
var ip *net.IP
// Only bother looking up the source address if we actually _have_ non 0.0.0.0 listeners.
// Otherwise, save some time.
r.mutex.Lock()
router := r.routes
r.mutex.Unlock()
if router != nil {
_, _, src, err := router.Route(raddr.IP)
if err == nil && !src.IsUnspecified() {
ip = &src
}
}
r.mutex.Lock()
defer r.mutex.Unlock()
tr, err := r.transportForDialLocked(network, ip)
if err != nil {
return nil, err
}
tr.IncreaseCount()
return tr, nil
}
func (r *reuse) transportForDialLocked(network string, source *net.IP) (*refcountedTransport, error) {
if source != nil {
// We already have at least one suitable transport...
if trs, ok := r.unicast[source.String()]; ok {
// ... we don't care which port we're dialing from. Just use the first.
for _, tr := range trs {
return tr, nil
}
}
}
// Use a transport listening on 0.0.0.0 (or ::).
// Again, we don't care about the port number.
for _, tr := range r.globalListeners {
return tr, nil
}
// Use a transport we've previously dialed from
for _, tr := range r.globalDialers {
return tr, nil
}
// We don't have a transport that we can use for dialing.
// Dial a new connection from a random port.
var addr *net.UDPAddr
switch network {
case "udp4":
addr = &net.UDPAddr{IP: net.IPv4zero, Port: 0}
case "udp6":
addr = &net.UDPAddr{IP: net.IPv6zero, Port: 0}
}
conn, err := listenAndOptimize(network, addr)
if err != nil {
return nil, err
}
tr := &refcountedTransport{Transport: quic.Transport{
Conn: conn,
StatelessResetKey: r.statelessResetKey,
}, packetConn: conn}
if r.metricsTracer != nil {
tr.Transport.Tracer = r.metricsTracer
}
r.globalDialers[conn.LocalAddr().(*net.UDPAddr).Port] = tr
return tr, nil
}
func (r *reuse) TransportForListen(network string, laddr *net.UDPAddr) (*refcountedTransport, error) {
r.mutex.Lock()
defer r.mutex.Unlock()
// Check if we can reuse a transport we have already dialed out from.
// We reuse a transport from globalDialers when the requested port is 0 or the requested
// port is already in the globalDialers.
// If we are reusing a transport from globalDialers, we move the globalDialers entry to
// globalListeners
if laddr.IP.IsUnspecified() {
var rTr *refcountedTransport
var localAddr *net.UDPAddr
if laddr.Port == 0 {
// the requested port is 0, we can reuse any transport
for _, tr := range r.globalDialers {
rTr = tr
localAddr = rTr.LocalAddr().(*net.UDPAddr)
delete(r.globalDialers, localAddr.Port)
break
}
} else if _, ok := r.globalDialers[laddr.Port]; ok {
rTr = r.globalDialers[laddr.Port]
localAddr = rTr.LocalAddr().(*net.UDPAddr)
delete(r.globalDialers, localAddr.Port)
}
// found a match
if rTr != nil {
rTr.IncreaseCount()
r.globalListeners[localAddr.Port] = rTr
return rTr, nil
}
}
conn, err := listenAndOptimize(network, laddr)
if err != nil {
return nil, err
}
localAddr := conn.LocalAddr().(*net.UDPAddr)
tr := &refcountedTransport{Transport: quic.Transport{
Conn: conn,
StatelessResetKey: r.statelessResetKey,
}, packetConn: conn}
if r.metricsTracer != nil {
tr.Transport.Tracer = r.metricsTracer
}
tr.IncreaseCount()
// Deal with listen on a global address
if localAddr.IP.IsUnspecified() {
// The kernel already checked that the laddr is not already listen
// so we need not check here (when we create ListenUDP).
r.globalListeners[localAddr.Port] = tr
return tr, nil
}
// Deal with listen on a unicast address
if _, ok := r.unicast[localAddr.IP.String()]; !ok {
r.unicast[localAddr.IP.String()] = make(map[int]*refcountedTransport)
// Assume the system's routes may have changed if we're adding a new listener.
// Ignore the error, there's nothing we can do.
r.routes, _ = netroute.New()
}
// The kernel already checked that the laddr is not already listen
// so we need not check here (when we create ListenUDP).
r.unicast[localAddr.IP.String()][localAddr.Port] = tr
return tr, nil
}
func (r *reuse) Close() error {
close(r.closeChan)
<-r.gcStopChan
return nil
}