Implement FD batching per protocol spec Section 4.1 (#4)

The jsonrpc-fdpass spec calls for a mechanism to handle more FDs
than can be sent in a single sendmsg() - add support for that.

Assisted-by: OpenCode (Claude Opus 4)

Signed-off-by: Colin Walters <walters@verbum.org>

Author: cgwalters

transport.go

@@ -14,8 +14,16 @@ import (
 )
 
 const (
-	// MaxFDsPerMessage is the maximum number of file descriptors per message.
-	MaxFDsPerMessage = 8
+	// DefaultMaxFDsPerSendmsg is the default maximum number of file descriptors
+	// per sendmsg() call. Platform limits for SCM_RIGHTS vary (e.g., ~253 on
+	// Linux, ~512 on macOS). We start with an optimistic value; if sendmsg()
+	// fails with EINVAL, the batch size is automatically reduced and retried.
+	DefaultMaxFDsPerSendmsg = 500
+
+	// MaxFDsPerRecvmsg is the maximum number of FDs to expect in a single
+	// recvmsg() call. Must be at least as large as the largest platform limit.
+	MaxFDsPerRecvmsg = 512
+
 	// ReadBufferSize is the size of the read buffer.
 	ReadBufferSize = 4096
 )
@@ -31,13 +39,27 @@ var (
 
 // Sender sends JSON-RPC messages with file descriptors over a Unix socket.
 type Sender struct {
-	conn *net.UnixConn
-	mu   sync.Mutex
+	conn            *net.UnixConn
+	mu              sync.Mutex
+	maxFDsPerSendmsg int
 }
 
 // NewSender creates a new Sender for the given Unix connection.
 func NewSender(conn *net.UnixConn) *Sender {
-	return &Sender{conn: conn}
+	return &Sender{
+		conn:            conn,
+		maxFDsPerSendmsg: DefaultMaxFDsPerSendmsg,
+	}
+}
+
+// SetMaxFDsPerSendmsg sets the maximum number of file descriptors to send per
+// sendmsg() call. This is primarily useful for testing FD batching behavior.
+// The value must be at least 1.
+func (s *Sender) SetMaxFDsPerSendmsg(max int) {
+	if max < 1 {
+		max = 1
+	}
+	s.maxFDsPerSendmsg = max
 }
 
 // Send sends a message with optional file descriptors.
@@ -79,37 +101,77 @@ func (s *Sender) Send(msg *MessageWithFds) error {
 }
 
 func (s *Sender) sendWithFDs(sockfd int, data []byte, files []*os.File) error {
-	bytesSent := 0
-	fdsSent := false
+	// Extract raw FD ints
+	allFDs := make([]int, len(files))
+	for i, f := range files {
+		allFDs[i] = int(f.Fd())
+	}
 
-	for bytesSent < len(data) {
-		remaining := data[bytesSent:]
+	bytesSent := 0
+	fdsSent := 0
+	currentMaxFDs := s.maxFDsPerSendmsg
+
+	// Send data and FDs in batches. Each sendmsg can only handle a limited
+	// number of FDs. After all data bytes are sent, remaining FDs are sent
+	// with whitespace padding bytes per the protocol spec (Section 4.1).
+	for bytesSent < len(data) || fdsSent < len(allFDs) {
+		remainingData := data[bytesSent:]
+		remainingFDs := allFDs[fdsSent:]
+
+		// Determine how many FDs to send in this batch
+		fdBatchSize := len(remainingFDs)
+		if fdBatchSize > currentMaxFDs {
+			fdBatchSize = currentMaxFDs
+		}
+		fdBatch := remainingFDs[:fdBatchSize]
 
 		var n int
 		var err error
 
-		if !fdsSent && len(files) > 0 {
-			// First chunk with FDs: use sendmsg with ancillary data
-			fds := make([]int, len(files))
-			for i, f := range files {
-				fds[i] = int(f.Fd())
+		if len(fdBatch) > 0 {
+			// Send with FDs using sendmsg with ancillary data
+			rights := unix.UnixRights(fdBatch...)
+
+			var payload []byte
+			if len(remainingData) > 0 {
+				payload = remainingData
+			} else {
+				// All data bytes already sent; send a whitespace padding byte.
+				// The receiver's JSON parser ignores inter-message whitespace
+				// per RFC 8259. This is required because some systems need
+				// non-empty data for ancillary data delivery.
+				payload = []byte{' '}
 			}
 
-			rights := unix.UnixRights(fds...)
-			n, err = unix.SendmsgN(sockfd, remaining, rights, nil, 0)
+			n, err = unix.SendmsgN(sockfd, payload, rights, nil, 0)
 			if err != nil {
+				// EINVAL with multiple FDs likely means we exceeded the
+				// kernel's SCM_MAX_FD limit. Halve the batch size and retry.
+				if errors.Is(err, unix.EINVAL) && fdBatchSize > 1 {
+					currentMaxFDs = fdBatchSize / 2
+					continue
+				}
 				return fmt.Errorf("sendmsg failed: %w", err)
 			}
-			fdsSent = true
-		} else {
-			// No FDs or FDs already sent: use regular send
-			n, err = unix.Write(sockfd, remaining)
+			fdsSent += fdBatchSize
+
+			// Only count actual data bytes, not the padding byte
+			if len(remainingData) > 0 {
+				bytesSent += n
+			}
+		} else if len(remainingData) > 0 {
+			// No FDs left, just send remaining data bytes
+			n, err = unix.Write(sockfd, remainingData)
 			if err != nil {
 				return fmt.Errorf("write failed: %w", err)
 			}
+			bytesSent += n
 		}
+	}
 
-		bytesSent += n
+	// If we discovered a lower limit, remember it for future sends
+	if currentMaxFDs < s.maxFDsPerSendmsg {
+		s.maxFDsPerSendmsg = currentMaxFDs
 	}
 
 	return nil
@@ -139,24 +201,37 @@ func (r *Receiver) Receive() (*MessageWithFds, error) {
 
 	for {
 		// Try to parse a complete message from the buffer
-		msg, err := r.tryParseMessage()
+		result, err := r.tryParseMessage()
 		if err != nil {
 			return nil, err
 		}
-		if msg != nil {
-			return msg, nil
+		if result.msg != nil {
+			return result.msg, nil
 		}
 
-		// Need more data
+		// Need more data — either incomplete JSON or waiting for
+		// batched FDs from continuation sendmsg() calls.
 		if err := r.readMoreData(); err != nil {
+			// If we had a parsed message waiting for FDs and the
+			// connection closed, that's a mismatched count error.
+			if result.needFDs && errors.Is(err, ErrConnectionClosed) {
+				return nil, fmt.Errorf("%w: connection closed while waiting for batched FDs", ErrMismatchedCount)
+			}
 			return nil, err
 		}
 	}
 }
 
-func (r *Receiver) tryParseMessage() (*MessageWithFds, error) {
+// tryParseResult is used internally to communicate between tryParseMessage
+// and Receive about whether more FDs are needed from batched sendmsg calls.
+type tryParseResult struct {
+	msg     *MessageWithFds
+	needFDs bool // true when message is parsed but FD queue is short
+}
+
+func (r *Receiver) tryParseMessage() (*tryParseResult, error) {
 	if len(r.buffer) == 0 {
-		return nil, nil
+		return &tryParseResult{}, nil
 	}
 
 	// Use streaming JSON decoder to find message boundaries
@@ -166,7 +241,7 @@ func (r *Receiver) tryParseMessage() (*MessageWithFds, error) {
 	err := decoder.Decode(&value)
 	if err == io.EOF || errors.Is(err, io.ErrUnexpectedEOF) {
 		// Incomplete JSON - need more data
-		return nil, nil
+		return &tryParseResult{}, nil
 	}
 	if err != nil {
 		// Actual parse error - framing error
@@ -178,20 +253,22 @@ func (r *Receiver) tryParseMessage() (*MessageWithFds, error) {
 	bytesConsumed := decoder.InputOffset()
 
 	// Extract the consumed bytes for re-parsing
-	consumedData := r.buffer[:bytesConsumed]
-
-	// Remove consumed bytes from buffer
-	r.buffer = r.buffer[bytesConsumed:]
+	consumedData := make([]byte, bytesConsumed)
+	copy(consumedData, r.buffer[:bytesConsumed])
 
 	// Read the fds count from the message
 	fdCount := GetFDCount(value)
 
-	// Check we have enough FDs
+	// Check we have enough FDs. When FD batching is in use, the sender
+	// sends continuation sendmsg() calls with whitespace padding and more
+	// FDs. We need to read more data to collect them.
 	if fdCount > len(r.fdQueue) {
-		return nil, fmt.Errorf("%w: expected %d FDs, have %d in queue",
-			ErrMismatchedCount, fdCount, len(r.fdQueue))
+		return &tryParseResult{needFDs: true}, nil
 	}
 
+	// Remove consumed bytes from buffer
+	r.buffer = r.buffer[bytesConsumed:]
+
 	// Dequeue FDs
 	fds := make([]*os.File, fdCount)
 	copy(fds, r.fdQueue[:fdCount])
@@ -203,9 +280,11 @@ func (r *Receiver) tryParseMessage() (*MessageWithFds, error) {
 		return nil, err
 	}
 
-	return &MessageWithFds{
-		Message:         msg,
-		FileDescriptors: fds,
+	return &tryParseResult{
+		msg: &MessageWithFds{
+			Message:         msg,
+			FileDescriptors: fds,
+		},
 	}, nil
 }
 
@@ -251,9 +330,9 @@ func (r *Receiver) readMoreData() error {
 
 func (r *Receiver) recvWithFDs(sockfd int) (int, []*os.File, error) {
 	buf := make([]byte, ReadBufferSize)
-	// Allocate space for control message (for up to MaxFDsPerMessage FDs)
+	// Allocate space for control message (for up to MaxFDsPerRecvmsg FDs)
 	// Each FD is 4 bytes (int32), use CmsgSpace to get properly aligned size
-	oob := make([]byte, unix.CmsgSpace(MaxFDsPerMessage*4))
+	oob := make([]byte, unix.CmsgSpace(MaxFDsPerRecvmsg*4))
 
 	n, oobn, _, _, err := unix.Recvmsg(sockfd, buf, oob, unix.MSG_CMSG_CLOEXEC)
 	if err != nil {