114 lines
2.6 KiB
Go
114 lines
2.6 KiB
Go
package tcpreuse
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
|
|
"github.com/libp2p/go-reuseport"
|
|
ma "github.com/multiformats/go-multiaddr"
|
|
manet "github.com/multiformats/go-multiaddr/net"
|
|
)
|
|
|
|
type dialer interface {
|
|
Dial(network, addr string) (net.Conn, error)
|
|
DialContext(ctx context.Context, network, addr string) (net.Conn, error)
|
|
}
|
|
|
|
// Dial dials the given multiaddr, reusing ports we're currently listening on if
|
|
// possible.
|
|
//
|
|
// Dial attempts to be smart about choosing the source port. For example, If
|
|
// we're dialing a loopback address and we're listening on one or more loopback
|
|
// ports, Dial will randomly choose one of the loopback ports and addresses and
|
|
// reuse it.
|
|
func (t *Transport) Dial(raddr ma.Multiaddr) (manet.Conn, error) {
|
|
return t.DialContext(context.Background(), raddr)
|
|
}
|
|
|
|
// DialContext is like Dial but takes a context.
|
|
func (t *Transport) DialContext(ctx context.Context, raddr ma.Multiaddr) (manet.Conn, error) {
|
|
network, addr, err := manet.DialArgs(raddr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var d dialer
|
|
switch network {
|
|
case "tcp4":
|
|
d = t.v4.getDialer(network)
|
|
case "tcp6":
|
|
d = t.v6.getDialer(network)
|
|
default:
|
|
return nil, ErrWrongProto
|
|
}
|
|
conn, err := d.DialContext(ctx, network, addr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
maconn, err := manet.WrapNetConn(conn)
|
|
if err != nil {
|
|
conn.Close()
|
|
return nil, err
|
|
}
|
|
return maconn, nil
|
|
}
|
|
|
|
func (n *network) getDialer(network string) dialer {
|
|
n.mu.RLock()
|
|
d := n.dialer
|
|
n.mu.RUnlock()
|
|
if d == nil {
|
|
n.mu.Lock()
|
|
defer n.mu.Unlock()
|
|
|
|
if n.dialer == nil {
|
|
n.dialer = n.makeDialer(network)
|
|
}
|
|
d = n.dialer
|
|
}
|
|
return d
|
|
}
|
|
|
|
func (n *network) makeDialer(network string) dialer {
|
|
if !reuseport.Available() {
|
|
log.Debug("reuseport not available")
|
|
return &net.Dialer{}
|
|
}
|
|
|
|
var unspec net.IP
|
|
switch network {
|
|
case "tcp4":
|
|
unspec = net.IPv4zero
|
|
case "tcp6":
|
|
unspec = net.IPv6unspecified
|
|
default:
|
|
panic("invalid network: must be either tcp4 or tcp6")
|
|
}
|
|
|
|
// How many ports are we listening on.
|
|
var port = 0
|
|
for l := range n.listeners {
|
|
newPort := l.Addr().(*net.TCPAddr).Port
|
|
switch {
|
|
case newPort == 0: // Any port, ignore (really, we shouldn't get this case...).
|
|
case port == 0: // Haven't selected a port yet, choose this one.
|
|
port = newPort
|
|
case newPort == port: // Same as the selected port, continue...
|
|
default: // Multiple ports, use the multi dialer
|
|
return newMultiDialer(unspec, n.listeners)
|
|
}
|
|
}
|
|
|
|
// None.
|
|
if port == 0 {
|
|
return &net.Dialer{}
|
|
}
|
|
|
|
// One. Always dial from the single port we're listening on.
|
|
laddr := &net.TCPAddr{
|
|
IP: unspec,
|
|
Port: port,
|
|
}
|
|
|
|
return (*singleDialer)(laddr)
|
|
}
|