From 49f84c7c457641702f0c049a62aba592303c32e4 Mon Sep 17 00:00:00 2001
From: Tom Swindell <t.swindell@rubyx.co.uk>
Date: Sat, 17 Dec 2016 00:17:22 +0000
Subject: [PATCH] Added UDP datagram packet connection support.

Signed-off-by: Tom Swindell <t.swindell@rubyx.co.uk>
---
 net.go      | 75 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 net_test.go | 53 +++++++++++++++++++++++++++++++++++++
 2 files changed, 128 insertions(+)

diff --git a/net.go b/net.go
index b747339..9ee4ade 100644
--- a/net.go
+++ b/net.go
@@ -248,6 +248,81 @@ func WrapNetListener(nl net.Listener) (Listener, error) {
 	}, nil
 }
 
+// A PacketConn is a generic packet oriented network connection which uses an
+// underlying net.PacketConn, wrapped with the locally bound Multiaddr.
+type PacketConn interface {
+	Connection() net.PacketConn
+
+	Multiaddr() ma.Multiaddr
+
+	ReadFrom(b []byte) (int, ma.Multiaddr, error)
+	WriteTo(b []byte, maddr ma.Multiaddr) (int, error)
+
+	Close() error
+}
+
+// maPacketConn implements PacketConn
+type maPacketConn struct {
+	net.PacketConn
+	laddr ma.Multiaddr
+}
+
+// Connection returns the embedded net.PacketConn.
+func (l *maPacketConn) Connection() net.PacketConn {
+	return l.PacketConn
+}
+
+// Multiaddr returns the bound local Multiaddr.
+func (l *maPacketConn) Multiaddr() ma.Multiaddr {
+	return l.laddr
+}
+
+func (l *maPacketConn) ReadFrom(b []byte) (int, ma.Multiaddr, error) {
+	n, addr, err := l.PacketConn.ReadFrom(b)
+	maddr, _ := FromNetAddr(addr)
+	return n, maddr, err
+}
+
+func (l *maPacketConn) WriteTo(b []byte, maddr ma.Multiaddr) (int, error) {
+	addr, err := ToNetAddr(maddr)
+	if err != nil {
+		return 0, err
+	}
+	return l.PacketConn.WriteTo(b, addr)
+}
+
+// ListenPacket announces on the local network address laddr.
+// The Multiaddr must be a packet driven network, like udp4 or udp6.
+// See Dial for the syntax of laddr.
+func ListenPacket(laddr ma.Multiaddr) (PacketConn, error) {
+	lnet, lnaddr, err := DialArgs(laddr)
+	if err != nil {
+		return nil, err
+	}
+
+	pc, err := net.ListenPacket(lnet, lnaddr)
+	if err != nil {
+		return nil, err
+	}
+
+	// We want to fetch the new multiaddr from the listener, as it may
+	// have resolved to some other value. WrapPacketConn does this.
+	return WrapPacketConn(pc)
+}
+
+// WrapPacketConn wraps a net.PacketConn with a manet.PacketConn.
+func WrapPacketConn(pc net.PacketConn) (PacketConn, error) {
+	laddr, err := FromNetAddr(pc.LocalAddr())
+	if err != nil {
+		return nil, err
+	}
+
+	return &maPacketConn{
+		PacketConn: pc,
+		laddr:      laddr,
+	}, nil
+}
+
 // InterfaceMultiaddrs will return the addresses matching net.InterfaceAddrs
 func InterfaceMultiaddrs() ([]ma.Multiaddr, error) {
 	addrs, err := net.InterfaceAddrs()
diff --git a/net_test.go b/net_test.go
index 76333d1..44bdfd7 100644
--- a/net_test.go
+++ b/net_test.go
@@ -244,6 +244,59 @@ func TestListenAndDial(t *testing.T) {
 	wg.Wait()
 }
 
+func TestListenPacketAndDial(t *testing.T) {
+	maddr := newMultiaddr(t, "/ip4/127.0.0.1/udp/4324")
+	pc, err := ListenPacket(maddr)
+	if err != nil {
+		t.Fatal("failed to listen", err)
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go func() {
+		if !pc.Multiaddr().Equal(maddr) {
+			t.Fatal("connection multiaddr not equal:", maddr, pc.Multiaddr())
+		}
+
+		buffer := make([]byte, 1024)
+		_, addr, err := pc.ReadFrom(buffer)
+		if err != nil {
+			t.Fatal("failed to read into buffer", err)
+		}
+		pc.WriteTo(buffer, addr)
+
+		wg.Done()
+	}()
+
+	cn, err := Dial(maddr)
+	if err != nil {
+		t.Fatal("failed to dial", err)
+	}
+
+	buf := make([]byte, 1024)
+	if _, err := cn.Write([]byte("beep boop")); err != nil {
+		t.Fatal("failed to write", err)
+	}
+
+	if _, err := cn.Read(buf); err != nil {
+		t.Fatal("failed to read:", buf, err)
+	}
+
+	if !bytes.Equal(buf[:9], []byte("beep boop")) {
+		t.Fatal("failed to echk:", buf)
+	}
+
+	maddr2 := cn.RemoteMultiaddr()
+	if !maddr2.Equal(maddr) {
+		t.Fatal("remote multiaddr not equal:", maddr, maddr2)
+	}
+
+	cn.Close()
+	pc.Close()
+	wg.Wait()
+}
+
 func TestIPLoopback(t *testing.T) {
 	if IP4Loopback.String() != "/ip4/127.0.0.1" {
 		t.Error("IP4Loopback incorrect:", IP4Loopback)