net/memnet: add optional Listener.NewConn config knob

Updates tailscale/corp#11620

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
This commit is contained in:
Brad Fitzpatrick 2023-06-12 20:43:14 +00:00 committed by Brad Fitzpatrick
parent 4451a7c364
commit a076213f58

View File

@ -22,6 +22,10 @@ type Listener struct {
ch chan Conn ch chan Conn
closeOnce sync.Once closeOnce sync.Once
closed chan struct{} 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. // 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, 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() { defer func() {
if err != nil { if err != nil {
c.Close() c.Close()