mirror of
				https://source.quilibrium.com/quilibrium/ceremonyclient.git
				synced 2025-11-04 07:07:26 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			257 lines
		
	
	
		
			6.3 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			257 lines
		
	
	
		
			6.3 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
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
 | 
						|
	enableDraft29   bool
 | 
						|
	enableReuseport bool
 | 
						|
	enableMetrics   bool
 | 
						|
 | 
						|
	serverConfig *quic.Config
 | 
						|
	clientConfig *quic.Config
 | 
						|
 | 
						|
	quicListenersMu sync.Mutex
 | 
						|
	quicListeners   map[string]quicListenerEntry
 | 
						|
 | 
						|
	srk quic.StatelessResetKey
 | 
						|
	mt  *metricsTracer
 | 
						|
}
 | 
						|
 | 
						|
type quicListenerEntry struct {
 | 
						|
	refCount int
 | 
						|
	ln       *quicListener
 | 
						|
}
 | 
						|
 | 
						|
func NewConnManager(statelessResetKey quic.StatelessResetKey, opts ...Option) (*ConnManager, error) {
 | 
						|
	cm := &ConnManager{
 | 
						|
		enableReuseport: true,
 | 
						|
		enableDraft29:   true,
 | 
						|
		quicListeners:   make(map[string]quicListenerEntry),
 | 
						|
		srk:             statelessResetKey,
 | 
						|
	}
 | 
						|
	for _, o := range opts {
 | 
						|
		if err := o(cm); err != nil {
 | 
						|
			return nil, err
 | 
						|
		}
 | 
						|
	}
 | 
						|
 | 
						|
	quicConf := quicConfig.Clone()
 | 
						|
 | 
						|
	if cm.enableMetrics {
 | 
						|
		cm.mt = newMetricsTracer()
 | 
						|
	}
 | 
						|
	quicConf.Tracer = func(ctx context.Context, p quiclogging.Perspective, ci quic.ConnectionID) quiclogging.ConnectionTracer {
 | 
						|
		tracers := make([]quiclogging.ConnectionTracer, 0, 2)
 | 
						|
		if qlogTracerDir != "" {
 | 
						|
			tracers = append(tracers, qloggerForDir(qlogTracerDir, p, ci))
 | 
						|
		}
 | 
						|
		if cm.mt != nil {
 | 
						|
			tracers = append(tracers, cm.mt.TracerForConnection(ctx, p, ci))
 | 
						|
		}
 | 
						|
		return quiclogging.NewMultiplexedConnectionTracer(tracers...)
 | 
						|
	}
 | 
						|
	serverConfig := quicConf.Clone()
 | 
						|
	if !cm.enableDraft29 {
 | 
						|
		serverConfig.Versions = []quic.VersionNumber{quic.Version1}
 | 
						|
	}
 | 
						|
 | 
						|
	cm.clientConfig = quicConf
 | 
						|
	cm.serverConfig = serverConfig
 | 
						|
	if cm.enableReuseport {
 | 
						|
		cm.reuseUDP4 = newReuse(&statelessResetKey, cm.mt)
 | 
						|
		cm.reuseUDP6 = newReuse(&statelessResetKey, cm.mt)
 | 
						|
	}
 | 
						|
	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) {
 | 
						|
	if !c.enableDraft29 {
 | 
						|
		if _, err := addr.ValueForProtocol(ma.P_QUIC); err == nil {
 | 
						|
			return nil, errors.New("can't listen on `/quic` multiaddr (QUIC draft 29 version) when draft 29 support is disabled")
 | 
						|
		}
 | 
						|
	}
 | 
						|
 | 
						|
	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
 | 
						|
		}
 | 
						|
		ln, err := newQuicListener(tr, c.serverConfig, c.enableDraft29)
 | 
						|
		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)
 | 
						|
	}
 | 
						|
 | 
						|
	conn, err := listenAndOptimize(network, laddr)
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
	tr := &singleOwnerTransport{Transport: quic.Transport{Conn: conn, StatelessResetKey: &c.srk}, packetConn: conn}
 | 
						|
	if c.mt != nil {
 | 
						|
		tr.Transport.Tracer = c.mt
 | 
						|
	}
 | 
						|
	return tr, nil
 | 
						|
}
 | 
						|
 | 
						|
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 if v == quic.VersionDraft29 {
 | 
						|
		quicConf.Versions = []quic.VersionNumber{quic.VersionDraft29}
 | 
						|
	} 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}
 | 
						|
	}
 | 
						|
	conn, err := listenAndOptimize(network, laddr)
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
	tr := &singleOwnerTransport{Transport: quic.Transport{Conn: conn, StatelessResetKey: &c.srk}, packetConn: conn}
 | 
						|
	if c.mt != nil {
 | 
						|
		tr.Transport.Tracer = c.mt
 | 
						|
	}
 | 
						|
 | 
						|
	return tr, nil
 | 
						|
}
 | 
						|
 | 
						|
func (c *ConnManager) Protocols() []int {
 | 
						|
	if c.enableDraft29 {
 | 
						|
		return []int{ma.P_QUIC, ma.P_QUIC_V1}
 | 
						|
	}
 | 
						|
	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()
 | 
						|
}
 | 
						|
 | 
						|
// listenAndOptimize same as net.ListenUDP, but also calls quic.OptimizeConn
 | 
						|
func listenAndOptimize(network string, laddr *net.UDPAddr) (net.PacketConn, error) {
 | 
						|
	conn, err := net.ListenUDP(network, laddr)
 | 
						|
	if err != nil {
 | 
						|
		return nil, err
 | 
						|
	}
 | 
						|
	return quic.OptimizeConn(conn)
 | 
						|
}
 |