listener.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package nettest
     6  
     7  import (
     8  	"context"
     9  	"internal/gate"
    10  	"net"
    11  	"net/netip"
    12  )
    13  
    14  // Listener is an in-memory test implementation of net.Listener.
    15  type Listener struct {
    16  	gate      gate.Gate
    17  	queue     queue[*Conn]
    18  	closed    bool
    19  	acceptErr error
    20  	closeErr  error
    21  
    22  	addr     net.Addr
    23  	nextaddr netip.AddrPort
    24  }
    25  
    26  // NewListener returns a new Listener.
    27  func NewListener() *Listener {
    28  	return &Listener{
    29  		gate:     gate.New(false),
    30  		addr:     net.TCPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:1000")),
    31  		nextaddr: netip.MustParseAddrPort("127.0.0.1:10001"),
    32  	}
    33  }
    34  
    35  // Close closes the listener.
    36  // Any blocked Accept operations will be unblocked and return errors.
    37  func (li *Listener) Close() error {
    38  	li.lock()
    39  	defer li.unlock()
    40  	err := li.closeErr
    41  	li.closed = true
    42  	li.acceptErr = net.ErrClosed
    43  	li.closeErr = net.ErrClosed
    44  	if err != nil {
    45  		return &net.OpError{
    46  			Op:   "close",
    47  			Net:  "tcp",
    48  			Addr: li.addr,
    49  			Err:  err,
    50  		}
    51  	}
    52  	return err
    53  }
    54  
    55  // Addr returns the listener's network address.
    56  //
    57  // The address is always a *net.TCPAddr.
    58  func (li *Listener) Addr() net.Addr {
    59  	li.lock()
    60  	defer li.unlock()
    61  	return li.addr
    62  }
    63  
    64  // SetAddr sets the listener's network address.
    65  func (li *Listener) SetAddr(addr net.Addr) {
    66  	li.lock()
    67  	defer li.unlock()
    68  	li.addr = addr
    69  }
    70  
    71  // NewConn returns a new connection to the listener.
    72  //
    73  // Accept will return the other side of the conn.
    74  func (li *Listener) NewConn() *Conn {
    75  	return li.NewConnConfig(func(*Conn) {})
    76  }
    77  
    78  // NewConnConfig returns a new connection to the listener.
    79  //
    80  // The function f is called with the new client connection.
    81  // After f returns, Accept will return the other side of the connection.
    82  //
    83  // For example, to create a connection from a specific IP address:
    84  //
    85  //	conn := li.NewConnConfig(func(conn *nettest.Conn) {
    86  //		conn.SetLocalAddr(net.TCPAddrFromAddrPort(netip.MustParseAddrPort("10.0.0.1:1234")))
    87  //	})
    88  func (li *Listener) NewConnConfig(f func(*Conn)) *Conn {
    89  	li.lock()
    90  	defer li.unlock()
    91  	cli, srv := newConnPair(
    92  		net.TCPAddrFromAddrPort(li.nextaddr),
    93  		li.addr,
    94  	)
    95  	li.nextaddr = netip.AddrPortFrom(li.nextaddr.Addr(), li.nextaddr.Port()+1)
    96  	f(cli)
    97  	li.queue.push(srv)
    98  	return cli
    99  }
   100  
   101  // Accept waits for and returns the next connection to the listener.
   102  //
   103  // The connections returned by Accept are always [*Conn]s.
   104  func (li *Listener) Accept() (net.Conn, error) {
   105  	li.gate.WaitAndLock(context.Background())
   106  	defer li.unlock()
   107  	if li.acceptErr != nil && li.queue.len() == 0 {
   108  		return nil, &net.OpError{
   109  			Op:   "accept",
   110  			Net:  "tcp",
   111  			Addr: li.addr,
   112  			Err:  li.acceptErr,
   113  		}
   114  	}
   115  	return li.queue.pop(), nil
   116  }
   117  
   118  // SetAcceptError causes any currently blocked and future Accept calls to return
   119  // a net.OpError wrapping err.
   120  // Accept will return any available connections before returning the error,
   121  // including connections created after the error is set.
   122  // A nil error restores the usual behavior.
   123  func (li *Listener) SetAcceptError(err error) {
   124  	li.gate.Lock()
   125  	defer li.unlock()
   126  	if !li.closed {
   127  		li.acceptErr = err
   128  	}
   129  }
   130  
   131  // SetCloseError sets the error returned by Close.
   132  // Close still closes the listener.
   133  // A nil error restores the usual behavior.
   134  func (li *Listener) SetCloseError(err error) {
   135  	li.gate.Lock()
   136  	defer li.unlock()
   137  	if !li.closed {
   138  		li.closeErr = err
   139  	}
   140  }
   141  
   142  func (li *Listener) lock() {
   143  	li.gate.Lock()
   144  }
   145  
   146  func (li *Listener) unlock() {
   147  	li.gate.Unlock(li.acceptErr != nil || li.queue.len() > 0)
   148  }
   149
View as plain text