internal/transport/frame.go

Ref: main v0.0.5   Size: 1.3 KiB   History

package transport

import (
	"encoding/binary"
	"fmt"
	"io"

	"google.golang.org/protobuf/proto"
)

// DefaultMaxFrame bounds a single framed message (protect against a corrupt or
// hostile length prefix). Snapshots/reports are small; 4 MiB is generous.
const DefaultMaxFrame = 4 << 20

// WriteMsg writes msg as a 4-byte big-endian length prefix followed by its
// protobuf-marshaled bytes.
func WriteMsg(w io.Writer, msg proto.Message) error {
	b, err := proto.Marshal(msg)
	if err != nil {
		return fmt.Errorf("marshal: %w", err)
	}
	var hdr [4]byte
	binary.BigEndian.PutUint32(hdr[:], uint32(len(b)))
	if _, err := w.Write(hdr[:]); err != nil {
		return err
	}
	_, err = w.Write(b)
	return err
}

// ReadMsg reads one framed message into out. It bounds-checks the length against
// maxFrame BEFORE allocating, then reads exactly that many bytes. Returns io.EOF
// when the stream is cleanly closed at a frame boundary.
func ReadMsg(r io.Reader, out proto.Message, maxFrame uint32) error {
	var hdr [4]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return err
	}
	n := binary.BigEndian.Uint32(hdr[:])
	if n > maxFrame {
		return fmt.Errorf("frame too large: %d > %d", n, maxFrame)
	}
	body := make([]byte, n)
	if _, err := io.ReadFull(r, body); err != nil {
		return err
	}
	return proto.Unmarshal(body, out)
}