ceremonyclient/go-libp2p/p2p/transport/quicreuse/connmgr.go

229 lines
5.4 KiB
Go
Raw Normal View History

2023-08-21 03:50:38 +00:00
package quicreuse
import (
"context"
"crypto/tls"
"errors"
"net"
"sync"
ma "github.com/multiformats/go-multiaddr"
manet "github.com/multiformats/go-multiaddr/net"
"github.com/quic-go/quic-go"
quiclogging "github.com/quic-go/quic-go/logging"
)
type ConnManager struct {
reuseUDP4 *reuse
reuseUDP6 *reuse
enableReuseport bool
enableMetrics bool
serverConfig *quic.Config
clientConfig *quic.Config
quicListenersMu sync.Mutex
quicListeners map[string]quicListenerEntry
2024-06-08 11:32:45 +00:00
srk quic.StatelessResetKey
tokenKey quic.TokenGeneratorKey
2023-08-21 03:50:38 +00:00
}
type quicListenerEntry struct {
refCount int
ln *quicListener
}
2024-06-08 11:32:45 +00:00
func NewConnManager(statelessResetKey quic.StatelessResetKey, tokenKey quic.TokenGeneratorKey, opts ...Option) (*ConnManager, error) {
2023-08-21 03:50:38 +00:00
cm := &ConnManager{
enableReuseport: true,
quicListeners: make(map[string]quicListenerEntry),
srk: statelessResetKey,
2024-06-08 11:32:45 +00:00
tokenKey: tokenKey,
2023-08-21 03:50:38 +00:00
}
for _, o := range opts {
if err := o(cm); err != nil {
return nil, err
}
}
quicConf := quicConfig.Clone()
2024-06-08 11:32:45 +00:00
quicConf.Tracer = func(ctx context.Context, p quiclogging.Perspective, ci quic.ConnectionID) *quiclogging.ConnectionTracer {
var tracer *quiclogging.ConnectionTracer
2023-08-21 03:50:38 +00:00
if qlogTracerDir != "" {
2024-06-08 11:32:45 +00:00
tracer = qloggerForDir(qlogTracerDir, p, ci)
2023-08-21 03:50:38 +00:00
}
2024-06-08 11:32:45 +00:00
return tracer
2023-08-21 03:50:38 +00:00
}
serverConfig := quicConf.Clone()
cm.clientConfig = quicConf
cm.serverConfig = serverConfig
if cm.enableReuseport {
2024-06-08 11:32:45 +00:00
cm.reuseUDP4 = newReuse(&statelessResetKey, &tokenKey)
cm.reuseUDP6 = newReuse(&statelessResetKey, &tokenKey)
2023-08-21 03:50:38 +00:00
}
return cm, nil
}
func (c *ConnManager) getReuse(network string) (*reuse, error) {
switch network {
case "udp4":
return c.reuseUDP4, nil
case "udp6":
return c.reuseUDP6, nil
default:
return nil, errors.New("invalid network: must be either udp4 or udp6")
}
}
func (c *ConnManager) ListenQUIC(addr ma.Multiaddr, tlsConf *tls.Config, allowWindowIncrease func(conn quic.Connection, delta uint64) bool) (Listener, error) {
netw, host, err := manet.DialArgs(addr)
if err != nil {
return nil, err
}
laddr, err := net.ResolveUDPAddr(netw, host)
if err != nil {
return nil, err
}
c.quicListenersMu.Lock()
defer c.quicListenersMu.Unlock()
key := laddr.String()
entry, ok := c.quicListeners[key]
if !ok {
tr, err := c.transportForListen(netw, laddr)
if err != nil {
return nil, err
}
2024-06-08 11:32:45 +00:00
ln, err := newQuicListener(tr, c.serverConfig)
2023-08-21 03:50:38 +00:00
if err != nil {
return nil, err
}
key = tr.LocalAddr().String()
entry = quicListenerEntry{ln: ln}
}
l, err := entry.ln.Add(tlsConf, allowWindowIncrease, func() { c.onListenerClosed(key) })
if err != nil {
if entry.refCount <= 0 {
entry.ln.Close()
}
return nil, err
}
entry.refCount++
c.quicListeners[key] = entry
return l, nil
}
func (c *ConnManager) onListenerClosed(key string) {
c.quicListenersMu.Lock()
defer c.quicListenersMu.Unlock()
entry := c.quicListeners[key]
entry.refCount = entry.refCount - 1
if entry.refCount <= 0 {
delete(c.quicListeners, key)
entry.ln.Close()
} else {
c.quicListeners[key] = entry
}
}
func (c *ConnManager) transportForListen(network string, laddr *net.UDPAddr) (refCountedQuicTransport, error) {
if c.enableReuseport {
reuse, err := c.getReuse(network)
if err != nil {
return nil, err
}
return reuse.TransportForListen(network, laddr)
}
2024-06-08 11:32:45 +00:00
conn, err := net.ListenUDP(network, laddr)
2023-08-21 03:50:38 +00:00
if err != nil {
return nil, err
}
2024-06-08 11:32:45 +00:00
return &singleOwnerTransport{
packetConn: conn,
Transport: quic.Transport{
Conn: conn,
StatelessResetKey: &c.srk,
TokenGeneratorKey: &c.tokenKey,
},
}, nil
2023-08-21 03:50:38 +00:00
}
func (c *ConnManager) DialQUIC(ctx context.Context, raddr ma.Multiaddr, tlsConf *tls.Config, allowWindowIncrease func(conn quic.Connection, delta uint64) bool) (quic.Connection, error) {
naddr, v, err := FromQuicMultiaddr(raddr)
if err != nil {
return nil, err
}
netw, _, err := manet.DialArgs(raddr)
if err != nil {
return nil, err
}
quicConf := c.clientConfig.Clone()
quicConf.AllowConnectionWindowIncrease = allowWindowIncrease
if v == quic.Version1 {
// The endpoint has explicit support for QUIC v1, so we'll only use that version.
quicConf.Versions = []quic.VersionNumber{quic.Version1}
} else {
return nil, errors.New("unknown QUIC version")
}
tr, err := c.TransportForDial(netw, naddr)
if err != nil {
return nil, err
}
conn, err := tr.Dial(ctx, naddr, tlsConf, quicConf)
if err != nil {
tr.DecreaseCount()
return nil, err
}
return conn, nil
}
func (c *ConnManager) TransportForDial(network string, raddr *net.UDPAddr) (refCountedQuicTransport, error) {
if c.enableReuseport {
reuse, err := c.getReuse(network)
if err != nil {
return nil, err
}
return reuse.TransportForDial(network, raddr)
}
var laddr *net.UDPAddr
switch network {
case "udp4":
laddr = &net.UDPAddr{IP: net.IPv4zero, Port: 0}
case "udp6":
laddr = &net.UDPAddr{IP: net.IPv6zero, Port: 0}
}
2024-06-08 11:32:45 +00:00
conn, err := net.ListenUDP(network, laddr)
2023-08-21 03:50:38 +00:00
if err != nil {
return nil, err
}
2024-06-08 11:32:45 +00:00
return &singleOwnerTransport{Transport: quic.Transport{Conn: conn, StatelessResetKey: &c.srk}, packetConn: conn}, nil
2023-08-21 03:50:38 +00:00
}
func (c *ConnManager) Protocols() []int {
return []int{ma.P_QUIC_V1}
}
func (c *ConnManager) Close() error {
if !c.enableReuseport {
return nil
}
if err := c.reuseUDP6.Close(); err != nil {
return err
}
return c.reuseUDP4.Close()
}
2024-06-08 11:32:45 +00:00
func (c *ConnManager) ClientConfig() *quic.Config {
return c.clientConfig
2023-08-21 03:50:38 +00:00
}