mirror of https://github.com/libp2p/go-libp2p.git
lnykww
6 years ago
committed by
Marten Seemann
Marten Seemann
4 changed files with 275 additions and 50 deletions
@ -0,0 +1,137 @@ |
|||
package libp2pquic |
|||
|
|||
import ( |
|||
"net" |
|||
"sync" |
|||
"sync/atomic" |
|||
|
|||
"github.com/vishvananda/netlink" |
|||
) |
|||
|
|||
type reuseConn struct { |
|||
net.PacketConn |
|||
refCount int32 // to be used as an atomic
|
|||
} |
|||
|
|||
func newReuseConn(conn net.PacketConn) *reuseConn { |
|||
return &reuseConn{PacketConn: conn} |
|||
} |
|||
|
|||
func (c *reuseConn) IncreaseCount() { atomic.AddInt32(&c.refCount, 1) } |
|||
func (c *reuseConn) DecreaseCount() { atomic.AddInt32(&c.refCount, -1) } |
|||
func (c *reuseConn) GetCount() int { return int(atomic.LoadInt32(&c.refCount)) } |
|||
|
|||
type reuse struct { |
|||
mutex sync.Mutex |
|||
|
|||
unicast map[string] /* IP.String() */ map[int] /* port */ *reuseConn |
|||
// global contains connections that are listening on 0.0.0.0 / ::
|
|||
global map[int]*reuseConn |
|||
} |
|||
|
|||
func newReuse() *reuse { |
|||
return &reuse{ |
|||
unicast: make(map[string]map[int]*reuseConn), |
|||
global: make(map[int]*reuseConn), |
|||
} |
|||
} |
|||
|
|||
func (r *reuse) getSourceIPs(network string, raddr *net.UDPAddr) ([]net.IP, error) { |
|||
// Determine the source address that the kernel would use for this IP address.
|
|||
// Note: This only works on Linux.
|
|||
// On other OSes, this will return a netlink.ErrNotImplemetned.
|
|||
routes, err := (&netlink.Handle{}).RouteGet(raddr.IP) |
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
|
|||
ips := make([]net.IP, 0, len(routes)) |
|||
for _, route := range routes { |
|||
ips = append(ips, route.Src) |
|||
} |
|||
return ips, nil |
|||
} |
|||
|
|||
func (r *reuse) Dial(network string, raddr *net.UDPAddr) (*reuseConn, error) { |
|||
ips, err := r.getSourceIPs(network, raddr) |
|||
if err != nil && err != netlink.ErrNotImplemented { |
|||
return nil, err |
|||
} |
|||
|
|||
r.mutex.Lock() |
|||
defer r.mutex.Unlock() |
|||
|
|||
conn, err := r.dialLocked(network, raddr, ips) |
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
conn.IncreaseCount() |
|||
return conn, nil |
|||
} |
|||
|
|||
func (r *reuse) dialLocked(network string, raddr *net.UDPAddr, ips []net.IP) (*reuseConn, error) { |
|||
for _, ip := range ips { |
|||
// We already have at least one suitable connection...
|
|||
if conns, ok := r.unicast[ip.String()]; ok { |
|||
// ... we don't care which port we're dialing from. Just use the first.
|
|||
for _, c := range conns { |
|||
return c, nil |
|||
} |
|||
} |
|||
} |
|||
|
|||
// Use a connection listening on 0.0.0.0 (or ::).
|
|||
// Again, we don't care about the port number.
|
|||
for _, conn := range r.global { |
|||
return conn, nil |
|||
} |
|||
|
|||
// We don't have a connection 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 := net.ListenUDP(network, addr) |
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
rconn := newReuseConn(conn) |
|||
r.global[conn.LocalAddr().(*net.UDPAddr).Port] = rconn |
|||
return rconn, nil |
|||
} |
|||
|
|||
func (r *reuse) Listen(network string, laddr *net.UDPAddr) (*reuseConn, error) { |
|||
conn, err := net.ListenUDP(network, laddr) |
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
localAddr := conn.LocalAddr().(*net.UDPAddr) |
|||
|
|||
rconn := newReuseConn(conn) |
|||
rconn.IncreaseCount() |
|||
|
|||
r.mutex.Lock() |
|||
defer r.mutex.Unlock() |
|||
|
|||
// Deal with listen on a global address
|
|||
if laddr.IP.IsUnspecified() { |
|||
// The kernel already checked that the laddr is not already listen
|
|||
// so we need not check here (when we create ListenUDP).
|
|||
r.global[laddr.Port] = rconn |
|||
return rconn, err |
|||
} |
|||
|
|||
// Deal with listen on a unicast address
|
|||
if _, ok := r.unicast[localAddr.IP.String()]; !ok { |
|||
r.unicast[laddr.IP.String()] = make(map[int]*reuseConn) |
|||
} |
|||
|
|||
// The kernel already checked that the laddr is not already listen
|
|||
// so we need not check here (when we create ListenUDP).
|
|||
r.unicast[laddr.IP.String()][localAddr.Port] = rconn |
|||
return rconn, err |
|||
} |
|||
@ -0,0 +1,79 @@ |
|||
package libp2pquic |
|||
|
|||
import ( |
|||
"net" |
|||
"runtime" |
|||
|
|||
. "github.com/onsi/ginkgo" |
|||
. "github.com/onsi/gomega" |
|||
) |
|||
|
|||
var _ = Describe("Reuse", func() { |
|||
var reuse *reuse |
|||
|
|||
BeforeEach(func() { |
|||
reuse = newReuse() |
|||
}) |
|||
|
|||
It("creates a new global connection when listening on 0.0.0.0", func() { |
|||
addr, err := net.ResolveUDPAddr("udp4", "0.0.0.0:0") |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
conn, err := reuse.Listen("udp4", addr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(conn.GetCount()).To(Equal(1)) |
|||
}) |
|||
|
|||
It("creates a new global connection when listening on [::]", func() { |
|||
addr, err := net.ResolveUDPAddr("udp6", "[::]:1234") |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
conn, err := reuse.Listen("udp6", addr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(conn.GetCount()).To(Equal(1)) |
|||
}) |
|||
|
|||
It("creates a new global connection when dialing", func() { |
|||
addr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234") |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
conn, err := reuse.Dial("udp4", addr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(conn.GetCount()).To(Equal(1)) |
|||
laddr := conn.LocalAddr().(*net.UDPAddr) |
|||
Expect(laddr.IP.String()).To(Equal("0.0.0.0")) |
|||
Expect(laddr.Port).ToNot(BeZero()) |
|||
}) |
|||
|
|||
It("reuses a connection it created for listening when dialing", func() { |
|||
// listen
|
|||
addr, err := net.ResolveUDPAddr("udp4", "0.0.0.0:0") |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
lconn, err := reuse.Listen("udp4", addr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(lconn.GetCount()).To(Equal(1)) |
|||
// dial
|
|||
raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234") |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
conn, err := reuse.Dial("udp4", raddr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(conn.GetCount()).To(Equal(2)) |
|||
}) |
|||
|
|||
if runtime.GOOS == "linux" { |
|||
It("reuses a connection it created for listening on a specific interface", func() { |
|||
raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234") |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
ips, err := reuse.getSourceIPs("udp4", raddr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(ips).ToNot(BeEmpty()) |
|||
// listen
|
|||
addr, err := net.ResolveUDPAddr("udp4", ips[0].String()+":0") |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
lconn, err := reuse.Listen("udp4", addr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(lconn.GetCount()).To(Equal(1)) |
|||
// dial
|
|||
conn, err := reuse.Dial("udp4", raddr) |
|||
Expect(err).ToNot(HaveOccurred()) |
|||
Expect(conn.GetCount()).To(Equal(2)) |
|||
}) |
|||
} |
|||
}) |
|||
Loading…
Reference in new issue