From 5bd3bf1d6d3e8a1b400079b7c37a322fa4ea1efb Mon Sep 17 00:00:00 2001
From: Brian Picciano <mediocregopher@gmail.com>
Date: Fri, 1 Mar 2019 14:21:43 -0500
Subject: [PATCH] mnet: fix PacketConn support

---
 mnet/mnet.go | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/mnet/mnet.go b/mnet/mnet.go
index dd032e7..b6c1ca7 100644
--- a/mnet/mnet.go
+++ b/mnet/mnet.go
@@ -9,13 +9,18 @@ import (
 
 	"github.com/mediocregopher/mediocre-go-lib/mcfg"
 	"github.com/mediocregopher/mediocre-go-lib/mctx"
+	"github.com/mediocregopher/mediocre-go-lib/merr"
 	"github.com/mediocregopher/mediocre-go-lib/mlog"
 	"github.com/mediocregopher/mediocre-go-lib/mrun"
 )
 
 // Listener is returned by WithListen and simply wraps a net.Listener.
 type Listener struct {
+	// One of these will be populated during the start hook, depending on the
+	// protocol configured.
 	net.Listener
+	net.PacketConn
+
 	ctx context.Context
 
 	// If set to true before mrun's stop event is run, the stop event will not
@@ -28,6 +33,13 @@ type listenerOpts struct {
 	defaultAddr string
 }
 
+func (lOpts listenerOpts) isPacketConn() bool {
+	proto := strings.ToLower(lOpts.proto)
+	return strings.HasPrefix(proto, "udp") ||
+		proto == "unixgram" ||
+		strings.HasPrefix(proto, "ip")
+}
+
 // ListenerOpt is a value which adjusts the behavior of WithListener.
 type ListenerOpt func(*listenerOpts)
 
@@ -69,12 +81,21 @@ func WithListener(ctx context.Context, opts ...ListenerOpt) (context.Context, *L
 
 	l.ctx = mrun.WithStartHook(l.ctx, func(context.Context) error {
 		var err error
-		if l.Listener, err = net.Listen(lOpts.proto, *addr); err != nil {
-			return err
-		}
+
 		l.ctx = mctx.Annotate(l.ctx,
 			"proto", lOpts.proto,
-			"addr", l.Addr().String())
+			"addr", *addr)
+
+		if lOpts.isPacketConn() {
+			l.PacketConn, err = net.ListenPacket(lOpts.proto, *addr)
+		} else {
+			l.Listener, err = net.Listen(lOpts.proto, *addr)
+		}
+		if err != nil {
+			return merr.Wrap(err, l.ctx)
+		}
+
+		l.ctx = mctx.Annotate(l.ctx, "addr", l.Addr().String())
 		mlog.Info("listening", l.ctx)
 		return nil
 	})