This commit is contained in:
2026-02-19 10:07:43 +00:00
parent 007438e372
commit 6e637ecf77
1763 changed files with 60820 additions and 279516 deletions

View File

@@ -22,6 +22,7 @@ type Listener struct {
ch chan Conn
closeOnce sync.Once
closed chan struct{}
onClose func() // or nil
// NewConn, if non-nil, is called to create a new pair of connections
// when dialing. If nil, NewConn is used.
@@ -38,24 +39,29 @@ func Listen(addr string) *Listener {
}
// Addr implements net.Listener.Addr.
func (l *Listener) Addr() net.Addr {
return l.addr
func (ln *Listener) Addr() net.Addr {
return ln.addr
}
// Close closes the pipe listener.
func (l *Listener) Close() error {
l.closeOnce.Do(func() {
close(l.closed)
func (ln *Listener) Close() error {
var cleanup func()
ln.closeOnce.Do(func() {
cleanup = ln.onClose
close(ln.closed)
})
if cleanup != nil {
cleanup()
}
return nil
}
// Accept blocks until a new connection is available or the listener is closed.
func (l *Listener) Accept() (net.Conn, error) {
func (ln *Listener) Accept() (net.Conn, error) {
select {
case c := <-l.ch:
case c := <-ln.ch:
return c, nil
case <-l.closed:
case <-ln.closed:
return nil, net.ErrClosed
}
}
@@ -64,18 +70,18 @@ func (l *Listener) Accept() (net.Conn, error) {
// The provided Context must be non-nil. If the context expires before the
// connection is complete, an error is returned. Once successfully connected
// any expiration of the context will not affect the connection.
func (l *Listener) Dial(ctx context.Context, network, addr string) (_ net.Conn, err error) {
func (ln *Listener) Dial(ctx context.Context, network, addr string) (_ net.Conn, err error) {
if !strings.HasSuffix(network, "tcp") {
return nil, net.UnknownNetworkError(network)
}
if connAddr(addr) != l.addr {
if connAddr(addr) != ln.addr {
return nil, &net.AddrError{
Err: "invalid address",
Addr: addr,
}
}
newConn := l.NewConn
newConn := ln.NewConn
if newConn == nil {
newConn = func(network, addr string, maxBuf int) (Conn, Conn) {
return NewConn(addr, maxBuf)
@@ -92,9 +98,9 @@ func (l *Listener) Dial(ctx context.Context, network, addr string) (_ net.Conn,
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-l.closed:
case <-ln.closed:
return nil, net.ErrClosed
case l.ch <- s:
case ln.ch <- s:
return c, nil
}
}

View File

@@ -6,3 +6,87 @@
// in tests and other situations where you don't want to use the
// network.
package memnet
import (
"context"
"fmt"
"net"
"net/netip"
"tailscale.com/net/netx"
"tailscale.com/syncs"
)
var _ netx.Network = (*Network)(nil)
// Network implements [Network] using an in-memory network, usually
// used for testing.
//
// As of 2025-04-08, it only supports TCP.
//
// Its zero value is a valid [netx.Network] implementation.
type Network struct {
mu syncs.Mutex
lns map[string]*Listener // address -> listener
}
func (m *Network) Listen(network, address string) (net.Listener, error) {
if network != "tcp" && network != "tcp4" && network != "tcp6" {
return nil, fmt.Errorf("memNetwork: Listen called with unsupported network %q", network)
}
ap, err := netip.ParseAddrPort(address)
if err != nil {
return nil, fmt.Errorf("memNetwork: Listen called with invalid address %q: %w", address, err)
}
m.mu.Lock()
defer m.mu.Unlock()
if m.lns == nil {
m.lns = make(map[string]*Listener)
}
port := ap.Port()
for {
if port == 0 {
port = 33000
}
key := net.JoinHostPort(ap.Addr().String(), fmt.Sprint(port))
_, ok := m.lns[key]
if ok {
if ap.Port() != 0 {
return nil, fmt.Errorf("memNetwork: Listen called with duplicate address %q", address)
}
port++
continue
}
ln := Listen(key)
m.lns[key] = ln
ln.onClose = func() {
m.mu.Lock()
delete(m.lns, key)
m.mu.Unlock()
}
return ln, nil
}
}
func (m *Network) NewLocalTCPListener() net.Listener {
ln, err := m.Listen("tcp", "127.0.0.1:0")
if err != nil {
panic(fmt.Sprintf("memNetwork: failed to create local TCP listener: %v", err))
}
return ln
}
func (m *Network) Dial(ctx context.Context, network, address string) (net.Conn, error) {
if network != "tcp" && network != "tcp4" && network != "tcp6" {
return nil, fmt.Errorf("memNetwork: Dial called with unsupported network %q", network)
}
m.mu.Lock()
ln, ok := m.lns[address]
m.mu.Unlock()
if !ok {
return nil, fmt.Errorf("memNetwork: Dial called on unknown address %q", address)
}
return ln.Dial(ctx, network, address)
}