src/internal/nettest/conn.go - go - Git at Google

 // Copyright 2026 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package nettest

 import (
 	"bytes"
 	"errors"
 	"io"
 	"math"
 	"net"
 	"net/netip"
 	"os"
 	"time"
 )

 // Conn is an in-memory test implementation of net.Conn.
 type Conn struct {
 	// Conns come in pairs.
 	// Writes to one Conn are read by its peer, and vice-versa.
 	//
 	// A connHalf handles one direction of data flow.
 	// A Conn consists of read and write halves.
 	// A Conn's peer has the same halves, only swapped.
 	//
 	// A Conn reads from r and writes to w.
 	r, w *connHalf

 	// peer is the other endpoint.
 	peer *Conn
 }

 // NewConnPair returns a pair of connected Conns.
 func NewConnPair() (*Conn, *Conn) {
 	return newConnPair(
 		net.TCPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:10000")),
 		net.TCPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:10001")),
 	)
 }

 func newConnPair(addr1, addr2 net.Addr) (*Conn, *Conn) {
 	h1 := newConnHalf(addr1)
 	h2 := newConnHalf(addr2)
 	c1 := &Conn{r: h1, w: h2}
 	c2 := &Conn{r: h2, w: h1}
 	c1.peer = c2
 	c2.peer = c1
 	c1.SetReadBufferSize(-1)
 	c2.SetReadBufferSize(-1)
 	return c1, c2
 }

 // Peer returns the other end of the connection.
 func (c *Conn) Peer() *Conn {
 	return c.peer
 }

 // Read reads data from the connection.
 func (c *Conn) Read(b []byte) (n int, err error) {
 	n, err = c.r.read(b)
 	if err != nil && err != io.EOF {
 		err = &net.OpError{
 			Op:     "read",
 			Net:    "tcp",
 			Source: c.RemoteAddr(),
 			Addr:   c.LocalAddr(),
 			Err:    err,
 		}
 	}
 	return n, err
 }

 // CanRead reports whether Read can proceed without blocking.
 func (c *Conn) CanRead() bool {
 	return c.r.canRead()
 }

 // Write writes data to the connection.
 func (c *Conn) Write(b []byte) (n int, err error) {
 	n, err = c.w.write(b)
 	if err != nil {
 		err = &net.OpError{
 			Op:     "write",
 			Net:    "tcp",
 			Source: c.LocalAddr(),
 			Addr:   c.RemoteAddr(),
 			Err:    err,
 		}
 	}
 	return n, err
 }

 // IsClosed reports whether the connection has been closed.
 // A connection is closed if [CloseRead] and [CloseWrite] are both called,
 // or if [Close] is called.
 //
 // To identify when the other side of the Conn has been closed,
 // use Conn.Peer().IsClosed().
 func (c *Conn) IsClosed() bool {
 	c.r.lock()
 	readClosed := c.r.readClosed
 	c.r.unlock()
 	c.w.lock()
 	writeClosed := c.w.writeClosed
 	c.w.unlock()
 	return readClosed && writeClosed
 }

 var errClosedByPeer = errors.New("connection closed by peer")

 // CloseRead shuts down the reading side of the connection.
 func (c *Conn) CloseRead() error {
 	c.r.lock()
 	defer c.r.unlock()
 	c.r.buf.Reset() // discard unread data
 	c.r.readClosed = true
 	return nil
 }

 // CloseWrite shuts down the writing side of the connection.
 func (c *Conn) CloseWrite() error {
 	c.w.lock()
 	defer c.w.unlock()
 	c.w.writeClosed = true
 	return nil
 }

 // Close closes the connection.
 func (c *Conn) Close() error {
 	c.r.lock()
 	readClosed := c.r.readClosed
 	c.r.buf.Reset() // discard unread data
 	c.r.readClosed = true
 	err := c.r.closeErr
 	c.r.unlock()

 	c.w.lock()
 	writeClosed := c.w.writeClosed
 	c.w.writeClosed = true
 	c.w.unlock()

 	if readClosed && writeClosed {
 		err = net.ErrClosed
 	}
 	if err != nil {
 		err = &net.OpError{
 			Op:   "close",
 			Net:  "tcp",
 			Addr: c.LocalAddr(),
 			Err:  err,
 		}
 	}
 	return err
 }

 // SetCloseError sets the error returned by Close.
 // Close still closes the connection.
 // A nil error restores the usual behavior.
 func (c *Conn) SetCloseError(err error) {
 	c.r.lock()
 	c.r.closeErr = err
 	c.r.unlock()
 }

 // LocalAddr returns the (fake) local network address.
 func (c *Conn) LocalAddr() net.Addr {
 	c.r.lock()
 	defer c.r.unlock()
 	return c.r.addr
 }

 // SetLocalAddr sets the local address.
 //
 // To set the remote address, set the local address of Conn's peer.
 func (c *Conn) SetLocalAddr(addr net.Addr) {
 	c.r.lock()
 	defer c.r.unlock()
 	c.r.addr = addr
 }

 // LocalAddr returns the (fake) remote network address.
 func (c *Conn) RemoteAddr() net.Addr {
 	c.r.lock()
 	defer c.r.unlock()
 	return c.w.addr
 }

 // SetDeadline sets the read and write deadlines for the connection.
 func (c *Conn) SetDeadline(t time.Time) error {
 	c.SetReadDeadline(t)
 	c.SetWriteDeadline(t)
 	return nil
 }

 // SetReadDeadline sets the read deadline for the connection.
 func (c *Conn) SetReadDeadline(t time.Time) error {
 	c.r.readDeadline.setDeadline(c.r, t)
 	return nil
 }

 // SetWriteDeadline sets the write deadline for the connection.
 func (c *Conn) SetWriteDeadline(t time.Time) error {
 	c.w.writeDeadline.setDeadline(c.w, t)
 	return nil
 }

 // SetReadBufferSize sets the connection's read buffer.
 // Writes to the other end of the connection will block so long as the buffer is full.
 // Setting the size to 0 blocks all writes until the size is increased.
 func (c *Conn) SetReadBufferSize(size int) {
 	if size < 0 {
 		size = math.MaxInt
 	}
 	c.r.setBufferSize(size)
 }

 // SetReadError causes any currently blocked and future Read calls to return
 // a net.OpError wrapping err. It does not affect the other half of the connection.
 // Reads will return any buffered data before returning the error,
 // including data written after the error is set and io.EOF after the other end is closed.
 // A nil error restores the usual behavior.
 func (c *Conn) SetReadError(err error) {
 	c.r.lock()
 	defer c.r.unlock()
 	c.r.readErr = err
 }

 // SetWriteError causes any currently blocked and future Write calls to return
 // a net.OpError wrapping err. It does not affect the other half of the connection.
 // Writes will not write data to the connection buffer while an error is set.
 // A nil error restores the usual behavior.
 func (c *Conn) SetWriteError(err error) {
 	c.w.lock()
 	defer c.w.unlock()
 	c.w.writeErr = err
 }

 // connHalf is one direction data flow in a Conn.
 // The connHalf contains a buffer.
 // Writes to the connHalf push to the buffer and reads pull from it.
 type connHalf struct {
 	addr net.Addr

 	// A half can be readable and/or writable.
 	//
 	// These four channels act as a lock,
 	// and allow waiting for readability/writability.
 	// When the half is unlocked, exactly one channel contains a value.
 	// When the half is locked, all channels are empty.
 	lockr  chan struct{} // readable
 	lockw  chan struct{} // writable
 	lockrw chan struct{} // readable and writable
 	lockc  chan struct{} // neither readable nor writable

 	// Read and write timeouts.
 	readDeadline, writeDeadline connDeadline

 	bufMax int // maximum buffer size
 	buf    bytes.Buffer

 	readClosed, writeClosed bool
 	readErr, writeErr       error // errors returned by reads/writes
 	closeErr                error // error returned by closing the conn reading from this half
 }

 func newConnHalf(addr net.Addr) *connHalf {
 	h := &connHalf{
 		addr:   addr,
 		lockw:  make(chan struct{}, 1),
 		lockr:  make(chan struct{}, 1),
 		lockrw: make(chan struct{}, 1),
 		lockc:  make(chan struct{}, 1),
 		bufMax: math.MaxInt, // unlimited
 	}
 	h.unlock()
 	return h
 }

 // lock locks h.
 func (h *connHalf) lock() {
 	select {
 	case <-h.lockw: // writable
 	case <-h.lockr: // readable
 	case <-h.lockrw: // readable and writable
 	case <-h.lockc: // neither readable nor writable
 	}
 }

 // unlock unlocks h.
 func (h *connHalf) unlock() {
 	canRead := h.canReadLocked()
 	canWrite := h.canWriteLocked()
 	switch {
 	case canRead && canWrite:
 		h.lockrw <- struct{}{} // readable and writable
 	case canRead:
 		h.lockr <- struct{}{} // readable
 	case canWrite:
 		h.lockw <- struct{}{} // writable
 	default:
 		h.lockc <- struct{}{} // neither readable nor writable
 	}
 }

 func (h *connHalf) canRead() bool {
 	h.lock()
 	defer h.unlock()
 	return h.canReadLocked()
 }

 func (h *connHalf) canReadLocked() bool {
 	return h.readErr != nil || h.readDeadline.expired || h.buf.Len() > 0 || h.readClosed || h.writeClosed
 }

 func (h *connHalf) canWriteLocked() bool {
 	return h.writeErr != nil || h.writeDeadline.expired || h.bufMax > h.buf.Len() || h.readClosed || h.writeClosed
 }

 // waitAndLockForRead waits until h is readable and locks it.
 func (h *connHalf) waitAndLockForRead() {
 	select {
 	case <-h.lockr:
 		// readable
 	case <-h.lockrw:
 		// readable and writable
 	}
 }

 // waitAndLockForWrite waits until h is writable and locks it.
 func (h *connHalf) waitAndLockForWrite() {
 	select {
 	case <-h.lockw:
 		// writable
 	case <-h.lockrw:
 		// readable and writable
 	}
 }

 func (h *connHalf) read(b []byte) (n int, err error) {
 	h.waitAndLockForRead()
 	defer h.unlock()
 	if h.readClosed {
 		return 0, net.ErrClosed
 	}
 	if h.readDeadline.expired {
 		return 0, os.ErrDeadlineExceeded
 	}
 	if h.buf.Len() > 0 {
 		return h.buf.Read(b)
 	}
 	if h.writeClosed {
 		return 0, io.EOF
 	}
 	return 0, h.readErr
 }

 func (h *connHalf) setBufferSize(size int) {
 	h.lock()
 	defer h.unlock()
 	h.bufMax = size
 }

 func (h *connHalf) write(b []byte) (n int, err error) {
 	for n < len(b) {
 		nn, err := h.writePartial(b[n:])
 		n += nn
 		if err != nil {
 			return n, err
 		}
 	}
 	return n, nil
 }

 func (h *connHalf) writePartial(b []byte) (n int, err error) {
 	h.waitAndLockForWrite()
 	defer h.unlock()
 	if h.writeClosed {
 		return 0, net.ErrClosed
 	}
 	if h.writeDeadline.expired {
 		return 0, os.ErrDeadlineExceeded
 	}
 	if h.readClosed {
 		return 0, errClosedByPeer
 	}
 	if h.writeErr != nil {
 		return 0, h.writeErr
 	}
 	writeMax := h.bufMax - h.buf.Len()
 	if writeMax < len(b) {
 		b = b[:writeMax]
 	}
 	return h.buf.Write(b)
 }

 type connDeadline struct {
 	timer   *time.Timer
 	expired bool
 }

 type locker interface {
 	lock()
 	unlock()
 }

 func (d *connDeadline) setDeadline(mu locker, t time.Time) {
 	mu.lock()
 	defer mu.unlock()
 	if d.timer != nil {
 		d.timer.Stop()
 		d.timer = nil
 	}
 	if t.IsZero() {
 		// No deadline.
 		d.expired = false
 		return
 	}
 	expiry := time.Until(t)
 	if expiry <= 0 {
 		// Deadline has already passed.
 		d.expired = true
 		return
 	}
 	// Deadline is in the future.
 	d.expired = false
 	var timer *time.Timer
 	timer = time.AfterFunc(expiry, func() {
 		mu.lock()
 		defer mu.unlock()
 		if d.timer == timer {
 			d.timer = nil
 			d.expired = true
 		}
 	})
 	d.timer = timer
 }
	// Copyright 2026 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package nettest

	import (
	"bytes"
	"errors"
	"io"
	"math"
	"net"
	"net/netip"
	"os"
	"time"
	)

	// Conn is an in-memory test implementation of net.Conn.
	type Conn struct {
	// Conns come in pairs.
	// Writes to one Conn are read by its peer, and vice-versa.
	//
	// A connHalf handles one direction of data flow.
	// A Conn consists of read and write halves.
	// A Conn's peer has the same halves, only swapped.
	//
	// A Conn reads from r and writes to w.
	r, w *connHalf

	// peer is the other endpoint.
	peer *Conn
	}

	// NewConnPair returns a pair of connected Conns.
	func NewConnPair() (Conn, Conn) {
	return newConnPair(
	net.TCPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:10000")),
	net.TCPAddrFromAddrPort(netip.MustParseAddrPort("127.0.0.1:10001")),
	)
	}

	func newConnPair(addr1, addr2 net.Addr) (Conn, Conn) {
	h1 := newConnHalf(addr1)
	h2 := newConnHalf(addr2)
	c1 := &Conn{r: h1, w: h2}
	c2 := &Conn{r: h2, w: h1}
	c1.peer = c2
	c2.peer = c1
	c1.SetReadBufferSize(-1)
	c2.SetReadBufferSize(-1)
	return c1, c2
	}

	// Peer returns the other end of the connection.
	func (c Conn) Peer() Conn {
	return c.peer
	}

	// Read reads data from the connection.
	func (c *Conn) Read(b []byte) (n int, err error) {
	n, err = c.r.read(b)
	if err != nil && err != io.EOF {
	err = &net.OpError{
	Op: "read",
	Net: "tcp",
	Source: c.RemoteAddr(),
	Addr: c.LocalAddr(),
	Err: err,
	}
	}
	return n, err
	}

	// CanRead reports whether Read can proceed without blocking.
	func (c *Conn) CanRead() bool {
	return c.r.canRead()
	}

	// Write writes data to the connection.
	func (c *Conn) Write(b []byte) (n int, err error) {
	n, err = c.w.write(b)
	if err != nil {
	err = &net.OpError{
	Op: "write",
	Net: "tcp",
	Source: c.LocalAddr(),
	Addr: c.RemoteAddr(),
	Err: err,
	}
	}
	return n, err
	}

	// IsClosed reports whether the connection has been closed.
	// A connection is closed if [CloseRead] and [CloseWrite] are both called,
	// or if [Close] is called.
	//
	// To identify when the other side of the Conn has been closed,
	// use Conn.Peer().IsClosed().
	func (c *Conn) IsClosed() bool {
	c.r.lock()
	readClosed := c.r.readClosed
	c.r.unlock()
	c.w.lock()
	writeClosed := c.w.writeClosed
	c.w.unlock()
	return readClosed && writeClosed
	}

	var errClosedByPeer = errors.New("connection closed by peer")

	// CloseRead shuts down the reading side of the connection.
	func (c *Conn) CloseRead() error {
	c.r.lock()
	defer c.r.unlock()
	c.r.buf.Reset() // discard unread data
	c.r.readClosed = true
	return nil
	}

	// CloseWrite shuts down the writing side of the connection.
	func (c *Conn) CloseWrite() error {
	c.w.lock()
	defer c.w.unlock()
	c.w.writeClosed = true
	return nil
	}

	// Close closes the connection.
	func (c *Conn) Close() error {
	c.r.lock()
	readClosed := c.r.readClosed
	c.r.buf.Reset() // discard unread data
	c.r.readClosed = true
	err := c.r.closeErr
	c.r.unlock()

	c.w.lock()
	writeClosed := c.w.writeClosed
	c.w.writeClosed = true
	c.w.unlock()

	if readClosed && writeClosed {
	err = net.ErrClosed
	}
	if err != nil {
	err = &net.OpError{
	Op: "close",
	Net: "tcp",
	Addr: c.LocalAddr(),
	Err: err,
	}
	}
	return err
	}

	// SetCloseError sets the error returned by Close.
	// Close still closes the connection.
	// A nil error restores the usual behavior.
	func (c *Conn) SetCloseError(err error) {
	c.r.lock()
	c.r.closeErr = err
	c.r.unlock()
	}

	// LocalAddr returns the (fake) local network address.
	func (c *Conn) LocalAddr() net.Addr {
	c.r.lock()
	defer c.r.unlock()
	return c.r.addr
	}

	// SetLocalAddr sets the local address.
	//
	// To set the remote address, set the local address of Conn's peer.
	func (c *Conn) SetLocalAddr(addr net.Addr) {
	c.r.lock()
	defer c.r.unlock()
	c.r.addr = addr
	}

	// LocalAddr returns the (fake) remote network address.
	func (c *Conn) RemoteAddr() net.Addr {
	c.r.lock()
	defer c.r.unlock()
	return c.w.addr
	}

	// SetDeadline sets the read and write deadlines for the connection.
	func (c *Conn) SetDeadline(t time.Time) error {
	c.SetReadDeadline(t)
	c.SetWriteDeadline(t)
	return nil
	}

	// SetReadDeadline sets the read deadline for the connection.
	func (c *Conn) SetReadDeadline(t time.Time) error {
	c.r.readDeadline.setDeadline(c.r, t)
	return nil
	}

	// SetWriteDeadline sets the write deadline for the connection.
	func (c *Conn) SetWriteDeadline(t time.Time) error {
	c.w.writeDeadline.setDeadline(c.w, t)
	return nil
	}

	// SetReadBufferSize sets the connection's read buffer.
	// Writes to the other end of the connection will block so long as the buffer is full.
	// Setting the size to 0 blocks all writes until the size is increased.
	func (c *Conn) SetReadBufferSize(size int) {
	if size < 0 {
	size = math.MaxInt
	}
	c.r.setBufferSize(size)
	}

	// SetReadError causes any currently blocked and future Read calls to return
	// a net.OpError wrapping err. It does not affect the other half of the connection.
	// Reads will return any buffered data before returning the error,
	// including data written after the error is set and io.EOF after the other end is closed.
	// A nil error restores the usual behavior.
	func (c *Conn) SetReadError(err error) {
	c.r.lock()
	defer c.r.unlock()
	c.r.readErr = err
	}

	// SetWriteError causes any currently blocked and future Write calls to return
	// a net.OpError wrapping err. It does not affect the other half of the connection.
	// Writes will not write data to the connection buffer while an error is set.
	// A nil error restores the usual behavior.
	func (c *Conn) SetWriteError(err error) {
	c.w.lock()
	defer c.w.unlock()
	c.w.writeErr = err
	}

	// connHalf is one direction data flow in a Conn.
	// The connHalf contains a buffer.
	// Writes to the connHalf push to the buffer and reads pull from it.
	type connHalf struct {
	addr net.Addr

	// A half can be readable and/or writable.
	//
	// These four channels act as a lock,
	// and allow waiting for readability/writability.
	// When the half is unlocked, exactly one channel contains a value.
	// When the half is locked, all channels are empty.
	lockr chan struct{} // readable
	lockw chan struct{} // writable
	lockrw chan struct{} // readable and writable
	lockc chan struct{} // neither readable nor writable

	// Read and write timeouts.
	readDeadline, writeDeadline connDeadline

	bufMax int // maximum buffer size
	buf bytes.Buffer

	readClosed, writeClosed bool
	readErr, writeErr error // errors returned by reads/writes
	closeErr error // error returned by closing the conn reading from this half
	}

	func newConnHalf(addr net.Addr) *connHalf {
	h := &connHalf{
	addr: addr,
	lockw: make(chan struct{}, 1),
	lockr: make(chan struct{}, 1),
	lockrw: make(chan struct{}, 1),
	lockc: make(chan struct{}, 1),
	bufMax: math.MaxInt, // unlimited
	}
	h.unlock()
	return h
	}

	// lock locks h.
	func (h *connHalf) lock() {
	select {
	case <-h.lockw: // writable
	case <-h.lockr: // readable
	case <-h.lockrw: // readable and writable
	case <-h.lockc: // neither readable nor writable
	}
	}

	// unlock unlocks h.
	func (h *connHalf) unlock() {
	canRead := h.canReadLocked()
	canWrite := h.canWriteLocked()
	switch {
	case canRead && canWrite:
	h.lockrw <- struct{}{} // readable and writable
	case canRead:
	h.lockr <- struct{}{} // readable
	case canWrite:
	h.lockw <- struct{}{} // writable
	default:
	h.lockc <- struct{}{} // neither readable nor writable
	}
	}

	func (h *connHalf) canRead() bool {
	h.lock()
	defer h.unlock()
	return h.canReadLocked()
	}

	func (h *connHalf) canReadLocked() bool {
	return h.readErr != nil \|\| h.readDeadline.expired \|\| h.buf.Len() > 0 \|\| h.readClosed \|\| h.writeClosed
	}

	func (h *connHalf) canWriteLocked() bool {
	return h.writeErr != nil \|\| h.writeDeadline.expired \|\| h.bufMax > h.buf.Len() \|\| h.readClosed \|\| h.writeClosed
	}

	// waitAndLockForRead waits until h is readable and locks it.
	func (h *connHalf) waitAndLockForRead() {
	select {
	case <-h.lockr:
	// readable
	case <-h.lockrw:
	// readable and writable
	}
	}

	// waitAndLockForWrite waits until h is writable and locks it.
	func (h *connHalf) waitAndLockForWrite() {
	select {
	case <-h.lockw:
	// writable
	case <-h.lockrw:
	// readable and writable
	}
	}

	func (h *connHalf) read(b []byte) (n int, err error) {
	h.waitAndLockForRead()
	defer h.unlock()
	if h.readClosed {
	return 0, net.ErrClosed
	}
	if h.readDeadline.expired {
	return 0, os.ErrDeadlineExceeded
	}
	if h.buf.Len() > 0 {
	return h.buf.Read(b)
	}
	if h.writeClosed {
	return 0, io.EOF
	}
	return 0, h.readErr
	}

	func (h *connHalf) setBufferSize(size int) {
	h.lock()
	defer h.unlock()
	h.bufMax = size
	}

	func (h *connHalf) write(b []byte) (n int, err error) {
	for n < len(b) {
	nn, err := h.writePartial(b[n:])
	n += nn
	if err != nil {
	return n, err
	}
	}
	return n, nil
	}

	func (h *connHalf) writePartial(b []byte) (n int, err error) {
	h.waitAndLockForWrite()
	defer h.unlock()
	if h.writeClosed {
	return 0, net.ErrClosed
	}
	if h.writeDeadline.expired {
	return 0, os.ErrDeadlineExceeded
	}
	if h.readClosed {
	return 0, errClosedByPeer
	}
	if h.writeErr != nil {
	return 0, h.writeErr
	}
	writeMax := h.bufMax - h.buf.Len()
	if writeMax < len(b) {
	b = b[:writeMax]
	}
	return h.buf.Write(b)
	}

	type connDeadline struct {
	timer *time.Timer
	expired bool
	}

	type locker interface {
	lock()
	unlock()
	}

	func (d *connDeadline) setDeadline(mu locker, t time.Time) {
	mu.lock()
	defer mu.unlock()
	if d.timer != nil {
	d.timer.Stop()
	d.timer = nil
	}
	if t.IsZero() {
	// No deadline.
	d.expired = false
	return
	}
	expiry := time.Until(t)
	if expiry <= 0 {
	// Deadline has already passed.
	d.expired = true
	return
	}
	// Deadline is in the future.
	d.expired = false
	var timer *time.Timer
	timer = time.AfterFunc(expiry, func() {
	mu.lock()
	defer mu.unlock()
	if d.timer == timer {
	d.timer = nil
	d.expired = true
	}
	})
	d.timer = timer
	}