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 }