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 }