mirror of
				https://source.quilibrium.com/quilibrium/ceremonyclient.git
				synced 2025-11-04 00:57:26 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			300 lines
		
	
	
		
			7.0 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			300 lines
		
	
	
		
			7.0 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
package basichost
 | 
						|
 | 
						|
import (
 | 
						|
	"context"
 | 
						|
	"io"
 | 
						|
	"net"
 | 
						|
	"net/netip"
 | 
						|
	"strconv"
 | 
						|
	"sync"
 | 
						|
	"time"
 | 
						|
 | 
						|
	"github.com/libp2p/go-libp2p/core/network"
 | 
						|
	inat "github.com/libp2p/go-libp2p/p2p/net/nat"
 | 
						|
 | 
						|
	ma "github.com/multiformats/go-multiaddr"
 | 
						|
	manet "github.com/multiformats/go-multiaddr/net"
 | 
						|
)
 | 
						|
 | 
						|
// NATManager is a simple interface to manage NAT devices.
 | 
						|
// It listens Listen and ListenClose notifications from the network.Network,
 | 
						|
// and tries to obtain port mappings for those.
 | 
						|
type NATManager interface {
 | 
						|
	GetMapping(ma.Multiaddr) ma.Multiaddr
 | 
						|
	HasDiscoveredNAT() bool
 | 
						|
	io.Closer
 | 
						|
}
 | 
						|
 | 
						|
// NewNATManager creates a NAT manager.
 | 
						|
func NewNATManager(net network.Network) NATManager {
 | 
						|
	return newNATManager(net)
 | 
						|
}
 | 
						|
 | 
						|
type entry struct {
 | 
						|
	protocol string
 | 
						|
	port     int
 | 
						|
}
 | 
						|
 | 
						|
type nat interface {
 | 
						|
	AddMapping(ctx context.Context, protocol string, port int) error
 | 
						|
	RemoveMapping(ctx context.Context, protocol string, port int) error
 | 
						|
	GetMapping(protocol string, port int) (netip.AddrPort, bool)
 | 
						|
	io.Closer
 | 
						|
}
 | 
						|
 | 
						|
// so we can mock it in tests
 | 
						|
var discoverNAT = func(ctx context.Context) (nat, error) { return inat.DiscoverNAT(ctx) }
 | 
						|
 | 
						|
// natManager takes care of adding + removing port mappings to the nat.
 | 
						|
// Initialized with the host if it has a NATPortMap option enabled.
 | 
						|
// natManager receives signals from the network, and check on nat mappings:
 | 
						|
//   - natManager listens to the network and adds or closes port mappings
 | 
						|
//     as the network signals Listen() or ListenClose().
 | 
						|
//   - closing the natManager closes the nat and its mappings.
 | 
						|
type natManager struct {
 | 
						|
	net   network.Network
 | 
						|
	natMx sync.RWMutex
 | 
						|
	nat   nat
 | 
						|
 | 
						|
	syncFlag chan struct{} // cap: 1
 | 
						|
 | 
						|
	tracked map[entry]bool // the bool is only used in doSync and has no meaning outside of that function
 | 
						|
 | 
						|
	refCount  sync.WaitGroup
 | 
						|
	ctx       context.Context
 | 
						|
	ctxCancel context.CancelFunc
 | 
						|
}
 | 
						|
 | 
						|
func newNATManager(net network.Network) *natManager {
 | 
						|
	ctx, cancel := context.WithCancel(context.Background())
 | 
						|
	nmgr := &natManager{
 | 
						|
		net:       net,
 | 
						|
		syncFlag:  make(chan struct{}, 1),
 | 
						|
		ctx:       ctx,
 | 
						|
		ctxCancel: cancel,
 | 
						|
		tracked:   make(map[entry]bool),
 | 
						|
	}
 | 
						|
	nmgr.refCount.Add(1)
 | 
						|
	go nmgr.background(ctx)
 | 
						|
	return nmgr
 | 
						|
}
 | 
						|
 | 
						|
// Close closes the natManager, closing the underlying nat
 | 
						|
// and unregistering from network events.
 | 
						|
func (nmgr *natManager) Close() error {
 | 
						|
	nmgr.ctxCancel()
 | 
						|
	nmgr.refCount.Wait()
 | 
						|
	return nil
 | 
						|
}
 | 
						|
 | 
						|
func (nmgr *natManager) HasDiscoveredNAT() bool {
 | 
						|
	nmgr.natMx.RLock()
 | 
						|
	defer nmgr.natMx.RUnlock()
 | 
						|
	return nmgr.nat != nil
 | 
						|
}
 | 
						|
 | 
						|
func (nmgr *natManager) background(ctx context.Context) {
 | 
						|
	defer nmgr.refCount.Done()
 | 
						|
 | 
						|
	defer func() {
 | 
						|
		nmgr.natMx.Lock()
 | 
						|
		defer nmgr.natMx.Unlock()
 | 
						|
 | 
						|
		if nmgr.nat != nil {
 | 
						|
			nmgr.nat.Close()
 | 
						|
		}
 | 
						|
	}()
 | 
						|
 | 
						|
	discoverCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
 | 
						|
	defer cancel()
 | 
						|
	natInstance, err := discoverNAT(discoverCtx)
 | 
						|
	if err != nil {
 | 
						|
		log.Info("DiscoverNAT error:", err)
 | 
						|
		return
 | 
						|
	}
 | 
						|
 | 
						|
	nmgr.natMx.Lock()
 | 
						|
	nmgr.nat = natInstance
 | 
						|
	nmgr.natMx.Unlock()
 | 
						|
 | 
						|
	// sign natManager up for network notifications
 | 
						|
	// we need to sign up here to avoid missing some notifs
 | 
						|
	// before the NAT has been found.
 | 
						|
	nmgr.net.Notify((*nmgrNetNotifiee)(nmgr))
 | 
						|
	defer nmgr.net.StopNotify((*nmgrNetNotifiee)(nmgr))
 | 
						|
 | 
						|
	nmgr.doSync() // sync one first.
 | 
						|
	for {
 | 
						|
		select {
 | 
						|
		case <-nmgr.syncFlag:
 | 
						|
			nmgr.doSync() // sync when our listen addresses chnage.
 | 
						|
		case <-ctx.Done():
 | 
						|
			return
 | 
						|
		}
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (nmgr *natManager) sync() {
 | 
						|
	select {
 | 
						|
	case nmgr.syncFlag <- struct{}{}:
 | 
						|
	default:
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
// doSync syncs the current NAT mappings, removing any outdated mappings and adding any
 | 
						|
// new mappings.
 | 
						|
func (nmgr *natManager) doSync() {
 | 
						|
	for e := range nmgr.tracked {
 | 
						|
		nmgr.tracked[e] = false
 | 
						|
	}
 | 
						|
	var newAddresses []entry
 | 
						|
	for _, maddr := range nmgr.net.ListenAddresses() {
 | 
						|
		// Strip the IP
 | 
						|
		maIP, rest := ma.SplitFirst(maddr)
 | 
						|
		if maIP == nil || rest == nil {
 | 
						|
			continue
 | 
						|
		}
 | 
						|
 | 
						|
		switch maIP.Protocol().Code {
 | 
						|
		case ma.P_IP6, ma.P_IP4:
 | 
						|
		default:
 | 
						|
			continue
 | 
						|
		}
 | 
						|
 | 
						|
		// Only bother if we're listening on an unicast / unspecified IP.
 | 
						|
		ip := net.IP(maIP.RawValue())
 | 
						|
		if !ip.IsGlobalUnicast() && !ip.IsUnspecified() {
 | 
						|
			continue
 | 
						|
		}
 | 
						|
 | 
						|
		// Extract the port/protocol
 | 
						|
		proto, _ := ma.SplitFirst(rest)
 | 
						|
		if proto == nil {
 | 
						|
			continue
 | 
						|
		}
 | 
						|
 | 
						|
		var protocol string
 | 
						|
		switch proto.Protocol().Code {
 | 
						|
		case ma.P_TCP:
 | 
						|
			protocol = "tcp"
 | 
						|
		case ma.P_UDP:
 | 
						|
			protocol = "udp"
 | 
						|
		default:
 | 
						|
			continue
 | 
						|
		}
 | 
						|
		port, err := strconv.ParseUint(proto.Value(), 10, 16)
 | 
						|
		if err != nil {
 | 
						|
			// bug in multiaddr
 | 
						|
			panic(err)
 | 
						|
		}
 | 
						|
		e := entry{protocol: protocol, port: int(port)}
 | 
						|
		if _, ok := nmgr.tracked[e]; ok {
 | 
						|
			nmgr.tracked[e] = true
 | 
						|
		} else {
 | 
						|
			newAddresses = append(newAddresses, e)
 | 
						|
		}
 | 
						|
	}
 | 
						|
 | 
						|
	var wg sync.WaitGroup
 | 
						|
	defer wg.Wait()
 | 
						|
 | 
						|
	// Close old mappings
 | 
						|
	for e, v := range nmgr.tracked {
 | 
						|
		if !v {
 | 
						|
			nmgr.nat.RemoveMapping(nmgr.ctx, e.protocol, e.port)
 | 
						|
			delete(nmgr.tracked, e)
 | 
						|
		}
 | 
						|
	}
 | 
						|
 | 
						|
	// Create new mappings.
 | 
						|
	for _, e := range newAddresses {
 | 
						|
		if err := nmgr.nat.AddMapping(nmgr.ctx, e.protocol, e.port); err != nil {
 | 
						|
			log.Errorf("failed to port-map %s port %d: %s", e.protocol, e.port, err)
 | 
						|
		}
 | 
						|
		nmgr.tracked[e] = false
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func (nmgr *natManager) GetMapping(addr ma.Multiaddr) ma.Multiaddr {
 | 
						|
	nmgr.natMx.Lock()
 | 
						|
	defer nmgr.natMx.Unlock()
 | 
						|
 | 
						|
	if nmgr.nat == nil { // NAT not yet initialized
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
 | 
						|
	var found bool
 | 
						|
	var proto int // ma.P_TCP or ma.P_UDP
 | 
						|
	transport, rest := ma.SplitFunc(addr, func(c ma.Component) bool {
 | 
						|
		if found {
 | 
						|
			return true
 | 
						|
		}
 | 
						|
		proto = c.Protocol().Code
 | 
						|
		found = proto == ma.P_TCP || proto == ma.P_UDP
 | 
						|
		return false
 | 
						|
	})
 | 
						|
	if !manet.IsThinWaist(transport) {
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
 | 
						|
	naddr, err := manet.ToNetAddr(transport)
 | 
						|
	if err != nil {
 | 
						|
		log.Error("error parsing net multiaddr %q: %s", transport, err)
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
 | 
						|
	var (
 | 
						|
		ip       net.IP
 | 
						|
		port     int
 | 
						|
		protocol string
 | 
						|
	)
 | 
						|
	switch naddr := naddr.(type) {
 | 
						|
	case *net.TCPAddr:
 | 
						|
		ip = naddr.IP
 | 
						|
		port = naddr.Port
 | 
						|
		protocol = "tcp"
 | 
						|
	case *net.UDPAddr:
 | 
						|
		ip = naddr.IP
 | 
						|
		port = naddr.Port
 | 
						|
		protocol = "udp"
 | 
						|
	default:
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
 | 
						|
	if !ip.IsGlobalUnicast() && !ip.IsUnspecified() {
 | 
						|
		// We only map global unicast & unspecified addresses ports, not broadcast, multicast, etc.
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
 | 
						|
	extAddr, ok := nmgr.nat.GetMapping(protocol, port)
 | 
						|
	if !ok {
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
 | 
						|
	var mappedAddr net.Addr
 | 
						|
	switch naddr.(type) {
 | 
						|
	case *net.TCPAddr:
 | 
						|
		mappedAddr = net.TCPAddrFromAddrPort(extAddr)
 | 
						|
	case *net.UDPAddr:
 | 
						|
		mappedAddr = net.UDPAddrFromAddrPort(extAddr)
 | 
						|
	}
 | 
						|
	mappedMaddr, err := manet.FromNetAddr(mappedAddr)
 | 
						|
	if err != nil {
 | 
						|
		log.Errorf("mapped addr can't be turned into a multiaddr %q: %s", mappedAddr, err)
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
	extMaddr := mappedMaddr
 | 
						|
	if rest != nil {
 | 
						|
		extMaddr = ma.Join(extMaddr, rest)
 | 
						|
	}
 | 
						|
	return extMaddr
 | 
						|
}
 | 
						|
 | 
						|
type nmgrNetNotifiee natManager
 | 
						|
 | 
						|
func (nn *nmgrNetNotifiee) natManager() *natManager                          { return (*natManager)(nn) }
 | 
						|
func (nn *nmgrNetNotifiee) Listen(network.Network, ma.Multiaddr)             { nn.natManager().sync() }
 | 
						|
func (nn *nmgrNetNotifiee) ListenClose(n network.Network, addr ma.Multiaddr) { nn.natManager().sync() }
 | 
						|
func (nn *nmgrNetNotifiee) Connected(network.Network, network.Conn)          {}
 | 
						|
func (nn *nmgrNetNotifiee) Disconnected(network.Network, network.Conn)       {}
 |