DebTik/internal/bwtest/tcp_frames.go

package bwtest

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
)

// A single framed stream multiplexes data and reports without competing readers
// or interleaved writes. The newline JSON handshake negotiates this format.
const tcpFramingVersion = "frames-v1"
const (
	tcpFrameData    byte = 1
	tcpFrameControl byte = 2
	tcpFrameMax          = 64 * 1024
)

type tcpFrameWriter struct {
	conn net.Conn
	mu   sync.Mutex
}

func (w *tcpFrameWriter) write(kind byte, payload []byte) (int, error) {
	if len(payload) == 0 || len(payload) > tcpFrameMax {
		return 0, fmt.Errorf("invalid TCP bandwidth frame size")
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	var h [5]byte
	h[0] = kind
	binary.BigEndian.PutUint32(h[1:], uint32(len(payload)))
	buffers := net.Buffers{h[:], payload}
	n, err := buffers.WriteTo(w.conn)
	if n < 5 {
		return 0, err
	}
	return int(n) - 5, err
}
func readTCPFrame(r io.Reader, buf []byte) (byte, []byte, error) {
	var h [5]byte
	if _, err := io.ReadFull(r, h[:]); err != nil {
		return 0, nil, err
	}
	n := int(binary.BigEndian.Uint32(h[1:]))
	if (h[0] != tcpFrameData && h[0] != tcpFrameControl) || n == 0 || n > tcpFrameMax || n > len(buf) {
		return 0, nil, fmt.Errorf("invalid TCP bandwidth frame")
	}
	_, err := io.ReadFull(r, buf[:n])
	return h[0], buf[:n], err
}