From a076213f5854e60d66376be2e57c2f63132bad16 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 12 Jun 2023 20:43:14 +0000
Subject: [PATCH] net/memnet: add optional Listener.NewConn config knob

Updates tailscale/corp#11620

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 net/memnet/listener.go | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/net/memnet/listener.go b/net/memnet/listener.go
index f6ac809bd..d84a2e443 100644
--- a/net/memnet/listener.go
+++ b/net/memnet/listener.go
@@ -22,6 +22,10 @@ type Listener struct {
 	ch        chan Conn
 	closeOnce sync.Once
 	closed    chan struct{}
+
+	// NewConn, if non-nil, is called to create a new pair of connections
+	// when dialing. If nil, NewConn is used.
+	NewConn func(network, addr string, maxBuf int) (Conn, Conn)
 }
 
 // Listen returns a new Listener for the provided address.
@@ -70,7 +74,14 @@ func (l *Listener) Dial(ctx context.Context, network, addr string) (_ net.Conn,
 			Addr: addr,
 		}
 	}
-	c, s := NewConn(addr, bufferSize)
+
+	newConn := l.NewConn
+	if newConn == nil {
+		newConn = func(network, addr string, maxBuf int) (Conn, Conn) {
+			return NewConn(addr, maxBuf)
+		}
+	}
+	c, s := newConn(network, addr, bufferSize)
 	defer func() {
 		if err != nil {
 			c.Close()