notify/record_codec.go

package notify

import (
	"encoding/binary"
	"errors"
)

const (
	recordFrameMagic            = "NRS1"
	recordFrameVersion          = 1
	recordFrameTypeBatch  uint8 = 1
	recordFrameTypeAck    uint8 = 2
	recordFrameTypeError  uint8 = 3
	recordFrameHeaderSize       = 8
	recordBatchHeaderSize       = 10
	recordErrorHeaderSize       = 16
)

var (
	errRecordFrameInvalid = errors.New("invalid record frame")
	errRecordSeqInvalid   = errors.New("invalid record sequence")
)

type recordOutboundMessage struct {
	Seq     uint64
	Payload []byte
}

type recordFrame struct {
	Type      uint8
	Batch     []recordOutboundMessage
	AckSeq    uint64
	Failure   RecordFailure
	Retryable bool
}

func encodeRecordBatchFrame(batch []recordOutboundMessage) ([]byte, error) {
	if len(batch) == 0 {
		return nil, nil
	}
	firstSeq := batch[0].Seq
	if firstSeq == 0 {
		return nil, errRecordSeqInvalid
	}
	size := recordFrameHeaderSize + recordBatchHeaderSize
	for index, item := range batch {
		wantSeq := firstSeq + uint64(index)
		if item.Seq != wantSeq {
			return nil, errRecordSeqInvalid
		}
		size += 4 + len(item.Payload)
	}
	frame := make([]byte, size)
	copy(frame[:4], recordFrameMagic)
	frame[4] = recordFrameVersion
	frame[5] = recordFrameTypeBatch
	binary.BigEndian.PutUint16(frame[8:10], uint16(len(batch)))
	binary.BigEndian.PutUint64(frame[10:18], firstSeq)
	offset := recordFrameHeaderSize + recordBatchHeaderSize
	for _, item := range batch {
		binary.BigEndian.PutUint32(frame[offset:offset+4], uint32(len(item.Payload)))
		offset += 4
		copy(frame[offset:offset+len(item.Payload)], item.Payload)
		offset += len(item.Payload)
	}
	return frame, nil
}

func encodeRecordAckFrame(ackSeq uint64) ([]byte, error) {
	frame := make([]byte, recordFrameHeaderSize+8)
	copy(frame[:4], recordFrameMagic)
	frame[4] = recordFrameVersion
	frame[5] = recordFrameTypeAck
	binary.BigEndian.PutUint64(frame[8:16], ackSeq)
	return frame, nil
}

func encodeRecordErrorFrame(failure RecordFailure) ([]byte, error) {
	if failure.FailedSeq == 0 {
		return nil, errRecordSeqInvalid
	}
	codeBytes := []byte(failure.Code)
	msgBytes := []byte(failure.Message)
	frame := make([]byte, recordFrameHeaderSize+recordErrorHeaderSize+len(codeBytes)+len(msgBytes))
	copy(frame[:4], recordFrameMagic)
	frame[4] = recordFrameVersion
	frame[5] = recordFrameTypeError
	if failure.Retryable {
		frame[6] = 1
	}
	binary.BigEndian.PutUint64(frame[8:16], failure.FailedSeq)
	binary.BigEndian.PutUint16(frame[16:18], uint16(len(codeBytes)))
	binary.BigEndian.PutUint32(frame[18:22], uint32(len(msgBytes)))
	offset := recordFrameHeaderSize + recordErrorHeaderSize
	copy(frame[offset:offset+len(codeBytes)], codeBytes)
	offset += len(codeBytes)
	copy(frame[offset:offset+len(msgBytes)], msgBytes)
	return frame, nil
}

func decodeRecordFrame(payload []byte) (recordFrame, error) {
	if len(payload) < recordFrameHeaderSize || string(payload[:4]) != recordFrameMagic {
		return recordFrame{}, errRecordFrameInvalid
	}
	if payload[4] != recordFrameVersion {
		return recordFrame{}, errRecordFrameInvalid
	}
	frameType := payload[5]
	switch frameType {
	case recordFrameTypeBatch:
		return decodeRecordBatchFrame(payload)
	case recordFrameTypeAck:
		if len(payload) != recordFrameHeaderSize+8 {
			return recordFrame{}, errRecordFrameInvalid
		}
		return recordFrame{
			Type:   recordFrameTypeAck,
			AckSeq: binary.BigEndian.Uint64(payload[8:16]),
		}, nil
	case recordFrameTypeError:
		return decodeRecordErrorFrame(payload)
	default:
		return recordFrame{}, errRecordFrameInvalid
	}
}

func decodeRecordBatchFrame(payload []byte) (recordFrame, error) {
	if len(payload) < recordFrameHeaderSize+recordBatchHeaderSize {
		return recordFrame{}, errRecordFrameInvalid
	}
	count := int(binary.BigEndian.Uint16(payload[8:10]))
	firstSeq := binary.BigEndian.Uint64(payload[10:18])
	if count <= 0 || firstSeq == 0 {
		return recordFrame{}, errRecordFrameInvalid
	}
	offset := recordFrameHeaderSize + recordBatchHeaderSize
	batch := make([]recordOutboundMessage, 0, count)
	for index := 0; index < count; index++ {
		if offset+4 > len(payload) {
			return recordFrame{}, errRecordFrameInvalid
		}
		itemLen := int(binary.BigEndian.Uint32(payload[offset : offset+4]))
		offset += 4
		if itemLen < 0 || offset+itemLen > len(payload) {
			return recordFrame{}, errRecordFrameInvalid
		}
		item := recordOutboundMessage{
			Seq:     firstSeq + uint64(index),
			Payload: append([]byte(nil), payload[offset:offset+itemLen]...),
		}
		offset += itemLen
		batch = append(batch, item)
	}
	if offset != len(payload) {
		return recordFrame{}, errRecordFrameInvalid
	}
	return recordFrame{
		Type:  recordFrameTypeBatch,
		Batch: batch,
	}, nil
}

func decodeRecordErrorFrame(payload []byte) (recordFrame, error) {
	if len(payload) < recordFrameHeaderSize+recordErrorHeaderSize {
		return recordFrame{}, errRecordFrameInvalid
	}
	failedSeq := binary.BigEndian.Uint64(payload[8:16])
	codeLen := int(binary.BigEndian.Uint16(payload[16:18]))
	msgLen := int(binary.BigEndian.Uint32(payload[18:22]))
	offset := recordFrameHeaderSize + recordErrorHeaderSize
	if failedSeq == 0 || offset+codeLen+msgLen != len(payload) {
		return recordFrame{}, errRecordFrameInvalid
	}
	failure := RecordFailure{
		FailedSeq: failedSeq,
		Retryable: payload[6] == 1,
		Code:      RecordErrorCode(string(payload[offset : offset+codeLen])),
		Message:   string(payload[offset+codeLen:]),
	}
	return recordFrame{
		Type:    recordFrameTypeError,
		Failure: failure,
	}, nil
}