lneto/tcp/conn.go

package tcp

import (
	"errors"
	"io"
	"log/slog"
	"net"
	"net/netip"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/soypat/lneto"
	"github.com/soypat/lneto/internal"
)

var (
	errDeadlineExceeded = os.ErrDeadlineExceeded
	errNoRemoteAddr     = errors.New("tcp: no remote address established")
)

// Conn builds on the [Handler] abstraction and adds IP header knowledge, time management, and familiar user facing API
// like Write and Read methods.
//
// Note that the complete emulation of [net.TCPConn] at this level of abstraction is yet a non-goal,
// even though the functionality provided is similar.
type Conn struct {
	mu         sync.Mutex
	h          Handler
	remoteAddr []byte

	// writeLock holds the connection ID ([Handler.connid]) of the goroutine
	// that currently owns the write side via [Conn.Write] or [Conn.Flush].
	// A zero value means no write is in progress. It serializes concurrent
	// writers and lets [Conn.Close] acquire the write side before closing so it
	// does not drop in-flight data (see issue #82).
	writeLock atomic.Uint64

	_backoff lneto.BackoffStrategy
	rdead    time.Time
	wdead    time.Time
	abortErr error
	logger

	ipID uint16
}

// reset must be called while holding [Conn.mu].
func (conn *Conn) reset(h Handler) {
	// Reset fields individually - DO NOT copy the mutex (undefined behavior in Go).
	// "A Mutex must not be copied after first use." - sync package docs.
	// Copying a locked mutex causes corruption on multi-core systems.
	conn.h = h
	conn.remoteAddr = conn.remoteAddr[:0]
	conn.rdead = time.Time{}
	conn.wdead = time.Time{}
	conn.abortErr = nil
	conn.ipID = 0
	conn.writeLock.Store(0)
}

// ConnConfig provides configuration parameters for [Conn].
type ConnConfig struct {
	RxBuf []byte // Fixed size buffer for incoming data via [Conn.Read].
	TxBuf []byte // Fixed size buffer for egress data via [Conn.Write].
	// TxPacketQueueSize is the maximum number of sent-but-unacknowledged segments
	// the connection tracks at once for retransmission. Each queued entry records the
	// sequence range of one outgoing segment and is released when the peer ACKs it;
	// once the queue is full no further segments are emitted until an ACK frees a slot.
	// Must be greater than zero and no larger than len(TxBuf).
	TxPacketQueueSize int
	// RWBackoff sets the backoff policy for backoff when data unavailable on Read or buffer full on Write.
	// If not set a default backoff strategy will be used. See [internal.BackoffConnRW].
	RWBackoff lneto.BackoffStrategy
	// Logger sets the [Conn] logger.
	// Lower level logging available at [Handler.SetLoggers] via [Conn.InternalHandler].
	Logger *slog.Logger
}

// Configure should be called on any newly created connection before usage. See [ConnConfig].
func (conn *Conn) Configure(config ConnConfig) (err error) {
	if config.RWBackoff == nil {
		return lneto.ErrMissingHALConfig
	}
	conn.mu.Lock()
	defer conn.mu.Unlock()
	err = conn.h.SetBuffers(config.TxBuf, config.RxBuf, config.TxPacketQueueSize)
	if err != nil {
		return err
	}
	conn._backoff = config.RWBackoff
	conn.logger.log = config.Logger
	return nil
}

// LocalPort returns the local port on which the socket is listening or connected to.
func (conn *Conn) LocalPort() uint16 {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.LocalPort()
}

// RemotePort returns the port of the incoming remote connection. Is non-zero if connection is established.
func (conn *Conn) RemotePort() uint16 {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.RemotePort()
}

// RemoteAddr returns the address of the peer Conn is exchanging data with.
func (conn *Conn) RemoteAddr() []byte {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.remoteAddr
}

// State returns the TCP state of the socket.
func (conn *Conn) State() State {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.State()
}

// AwaitingSynSend reports whether the connection has been opened actively but has
// not yet emitted its SYN. It locks the connection so it is safe to poll from a
// goroutine other than the one driving the stack (unlike reaching through
// [Conn.InternalHandler]).
func (conn *Conn) AwaitingSynSend() bool {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.AwaitingSynSend()
}

// BufferedInput returns the number of bytes in the socket's receive(input) buffer
// and available to read via a [Conn.Read] call.
func (conn *Conn) BufferedInput() int {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.BufferedInput()
}

// BufferedUnsent returns the number of bytes in the socket's transmit(output) buffer
// that has yet to be sent.
func (conn *Conn) BufferedUnsent() int {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.BufferedUnsent()
}

// FreeInput returns the number of free bytes in the socket's receive(input) buffer.
// The TCP window mechanism advertises this space to the remote peer,
// preventing the sender from transmitting more data than the buffer can hold.
func (conn *Conn) FreeInput() int {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.FreeInput()
}

// FreeOutput returns the number of free bytes in the socket's transmit(output) buffer.
// This is the amount of data that can be written via [Conn.Write] before it blocks.
func (conn *Conn) FreeOutput() int {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.h.FreeOutput()
}

// OpenActive opens a connection to a remote peer with a known IP address and port combination.
// iss is the initial send sequence number which is ideally a random number which is far away from the last sequence number used on a connection to the same host.
func (conn *Conn) OpenActive(localPort uint16, remote netip.AddrPort, iss Value) error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if !remote.IsValid() {
		return lneto.ErrInvalidAddr
	}
	rport := remote.Port()
	err := conn.h.OpenActive(localPort, rport, iss)
	if err != nil {
		return err
	}
	conn.reset(conn.h)
	raddr := remote.Addr()
	if raddr.Is4() {
		addr4 := raddr.As4()
		conn.remoteAddr = append(conn.remoteAddr[:0], addr4[:]...)
	} else if raddr.Is6() {
		addr6 := raddr.As16()
		conn.remoteAddr = append(conn.remoteAddr[:0], addr6[:]...)
	}
	conn.debug("conn:dial", slog.Uint64("lport", uint64(localPort)), slog.Uint64("rport", uint64(rport)))
	return nil
}

// OpenListen opens a passive connection which listens for the first SYN packet to be received on a local port.
// iss is the initial send sequence number which is usually a randomly chosen number.
func (conn *Conn) OpenListen(localPort uint16, iss Value) error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	err := conn.h.OpenListen(localPort, iss)
	if err != nil {
		return err
	}
	conn.reset(conn.h)
	conn.debug("conn:listen", slog.Uint64("lport", uint64(localPort)))
	return nil
}

// CloseRead activates local discard mode on the connection. Incoming data is
// still ACKed normally but payload is dropped; future Read calls return io.EOF.
// The write side is unaffected.
// If [Conn.Close] is also called the connection is terminated.
func (conn *Conn) CloseRead() error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	err := conn.checkPipeOpen()
	if err != nil {
		return err
	}
	conn.h.ShutdownRead()
	return nil
}

// Close will initiate TCP close sequence. After Close is called future [Conn.Write] calls will fail with [net.ErrClosed].
func (conn *Conn) Close() error {
	connid, err := conn.acquireWriteLock(nil)
	if err != nil {
		return err
	}
	defer conn.releaseWriteLock(connid)
	conn.mu.Lock()
	defer conn.mu.Unlock()
	conn.trace("TCPConn.Close", slog.Uint64("lport", uint64(conn.h.localPort)), slog.Uint64("rport", uint64(conn.h.remotePort)))
	return conn.h.Close()
}

// Abort terminates all state of the connection forcibly.
func (conn *Conn) Abort() {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	conn.trace("TCPConn.Abort", slog.Uint64("lport", uint64(conn.h.localPort)), slog.Uint64("rport", uint64(conn.h.remotePort)))
	conn.h.Abort()
	conn.reset(conn.h)
}

// InternalHandler returns the internal [Handler] instance. The Handler contains lower level implementation logic for a TCP connection.
// Typical users should not be using this method unless implementing a stack which manages several TCP connections and thus need
// access to low level internals for careful memory management.
func (conn *Conn) InternalHandler() *Handler {
	return &conn.h
}

// Write writes argument data to the TCPConns's output buffer which is queued to be sent.
func (conn *Conn) Write(b []byte) (int, error) {
	connid, err := conn.acquireWriteLock(&conn.wdead)
	if err != nil {
		return 0, err
	}
	defer conn.releaseWriteLock(connid)
	rport := conn.RemotePort()
	plen := len(b)
	lport := conn.LocalPort()
	conn.trace("TCPConn.Write:start", slog.Uint64("lport", uint64(lport)), slog.Uint64("rport", uint64(rport)))
	if conn.deadlineExceeded(&conn.wdead) {
		return 0, errDeadlineExceeded
	} else if plen == 0 {
		return 0, nil
	}
	n := 0
	var backoffs uint
	for len(b) > 0 {
		if err := conn.checkPipe(connid, &conn.wdead); err != nil {
			return n, err
		}
		conn.mu.Lock()
		towrite := min(conn.h.FreeOutput(), len(b))
		var ngot int
		if towrite > 0 {
			ngot, err = conn.h.Write(b[:towrite])
			conn.mu.Unlock()
			if err != nil && err != internal.ErrRingBufferFull {
				break
			} else if ngot != towrite {
				panic("unreachable")
			}
			n += ngot
			b = b[ngot:]
			backoffs = 0
		} else {
			// No data can be written.
			conn.mu.Unlock()
			conn.trace("TCPConn.Write:insuf-buf", slog.Int("missing", plen-n), slog.Uint64("lport", uint64(lport)), slog.Uint64("rport", uint64(rport)))
			if conn.deadlineExceeded(&conn.wdead) {
				return n, errDeadlineExceeded
			}
			conn.backoff(backoffs)
			backoffs++
		}
	}
	return n, err
}

// Flush blocks until all buffered TCP data has been sent.
func (conn *Conn) Flush() error {
	connid, err := conn.acquireWriteLock(&conn.wdead)
	if err != nil {
		return err
	}
	defer conn.releaseWriteLock(connid)
	var backoffs uint
	for conn.BufferedUnsent() != 0 {
		if err := conn.checkPipe(connid, &conn.wdead); err != nil {
			return err
		}
		conn.backoff(backoffs)
		backoffs++
	}
	return nil
}

// Read reads data from the socket's input buffer. If the buffer is empty,
// Read will block until data is available or connection closes.
// Returns io.EOF when the remote has closed the connection and all buffered data has been read.
func (conn *Conn) Read(b []byte) (int, error) {
	conn.mu.Lock()
	connID := conn.h.connid
	lport := conn.h.localPort
	rport := conn.h.remotePort
	conn.mu.Unlock()
	conn.trace("TCPConn.Read:start", slog.Uint64("lport", uint64(lport)), slog.Uint64("rport", uint64(rport)))
	var backoffs uint
	n := 0
	for len(b) > 0 {
		conn.mu.Lock()
		if connID != conn.h.connid {
			conn.mu.Unlock()
			return n, net.ErrClosed
		}
		avail := conn.h.BufferedInput()
		if avail > 0 {
			// Read branch.
			ngot, err := conn.h.Read(b)
			conn.mu.Unlock()
			n += ngot
			if err != nil {
				return n, err
			}
			b = b[ngot:]
		} else if n > 0 {
			conn.mu.Unlock()
			break
		} else {
			state := conn.h.State()
			rxRefuse := conn.h.shutdownRx
			conn.mu.Unlock()
			if state.IsClosed() {
				return n, net.ErrClosed
			} else if !state.RxDataOpen() || rxRefuse {
				return n, io.EOF
			} else if conn.deadlineExceeded(&conn.rdead) {
				return n, errDeadlineExceeded
			}
			conn.backoff(backoffs)
			backoffs++
		}
	}
	return n, nil
}

// acquireWriteLock validates the connection is open and acquires the write lock
// for the current connection, returning its connection ID. It serializes
// concurrent writers: while another goroutine owns the write side it blocks with
// backoff until that writer releases, the connection closes, or deadline elapses.
// The returned connID must be released with [Conn.releaseWriteLock].
func (conn *Conn) acquireWriteLock(deadline *time.Time) (connID uint64, err error) {
	conn.mu.Lock()
	connID = conn.h.connid
	conn.mu.Unlock()
	var backoffs uint
	for {
		if err := conn.checkPipe(connID, deadline); err != nil {
			return 0, err
		}
		if conn.writeLock.CompareAndSwap(0, connID) {
			return connID, nil
		}
		conn.backoff(backoffs)
		backoffs++
	}
}

// releaseWriteLock releases a write lock previously acquired by [Conn.acquireWriteLock].
func (conn *Conn) releaseWriteLock(connID uint64) {
	conn.writeLock.CompareAndSwap(connID, 0)
}

func (conn *Conn) checkPipe(connID uint64, deadline *time.Time) (err error) {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if conn.abortErr != nil {
		err = conn.abortErr
	} else if connID != conn.h.connid || conn.h.State().IsClosed() {
		err = net.ErrClosed
	} else if deadline != nil && !deadline.IsZero() && time.Since(*deadline) > 0 {
		err = errDeadlineExceeded
	}
	return err
}

func (conn *Conn) checkPipeOpen() error {
	if conn.abortErr != nil {
		return conn.abortErr
	}
	state := conn.h.State()
	if state.IsClosed() {
		return net.ErrClosed
	}
	return nil
}

// Demux implements [lneto.StackNode].
func (conn *Conn) Demux(buf []byte, off int) (err error) {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if off >= len(buf) {
		// off is the IP header length; if it equals or exceeds the frame
		// length, there are zero bytes of TCP payload — drop the frame.
		return lneto.ErrTruncatedFrame
	}
	raddr, _, id, _, err := internal.GetIPAddr(buf[:off])
	if err != nil {
		return err
	}
	if conn.isRaddrSet() && !internal.BytesEqual(conn.remoteAddr, raddr) {
		return lneto.ErrMismatch
	}
	conn.trace("tcpconn.Recv", slog.Uint64("lport", uint64(conn.h.LocalPort())), slog.Uint64("rport", uint64(conn.h.remotePort)))
	err = conn.h.Recv(buf[off:])
	if err != nil {
		return err
	}
	if !conn.isRaddrSet() && conn.h.RemotePort() != 0 {
		conn.remoteAddr = append(conn.remoteAddr[:0], raddr...)
		conn.ipID = ^(id - 1)
	}
	return nil
}

// Encapsulate implements [lneto.StackNode].
func (conn *Conn) Encapsulate(carrierData []byte, offsetToIP, offsetToFrame int) (n int, err error) {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	if len(conn.remoteAddr) == 0 {
		return 0, errNoRemoteAddr
	}
	if offsetToIP < 0 {
		return 0, errNoRemoteAddr // No IP layer present.
	}
	ipFrame := carrierData[offsetToIP:offsetToFrame]
	raddr, _, _, _, err := internal.GetIPAddr(ipFrame)
	if err != nil {
		return 0, err
	} else if len(raddr) != len(conn.remoteAddr) {
		return 0, lneto.ErrMismatchLen
	}
	n, err = conn.h.Send(carrierData[offsetToFrame:])
	if err != nil || n == 0 {
		return 0, err
	}
	conn.trace("TCPConn.encaps", slog.Uint64("lport", uint64(conn.h.LocalPort())), slog.Uint64("rport", uint64(conn.h.remotePort)))
	err = internal.SetIPAddrs(ipFrame, conn.ipID, nil, conn.remoteAddr)
	if err != nil {
		return 0, err
	}
	conn.ipID++
	return n, nil
}

// Protocol implements [lneto.StackNode].
func (conn *Conn) Protocol() uint64 {
	return uint64(lneto.IPProtoTCP)
}

func (conn *Conn) isRaddrSet() bool {
	return len(conn.remoteAddr) != 0
}

// SetDeadline sets the read and write deadlines associated
// with the connection. It is equivalent to calling both
// SetReadDeadline and SetWriteDeadline. Implements [net.Conn].
func (conn *Conn) SetDeadline(t time.Time) error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	err := conn.setReadDeadline(t)
	if err != nil {
		return err
	}
	return conn.setWriteDeadline(t)
}

// SetReadDeadline sets the deadline for future Read calls
// and any currently-blocked Read call. A zero value for t means Read will not time out.
func (conn *Conn) SetReadDeadline(t time.Time) error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.setReadDeadline(t)
}

func (conn *Conn) setReadDeadline(t time.Time) error {
	conn.trace("TCPConn.setReadDeadline:start")
	err := conn.checkPipeOpen()
	if err == nil {
		conn.rdead = t
	}
	return err
}

// SetWriteDeadline sets the deadline for future Write calls
// and any currently-blocked Write call.
// Even if write times out, it may return n > 0, indicating that
// some of the data was successfully written.
// A zero value for t means Write will not time out.
func (conn *Conn) SetWriteDeadline(t time.Time) error {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return conn.setWriteDeadline(t)
}

func (conn *Conn) setWriteDeadline(t time.Time) error {
	conn.trace("TCPConn.SetWriteDeadline:start")
	err := conn.checkPipeOpen()
	if err == nil {
		conn.wdead = t
	}
	return err
}

func (conn *Conn) deadlineExceeded(deadline *time.Time) bool {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	return !deadline.IsZero() && time.Since(*deadline) > 0
}

// ConnectionID implements [lneto.StackNode].
func (conn *Conn) ConnectionID() *uint64 {
	return conn.h.ConnectionID()
}

func (conn *Conn) backoff(consecutiveBackoffs uint) {
	conn._backoff.Do(consecutiveBackoffs)
}