coder/agent/boundarylogproxy/codec/codec.go

// Package codec implements the wire format for agent <-> boundary communication.
//
// Wire Format:
//   - 8 bits: big-endian tag
//   - 24 bits: big-endian length of the protobuf data (bit usage depends on tag)
//   - length bytes: encoded protobuf data
//
// Note that while there are 24 bits available for the length, the actual maximum
// length depends on the tag. For TagV1, only 15 bits are used (MaxMessageSizeV1).
package codec

import (
	"encoding/binary"
	"io"

	"golang.org/x/xerrors"
	"google.golang.org/protobuf/proto"

	agentproto "github.com/coder/coder/v2/agent/proto"
)

type Tag uint8

const (
	// TagV1 identifies the first revision of the protocol. The payload is a
	// bare ReportBoundaryLogsRequest. This version has a maximum data length
	// of MaxMessageSizeV1.
	TagV1 Tag = 1

	// TagV2 identifies the second revision of the protocol. The payload is
	// a BoundaryMessage envelope. This version has a maximum data length of
	// MaxMessageSizeV2.
	TagV2 Tag = 2
)

const (
	// DataLength is the number of bits used for the length of encoded protobuf data.
	DataLength = 24

	// tagLength is the number of bits used for the tag.
	tagLength = 8

	// MaxMessageSizeV1 is the maximum size of the encoded protobuf messages sent
	// over the wire for the TagV1 tag. While the wire format allows 24 bits for
	// length, TagV1 only uses 15 bits.
	MaxMessageSizeV1 uint32 = 1 << 15

	// MaxMessageSizeV2 is the maximum data length for TagV2.
	MaxMessageSizeV2 = MaxMessageSizeV1
)

var (
	// ErrMessageTooLarge is returned when the message exceeds the maximum size
	// allowed for the tag.
	ErrMessageTooLarge = xerrors.New("message too large")
	// ErrUnsupportedTag is returned when an unrecognized tag is encountered.
	ErrUnsupportedTag = xerrors.New("unsupported tag")
)

// WriteFrame writes a framed message with the given tag and data. The data
// must not exceed 2^DataLength in length.
func WriteFrame(w io.Writer, tag Tag, data []byte) error {
	maxSize, err := maxSizeForTag(tag)
	if err != nil {
		return err
	}

	if len(data) > int(maxSize) {
		return xerrors.Errorf("%w for tag %d: %d > %d", ErrMessageTooLarge, tag, len(data), maxSize)
	}

	var header uint32
	//nolint:gosec // The length check above ensures there's no overflow.
	header |= uint32(len(data))
	header |= uint32(tag) << DataLength

	if err := binary.Write(w, binary.BigEndian, header); err != nil {
		return xerrors.Errorf("write header error: %w", err)
	}
	if _, err := w.Write(data); err != nil {
		return xerrors.Errorf("write data error: %w", err)
	}

	return nil
}

// ReadFrame reads a framed message, returning the decoded tag and data. If the
// message size exceeds MaxMessageSizeV1, ErrMessageTooLarge is returned. The
// provided buf is used if it has sufficient capacity; otherwise a new buffer is
// allocated. To reuse the buffer across calls, pass in the returned data slice:
//
//	buf := make([]byte, initialSize)
//	for {
//	    _, buf, _ = ReadFrame(r, buf)
//	}
func ReadFrame(r io.Reader, buf []byte) (Tag, []byte, error) {
	var header uint32
	if err := binary.Read(r, binary.BigEndian, &header); err != nil {
		return 0, nil, xerrors.Errorf("read header error: %w", err)
	}

	const lengthMask = (1 << DataLength) - 1
	length := header & lengthMask
	const tagMask = (1 << tagLength) - 1 // 0xFF
	shifted := (header >> DataLength) & tagMask
	if shifted > tagMask {
		// This is really only here to satisfy the gosec linter. We know from above that
		// shifted <= tagMask.
		return 0, nil, xerrors.Errorf("invalid tag: %d", shifted)
	}
	tag := Tag(shifted)

	maxSize, err := maxSizeForTag(tag)
	if err != nil {
		return 0, nil, err
	}

	if length > maxSize {
		return 0, nil, ErrMessageTooLarge
	}

	if cap(buf) < int(length) {
		buf = make([]byte, length)
	} else {
		buf = buf[:length:cap(buf)]
	}

	if _, err := io.ReadFull(r, buf[:length]); err != nil {
		return 0, nil, xerrors.Errorf("read full error: %w", err)
	}

	return tag, buf[:length], nil
}

// maxSizeForTag returns the maximum payload size for the given tag.
func maxSizeForTag(tag Tag) (uint32, error) {
	switch tag {
	case TagV1:
		return MaxMessageSizeV1, nil
	case TagV2:
		return MaxMessageSizeV2, nil
	default:
		return 0, xerrors.Errorf("%w: %d", ErrUnsupportedTag, tag)
	}
}

// ReadMessage reads a framed message and unmarshals it based on tag. The
// returned buf should be passed back on the next call for buffer reuse.
func ReadMessage(r io.Reader, buf []byte) (proto.Message, []byte, error) {
	tag, data, err := ReadFrame(r, buf)
	if err != nil {
		return nil, data, err
	}

	var msg proto.Message
	switch tag {
	case TagV1:
		var req agentproto.ReportBoundaryLogsRequest
		if err := proto.Unmarshal(data, &req); err != nil {
			return nil, data, xerrors.Errorf("unmarshal TagV1: %w", err)
		}
		msg = &req
	case TagV2:
		var envelope BoundaryMessage
		if err := proto.Unmarshal(data, &envelope); err != nil {
			return nil, data, xerrors.Errorf("unmarshal TagV2: %w", err)
		}
		msg = &envelope
	default:
		// maxSizeForTag already rejects unknown tags during ReadFrame,
		// but handle it here for safety.
		return nil, data, xerrors.Errorf("%w: %d", ErrUnsupportedTag, tag)
	}

	return msg, data, nil
}

// WriteMessage marshals a proto message and writes it as a framed message
// with the given tag.
func WriteMessage(w io.Writer, tag Tag, msg proto.Message) error {
	data, err := proto.Marshal(msg)
	if err != nil {
		return xerrors.Errorf("marshal: %w", err)
	}
	return WriteFrame(w, tag, data)
}