notify/record_codec.go

package notify

import (
	"encoding/binary"
	"errors"
)

const (
	recordFrameMagic              = "NRS1"
	recordFrameVersionV1          = 1
	recordFrameVersionV2          = 2
	recordFrameTypeBatch    uint8 = 1
	recordFrameTypeAck      uint8 = 2
	recordFrameTypeError    uint8 = 3
	recordFrameTypeFIN      uint8 = 4
	recordFrameHeaderSize         = 8
	recordBatchHeaderV1Size       = 10
	recordBatchHeaderV2Size       = 18
	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 {
	Version   uint8
	Type      uint8
	Batch     []recordOutboundMessage
	AckSeq    uint64
	FinalSeq  uint64
	Failure   RecordFailure
	Retryable bool
}

func encodeRecordBatchFrame(batch []recordOutboundMessage, ackSeq uint64, useV2 bool) ([]byte, error) {
	if len(batch) == 0 {
		return nil, nil
	}
	if len(batch) > recordMaxBatchRecords {
		return nil, errRecordFrameInvalid
	}
	firstSeq := batch[0].Seq
	if firstSeq == 0 || uint64(len(batch)-1) > ^uint64(0)-firstSeq {
		return nil, errRecordSeqInvalid
	}
	version := uint8(recordFrameVersionV1)
	batchHeaderSize := recordBatchHeaderV1Size
	if useV2 {
		version = recordFrameVersionV2
		batchHeaderSize = recordBatchHeaderV2Size
	}
	size := recordFrameHeaderSize + batchHeaderSize
	for index, item := range batch {
		wantSeq := firstSeq + uint64(index)
		if item.Seq != wantSeq {
			return nil, errRecordSeqInvalid
		}
		if len(item.Payload) > recordMaxPayloadBytes || size > transferFrameMaxPayloadBytes-4-len(item.Payload) {
			return nil, errRecordFrameInvalid
		}
		size += 4 + len(item.Payload)
	}
	frame := make([]byte, size)
	copy(frame[:4], recordFrameMagic)
	frame[4] = version
	frame[5] = recordFrameTypeBatch
	binary.BigEndian.PutUint16(frame[8:10], uint16(len(batch)))
	binary.BigEndian.PutUint64(frame[10:18], firstSeq)
	offset := recordFrameHeaderSize + batchHeaderSize
	if useV2 {
		binary.BigEndian.PutUint64(frame[18:26], ackSeq)
	}
	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] = recordFrameVersionV1
	frame[5] = recordFrameTypeAck
	binary.BigEndian.PutUint64(frame[8:16], ackSeq)
	return frame, nil
}

func encodeRecordFINFrame(finalSeq uint64) []byte {
	frame, _ := encodeRecordAckFrame(finalSeq)
	frame[5] = recordFrameTypeFIN
	return frame
}

func encodeRecordErrorFrame(failure RecordFailure) ([]byte, error) {
	if failure.FailedSeq == 0 {
		return nil, errRecordSeqInvalid
	}
	const headerSize = recordFrameHeaderSize + recordErrorHeaderSize
	if len(failure.Code) > int(^uint16(0)) || len(failure.Message) > transferFrameMaxPayloadBytes-headerSize-len(failure.Code) {
		return nil, errRecordFrameInvalid
	}
	frame := make([]byte, headerSize+len(failure.Code)+len(failure.Message))
	copy(frame[:4], recordFrameMagic)
	frame[4] = recordFrameVersionV1
	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(failure.Code)))
	binary.BigEndian.PutUint32(frame[18:22], uint32(len(failure.Message)))
	offset := recordFrameHeaderSize + recordErrorHeaderSize
	copy(frame[offset:], failure.Code)
	offset += len(failure.Code)
	copy(frame[offset:], failure.Message)
	return frame, nil
}

func decodeRecordFrame(payload []byte) (recordFrame, error) {
	if len(payload) < recordFrameHeaderSize || string(payload[:4]) != recordFrameMagic {
		return recordFrame{}, errRecordFrameInvalid
	}
	version := payload[4]
	frameType := payload[5]
	if frameType == recordFrameTypeFIN && version == recordFrameVersionV1 {
		if len(payload) != recordFrameHeaderSize+8 {
			return recordFrame{}, errRecordFrameInvalid
		}
		return recordFrame{Version: version, Type: frameType, FinalSeq: binary.BigEndian.Uint64(payload[8:16])}, nil
	}
	switch version {
	case recordFrameVersionV1:
		switch frameType {
		case recordFrameTypeBatch:
			return decodeRecordBatchFrameV1(payload)
		case recordFrameTypeAck:
			if len(payload) != recordFrameHeaderSize+8 {
				return recordFrame{}, errRecordFrameInvalid
			}
			return recordFrame{
				Version: recordFrameVersionV1,
				Type:    recordFrameTypeAck,
				AckSeq:  binary.BigEndian.Uint64(payload[8:16]),
			}, nil
		case recordFrameTypeError:
			frame, err := decodeRecordErrorFrame(payload)
			if err != nil {
				return recordFrame{}, err
			}
			frame.Version = recordFrameVersionV1
			return frame, nil
		default:
			return recordFrame{}, errRecordFrameInvalid
		}
	case recordFrameVersionV2:
		switch frameType {
		case recordFrameTypeBatch:
			return decodeRecordBatchFrameV2(payload)
		case recordFrameTypeAck:
			if len(payload) != recordFrameHeaderSize+8 {
				return recordFrame{}, errRecordFrameInvalid
			}
			return recordFrame{
				Version: recordFrameVersionV2,
				Type:    recordFrameTypeAck,
				AckSeq:  binary.BigEndian.Uint64(payload[8:16]),
			}, nil
		case recordFrameTypeError:
			frame, err := decodeRecordErrorFrame(payload)
			if err != nil {
				return recordFrame{}, err
			}
			frame.Version = recordFrameVersionV2
			return frame, nil
		default:
			return recordFrame{}, errRecordFrameInvalid
		}
	default:
		return recordFrame{}, errRecordFrameInvalid
	}
}

func decodeRecordBatchFrameV1(payload []byte) (recordFrame, error) {
	if len(payload) < recordFrameHeaderSize+recordBatchHeaderV1Size {
		return recordFrame{}, errRecordFrameInvalid
	}
	count := int(binary.BigEndian.Uint16(payload[8:10]))
	firstSeq := binary.BigEndian.Uint64(payload[10:18])
	if count <= 0 || firstSeq == 0 || uint64(count-1) > ^uint64(0)-firstSeq {
		return recordFrame{}, errRecordFrameInvalid
	}
	offset := recordFrameHeaderSize + recordBatchHeaderV1Size
	if count > (len(payload)-offset)/4 {
		return recordFrame{}, errRecordFrameInvalid
	}
	batch := make([]recordOutboundMessage, 0, count)
	for index := 0; index < count; index++ {
		if len(payload)-offset < 4 {
			return recordFrame{}, errRecordFrameInvalid
		}
		wireItemLen := binary.BigEndian.Uint32(payload[offset : offset+4])
		offset += 4
		if uint64(wireItemLen) > uint64(len(payload)-offset) {
			return recordFrame{}, errRecordFrameInvalid
		}
		itemLen := int(wireItemLen)
		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{
		Version: recordFrameVersionV1,
		Type:    recordFrameTypeBatch,
		Batch:   batch,
	}, nil
}

func decodeRecordBatchFrameV2(payload []byte) (recordFrame, error) {
	if len(payload) < recordFrameHeaderSize+recordBatchHeaderV2Size {
		return recordFrame{}, errRecordFrameInvalid
	}
	count := int(binary.BigEndian.Uint16(payload[8:10]))
	firstSeq := binary.BigEndian.Uint64(payload[10:18])
	ackSeq := binary.BigEndian.Uint64(payload[18:26])
	if count <= 0 || firstSeq == 0 || uint64(count-1) > ^uint64(0)-firstSeq {
		return recordFrame{}, errRecordFrameInvalid
	}
	offset := recordFrameHeaderSize + recordBatchHeaderV2Size
	if count > (len(payload)-offset)/4 {
		return recordFrame{}, errRecordFrameInvalid
	}
	batch := make([]recordOutboundMessage, 0, count)
	for index := 0; index < count; index++ {
		if len(payload)-offset < 4 {
			return recordFrame{}, errRecordFrameInvalid
		}
		wireItemLen := binary.BigEndian.Uint32(payload[offset : offset+4])
		offset += 4
		if uint64(wireItemLen) > uint64(len(payload)-offset) {
			return recordFrame{}, errRecordFrameInvalid
		}
		itemLen := int(wireItemLen)
		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{
		Version: recordFrameVersionV2,
		Type:    recordFrameTypeBatch,
		Batch:   batch,
		AckSeq:  ackSeq,
	}, 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]))
	wireMsgLen := binary.BigEndian.Uint32(payload[18:22])
	offset := recordFrameHeaderSize + recordErrorHeaderSize
	if failedSeq == 0 || len(payload)-offset < codeLen {
		return recordFrame{}, errRecordFrameInvalid
	}
	if uint64(len(payload)-offset-codeLen) != uint64(wireMsgLen) {
		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
}