api/zap/transport.go

// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package zap

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var (
	ErrClosed         = errors.New("zap: connection closed")
	ErrTimeout        = errors.New("zap: request timeout")
	ErrResponseFailed = errors.New("zap: response failed")
)

// Config contains transport configuration
type Config struct {
	// ReadTimeout is the timeout for reading a message
	ReadTimeout time.Duration
	// WriteTimeout is the timeout for writing a message
	WriteTimeout time.Duration
	// MaxConcurrent is the maximum number of concurrent requests
	MaxConcurrent int
	// BufferSize is the read/write buffer size
	BufferSize int
}

// DefaultConfig returns a Config with reasonable defaults
func DefaultConfig() *Config {
	return &Config{
		ReadTimeout:   30 * time.Second,
		WriteTimeout:  10 * time.Second,
		MaxConcurrent: 1000,
		BufferSize:    64 * 1024,
	}
}

// Conn is a ZAP connection that multiplexes requests
type Conn struct {
	conn   net.Conn
	config *Config

	reader *bufio.Reader
	writer *bufio.Writer

	mu       sync.Mutex
	requests map[uint32]chan *response
	nextID   uint32

	closed atomic.Bool
	done   chan struct{}
}

type response struct {
	msgType MessageType
	payload []byte
	err     error
}

// Dial connects to a ZAP server. The network transport is inferred from addr:
// a filesystem path ("/tmp/vm.sock", "@abstract", or anything containing no
// ':' host:port) dials a unix-domain socket — no TCP/UDP stack, the in-proc /
// same-host fast path; a "host:port" addr dials TCP for a genuinely remote
// server. Backward compatible: existing "127.0.0.1:port" addrs still dial tcp.
func Dial(ctx context.Context, addr string, config *Config) (*Conn, error) {
	if config == nil {
		config = DefaultConfig()
	}

	dialer := &net.Dialer{}
	conn, err := dialer.DialContext(ctx, transportNetwork(addr), addr)
	if err != nil {
		return nil, fmt.Errorf("zap: dial failed: %w", err)
	}

	return newConn(conn, config), nil
}

// transportNetwork infers the net.Dial/net.Listen network from an address:
// unix-domain socket for a filesystem/abstract path (same-host, no TCP stack),
// tcp for a host:port. A ':' with no '/' means host:port ⇒ tcp; anything that
// looks like a path ('/', leading '@' abstract socket, or no ':') ⇒ unix.
func transportNetwork(addr string) string {
	if strings.HasPrefix(addr, "@") || strings.HasPrefix(addr, "/") {
		return "unix"
	}
	if strings.Contains(addr, ":") && !strings.Contains(addr, "/") {
		return "tcp"
	}
	if strings.Contains(addr, "/") {
		return "unix"
	}
	return "tcp"
}

// NewConn wraps an existing net.Conn as a ZAP connection
func NewConn(conn net.Conn, config *Config) *Conn {
	if config == nil {
		config = DefaultConfig()
	}
	return newConn(conn, config)
}

func newConn(conn net.Conn, config *Config) *Conn {
	c := &Conn{
		conn:     conn,
		config:   config,
		reader:   bufio.NewReaderSize(conn, config.BufferSize),
		writer:   bufio.NewWriterSize(conn, config.BufferSize),
		requests: make(map[uint32]chan *response),
		done:     make(chan struct{}),
	}
	go c.readLoop()
	return c
}

// Close closes the connection
func (c *Conn) Close() error {
	if c.closed.Swap(true) {
		return nil
	}
	close(c.done)

	c.mu.Lock()
	for _, ch := range c.requests {
		ch <- &response{err: ErrClosed}
	}
	c.requests = nil
	c.mu.Unlock()

	return c.conn.Close()
}

// Call sends a request and waits for a response
func (c *Conn) Call(ctx context.Context, msgType MessageType, payload []byte) (MessageType, []byte, error) {
	if c.closed.Load() {
		return 0, nil, ErrClosed
	}

	// Allocate request ID and response channel
	id := atomic.AddUint32(&c.nextID, 1)
	respCh := make(chan *response, 1)

	c.mu.Lock()
	if c.requests == nil {
		c.mu.Unlock()
		return 0, nil, ErrClosed
	}
	c.requests[id] = respCh
	c.mu.Unlock()

	defer func() {
		c.mu.Lock()
		delete(c.requests, id)
		c.mu.Unlock()
	}()

	// Write request with request ID prefix
	reqBuf := GetBuffer()
	defer PutBuffer(reqBuf)

	reqBuf.WriteUint32(id)
	reqBuf.Grow(reqBuf.Len() + len(payload))
	copy(reqBuf.Data[reqBuf.offset:], payload)
	reqBuf.offset += len(payload)

	if err := c.write(msgType, reqBuf.Bytes()); err != nil {
		return 0, nil, err
	}

	// Wait for response
	select {
	case <-ctx.Done():
		return 0, nil, ctx.Err()
	case <-c.done:
		return 0, nil, ErrClosed
	case resp := <-respCh:
		if resp.err != nil {
			return 0, nil, resp.err
		}
		return resp.msgType, resp.payload, nil
	}
}

// Send sends a one-way message (no response expected)
func (c *Conn) Send(msgType MessageType, payload []byte) error {
	if c.closed.Load() {
		return ErrClosed
	}
	return c.write(msgType, payload)
}

func (c *Conn) write(msgType MessageType, payload []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.config.WriteTimeout > 0 {
		_ = c.conn.SetWriteDeadline(time.Now().Add(c.config.WriteTimeout))
	}

	if err := WriteMessage(c.writer, msgType, payload); err != nil {
		return err
	}
	return c.writer.Flush()
}

func (c *Conn) readLoop() {
	for {
		if c.config.ReadTimeout > 0 {
			_ = c.conn.SetReadDeadline(time.Now().Add(c.config.ReadTimeout))
		}

		msgType, payload, err := ReadMessage(c.reader)
		if err != nil {
			if !c.closed.Load() {
				c.Close()
			}
			return
		}

		// Extract request ID from response
		if len(payload) < 4 {
			continue
		}
		reqID := binary.BigEndian.Uint32(payload[:4])
		respPayload := payload[4:]

		c.mu.Lock()
		ch, ok := c.requests[reqID]
		c.mu.Unlock()

		if ok {
			// Error responses carry MsgErrorFlag and have the error
			// string as their payload. Success responses carry only
			// MsgResponseFlag and have the application payload.
			if msgType&MsgErrorFlag != 0 {
				ch <- &response{
					err: fmt.Errorf("remote error: %s", string(respPayload)),
				}
			} else {
				ch <- &response{
					msgType: msgType,
					payload: respPayload,
				}
			}
		}
	}
}

// Listener accepts ZAP connections
type Listener struct {
	listener net.Listener
	config   *Config
}

// Listen creates a new ZAP listener
func Listen(addr string, config *Config) (*Listener, error) {
	if config == nil {
		config = DefaultConfig()
	}

	listener, err := net.Listen(transportNetwork(addr), addr)
	if err != nil {
		return nil, fmt.Errorf("zap: listen failed: %w", err)
	}

	return &Listener{
		listener: listener,
		config:   config,
	}, nil
}

// NewListener wraps an existing net.Listener as a ZAP listener. Any
// transport that yields net.Conn — Unix sockets, in-memory pipes, a
// post-handshake Z-Wing listener — can be served as ZAP this way.
func NewListener(raw net.Listener, config *Config) *Listener {
	if config == nil {
		config = DefaultConfig()
	}
	return &Listener{listener: raw, config: config}
}

// Accept accepts a new connection
func (l *Listener) Accept() (*ServerConn, error) {
	conn, err := l.listener.Accept()
	if err != nil {
		return nil, err
	}
	return newServerConn(conn, l.config), nil
}

// Close closes the listener
func (l *Listener) Close() error {
	return l.listener.Close()
}

// Addr returns the listener's address
func (l *Listener) Addr() net.Addr {
	return l.listener.Addr()
}

// ServerConn is a server-side ZAP connection
type ServerConn struct {
	conn   net.Conn
	config *Config

	reader *bufio.Reader
	writer *bufio.Writer

	writeMu sync.Mutex
	closed  atomic.Bool
}

func newServerConn(conn net.Conn, config *Config) *ServerConn {
	return &ServerConn{
		conn:   conn,
		config: config,
		reader: bufio.NewReaderSize(conn, config.BufferSize),
		writer: bufio.NewWriterSize(conn, config.BufferSize),
	}
}

// Read reads the next request from the connection
func (c *ServerConn) Read() (uint32, MessageType, []byte, error) {
	if c.closed.Load() {
		return 0, 0, nil, ErrClosed
	}

	if c.config.ReadTimeout > 0 {
		_ = c.conn.SetReadDeadline(time.Now().Add(c.config.ReadTimeout))
	}

	msgType, payload, err := ReadMessage(c.reader)
	if err != nil {
		return 0, 0, nil, err
	}

	// Extract request ID
	if len(payload) < 4 {
		return 0, 0, nil, ErrInvalidMessage
	}
	reqID := binary.BigEndian.Uint32(payload[:4])

	return reqID, msgType, payload[4:], nil
}

// Write writes a response
func (c *ServerConn) Write(reqID uint32, msgType MessageType, payload []byte) error {
	if c.closed.Load() {
		return ErrClosed
	}

	// Prepend request ID to response
	respBuf := GetBuffer()
	defer PutBuffer(respBuf)

	respBuf.WriteUint32(reqID)
	respBuf.Grow(respBuf.Len() + len(payload))
	copy(respBuf.Data[respBuf.offset:], payload)
	respBuf.offset += len(payload)

	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	if c.config.WriteTimeout > 0 {
		_ = c.conn.SetWriteDeadline(time.Now().Add(c.config.WriteTimeout))
	}

	if err := WriteMessage(c.writer, msgType, respBuf.Bytes()); err != nil {
		return err
	}
	return c.writer.Flush()
}

// Close closes the server connection
func (c *ServerConn) Close() error {
	if c.closed.Swap(true) {
		return nil
	}
	return c.conn.Close()
}

// RemoteAddr returns the remote address
func (c *ServerConn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// Handler processes ZAP requests
type Handler interface {
	// Handle processes a request and returns a response
	Handle(ctx context.Context, msgType MessageType, payload []byte) (MessageType, []byte, error)
}

// HandlerFunc is a function that implements Handler
type HandlerFunc func(ctx context.Context, msgType MessageType, payload []byte) (MessageType, []byte, error)

// Handle implements Handler
func (f HandlerFunc) Handle(ctx context.Context, msgType MessageType, payload []byte) (MessageType, []byte, error) {
	return f(ctx, msgType, payload)
}

// Server serves ZAP requests
type Server struct {
	listener *Listener
	handler  Handler

	mu     sync.Mutex
	conns  map[*ServerConn]struct{}
	closed atomic.Bool
	done   chan struct{}
}

// NewServer creates a new ZAP server
func NewServer(listener *Listener, handler Handler) *Server {
	return &Server{
		listener: listener,
		handler:  handler,
		conns:    make(map[*ServerConn]struct{}),
		done:     make(chan struct{}),
	}
}

// Serve accepts and processes connections
func (s *Server) Serve(ctx context.Context) error {
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-s.done:
			return nil
		default:
		}

		conn, err := s.listener.Accept()
		if err != nil {
			if s.closed.Load() {
				return nil
			}
			continue
		}

		s.mu.Lock()
		// Close() sets conns to nil under this lock. Accept can return a
		// connection that was already in flight when that happened, and
		// assigning into the nil map panics the Serve goroutine — taking the
		// whole process with it, since this runs detached. Re-check under the
		// lock and hang up instead: after Close there is nobody left to serve.
		if s.conns == nil {
			s.mu.Unlock()
			conn.Close()
			return nil
		}
		s.conns[conn] = struct{}{}
		s.mu.Unlock()

		go s.handleConn(ctx, conn)
	}
}

func (s *Server) handleConn(ctx context.Context, conn *ServerConn) {
	defer func() {
		conn.Close()
		s.mu.Lock()
		delete(s.conns, conn)
		s.mu.Unlock()
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case <-s.done:
			return
		default:
		}

		reqID, msgType, payload, err := conn.Read()
		if err != nil {
			if err == io.EOF || errors.Is(err, net.ErrClosed) {
				return
			}
			continue
		}

		// Handle request concurrently to avoid blocking on long-running
		// operations like WaitForEvent. Write is protected by writeMu.
		go func(reqID uint32, msgType MessageType, payload []byte) {
			defer func() {
				if r := recover(); r != nil {
					errMsg := fmt.Sprintf("handler panic: %v", r)
					conn.Write(reqID, msgType|MsgResponseFlag|MsgErrorFlag, []byte(errMsg))
				}
			}()

			respType, respPayload, err := s.handler.Handle(ctx, msgType, payload)
			if err != nil {
				conn.Write(reqID, msgType|MsgResponseFlag|MsgErrorFlag, []byte(err.Error()))
				return
			}

			// Success responses carry only MsgResponseFlag — the client
			// demultiplexer routes by request ID and distinguishes
			// success from error via MsgErrorFlag.
			conn.Write(reqID, respType|MsgResponseFlag, respPayload)
		}(reqID, msgType, payload)
	}
}

// Close closes the server
func (s *Server) Close() error {
	if s.closed.Swap(true) {
		return nil
	}
	close(s.done)

	s.mu.Lock()
	for conn := range s.conns {
		conn.Close()
	}
	s.conns = nil
	s.mu.Unlock()

	return s.listener.Close()
}