batch.go

package fastwire

import "encoding/binary"

const (
	// BatchHeaderSize is the size of the packet count byte in a batched datagram.
	BatchHeaderSize = 1
	// BatchLenSize is the size of each packet's length prefix.
	BatchLenSize = 2
)

// MarshalBatch packs multiple encrypted packets into a single datagram.
// Format: [count:1B][len1:2B LE][pkt1][len2:2B LE][pkt2]...
func MarshalBatch(dst []byte, packets [][]byte) (int, error) {
	if len(packets) == 0 || len(packets) > 255 {
		return 0, ErrInvalidBatchFrame
	}

	needed := BatchHeaderSize
	for _, p := range packets {
		needed += BatchLenSize + len(p)
	}
	if len(dst) < needed {
		return 0, ErrBufferTooSmall
	}

	n := 0
	dst[n] = byte(len(packets))
	n++

	for _, p := range packets {
		binary.LittleEndian.PutUint16(dst[n:], uint16(len(p)))
		n += BatchLenSize
		copy(dst[n:], p)
		n += len(p)
	}

	return n, nil
}

// UnmarshalBatch unpacks a batched datagram into individual encrypted packets.
func UnmarshalBatch(data []byte) ([][]byte, error) {
	if len(data) < BatchHeaderSize {
		return nil, ErrInvalidBatchFrame
	}

	count := int(data[0])
	if count == 0 {
		return nil, ErrInvalidBatchFrame
	}

	packets := make([][]byte, 0, count)
	offset := BatchHeaderSize

	for range count {
		if offset+BatchLenSize > len(data) {
			return nil, ErrInvalidBatchFrame
		}
		pktLen := int(binary.LittleEndian.Uint16(data[offset:]))
		offset += BatchLenSize
		if offset+pktLen > len(data) {
			return nil, ErrInvalidBatchFrame
		}
		packets = append(packets, data[offset:offset+pktLen])
		offset += pktLen
	}

	return packets, nil
}

// BatchOverhead returns the wire overhead in bytes for a batch of n packets.
func BatchOverhead(n int) int {
	return BatchHeaderSize + n*BatchLenSize
}

// flushBatch packs the connection's batch buffer into batched datagrams and sends them.
// All pending buffers are returned to the pool on return, regardless of error.
func (c *Connection) flushBatch(mtu int) error {
	c.batchMu.Lock()
	if len(c.batchBuf) == 0 {
		c.batchMu.Unlock()
		return nil
	}
	// Swap out the pending packets under lock; process without lock.
	pending := make([][]byte, len(c.batchBuf))
	copy(pending, c.batchBuf)
	c.batchBuf = c.batchBuf[:0]
	c.batchMu.Unlock()

	defer func() {
		for _, p := range pending {
			putSendBuffer(p[:cap(p)])
		}
	}()

	var dgBuf []byte
	pktCount := 0

	flush := func() error {
		if pktCount == 0 {
			return nil
		}
		dgBuf[0] = byte(pktCount)
		err := c.sendFunc(dgBuf)
		putSendBuffer(dgBuf[:cap(dgBuf)])
		dgBuf = nil
		pktCount = 0
		return err
	}

	for _, pkt := range pending {
		entrySize := BatchLenSize + len(pkt)

		// If a single packet doesn't fit in a fresh datagram, send as a single-packet batch.
		if BatchHeaderSize+entrySize > mtu {
			if err := flush(); err != nil {
				return err
			}
			single := getSendBuffer(BatchHeaderSize + entrySize)
			single[0] = 1
			binary.LittleEndian.PutUint16(single[1:], uint16(len(pkt)))
			copy(single[3:], pkt)
			err := c.sendFunc(single[:BatchHeaderSize+entrySize])
			putSendBuffer(single[:cap(single)])
			if err != nil {
				return err
			}
			continue
		}

		if dgBuf == nil {
			dgBuf = getSendBuffer(mtu)[:1]
		}

		// Check if adding this packet would exceed MTU or max count.
		if len(dgBuf)+entrySize > mtu || pktCount >= 255 {
			if err := flush(); err != nil {
				return err
			}
			dgBuf = getSendBuffer(mtu)[:1]
		}

		dgBuf = binary.LittleEndian.AppendUint16(dgBuf, uint16(len(pkt)))
		dgBuf = append(dgBuf, pkt...)
		pktCount++
	}

	return flush()
}