github.com/cloudwego/kitex@v0.9.0/pkg/remote/codec/protobuf/protobuf.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package protobuf

import (
	"context"
	"errors"
	"fmt"

	"github.com/cloudwego/fastpb"

	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec"
	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
)

/**
 *  Kitex Protobuf Protocol
 *  |----------Len--------|--------------------------------MetaInfo--------------------------------|
 *  |---------4Byte-------|----2Byte----|----2Byte----|---------String-------|---------4Byte-------|
 *	+----------------------------------------------------------------------------------------------+
 *	|      PayloadLen     |    Magic    |   MsgType   |      MethodName      |        SeqID        |
 *	+----------------------------------------------------------------------------------------------+
 *	|  									 												           |
 *	|                         Protobuf  Argument/Result/Error   			                       |
 *	|   							 													           |
 *	+----------------------------------------------------------------------------------------------+
 */

const (
	metaInfoFixLen = 8
)

// NewProtobufCodec ...
func NewProtobufCodec() remote.PayloadCodec {
	return &protobufCodec{}
}

// protobufCodec implements  PayloadMarshaler
type protobufCodec struct{}

// Len encode outside not here
func (c protobufCodec) Marshal(ctx context.Context, message remote.Message, out remote.ByteBuffer) error {
	// 1. prepare info
	methodName := message.RPCInfo().Invocation().MethodName()
	if methodName == "" {
		return errors.New("empty methodName in protobuf Marshal")
	}
	data, err := getValidData(methodName, message)
	if err != nil {
		return err
	}

	// 3. encode metainfo
	// 3.1 magic && msgType
	if err := codec.WriteUint32(codec.ProtobufV1Magic+uint32(message.MessageType()), out); err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("protobuf marshal, write meta info failed: %s", err.Error()))
	}
	// 3.2 methodName
	if _, err := codec.WriteString(methodName, out); err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("protobuf marshal, write method name failed: %s", err.Error()))
	}
	// 3.3 seqID
	if err := codec.WriteUint32(uint32(message.RPCInfo().Invocation().SeqID()), out); err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("protobuf marshal, write seqID failed: %s", err.Error()))
	}

	// 4. write actual message buf
	msg, ok := data.(ProtobufMsgCodec)
	if !ok {
		// If Using Generics
		// if data is a MessageWriterWithContext
		// Do msg.WritePb(ctx context.Context, out remote.ByteBuffer)
		genmsg, isgen := data.(MessageWriterWithContext)
		if isgen {
			actualMsg, err := genmsg.WritePb(ctx)
			if err != nil {
				return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("protobuf marshal message failed: %s", err.Error()))
			}
			actualMsgBuf, ok := actualMsg.([]byte)
			if !ok {
				return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("protobuf marshal message failed: %s", err.Error()))
			}
			_, err = out.WriteBinary(actualMsgBuf)
			if err != nil {
				return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("protobuf marshal, write message buffer failed: %s", err.Error()))
			}
			return nil
		}
		// return error otherwise
		return remote.NewTransErrorWithMsg(remote.InvalidProtocol, "encode failed, codec msg type not match with protobufCodec")
	}

	// 2. encode pb struct
	// fast write
	if msg, ok := data.(fastpb.Writer); ok {
		msgsize := msg.Size()
		actualMsgBuf, err := out.Malloc(msgsize)
		if err != nil {
			return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("protobuf malloc size %d failed: %s", msgsize, err.Error()))
		}
		msg.FastWrite(actualMsgBuf)
		return nil
	}

	var actualMsgBuf []byte
	if actualMsgBuf, err = msg.Marshal(actualMsgBuf); err != nil {
		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("protobuf marshal message failed: %s", err.Error()))
	}
	if _, err = out.WriteBinary(actualMsgBuf); err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("protobuf marshal, write message buffer failed: %s", err.Error()))
	}
	return nil
}

func (c protobufCodec) Unmarshal(ctx context.Context, message remote.Message, in remote.ByteBuffer) error {
	payloadLen := message.PayloadLen()
	magicAndMsgType, err := codec.ReadUint32(in)
	if err != nil {
		return err
	}
	if magicAndMsgType&codec.MagicMask != codec.ProtobufV1Magic {
		return perrors.NewProtocolErrorWithType(perrors.BadVersion, "Bad version in protobuf Unmarshal")
	}
	msgType := magicAndMsgType & codec.FrontMask
	if err := codec.UpdateMsgType(msgType, message); err != nil {
		return err
	}

	methodName, methodFieldLen, err := codec.ReadString(in)
	if err != nil {
		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("protobuf unmarshal, read method name failed: %s", err.Error()))
	}
	if err = codec.SetOrCheckMethodName(methodName, message); err != nil && msgType != uint32(remote.Exception) {
		return err
	}
	seqID, err := codec.ReadUint32(in)
	if err != nil {
		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("protobuf unmarshal, read seqID failed: %s", err.Error()))
	}
	if err = codec.SetOrCheckSeqID(int32(seqID), message); err != nil && msgType != uint32(remote.Exception) {
		return err
	}
	actualMsgLen := payloadLen - metaInfoFixLen - methodFieldLen
	actualMsgBuf, err := in.Next(actualMsgLen)
	if err != nil {
		return perrors.NewProtocolErrorWithErrMsg(err, fmt.Sprintf("protobuf unmarshal, read message buffer failed: %s", err.Error()))
	}
	// exception message
	if message.MessageType() == remote.Exception {
		var exception pbError
		if err := exception.Unmarshal(actualMsgBuf); err != nil {
			return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("protobuf unmarshal Exception failed: %s", err.Error()))
		}
		return remote.NewTransError(exception.TypeID(), &exception)
	}

	if err = codec.NewDataIfNeeded(methodName, message); err != nil {
		return err
	}
	data := message.Data()

	// fast read
	if msg, ok := data.(fastpb.Reader); ok {
		if len(actualMsgBuf) == 0 {
			// if all fields of a struct is default value, actualMsgLen will be zero and actualMsgBuf will be nil
			// In the implementation of fastpb, if actualMsgBuf is nil, then fastpb will skip creating this struct, as a result user will get a nil pointer which is not expected.
			// So, when actualMsgBuf is nil, use default protobuf unmarshal method to decode the struct.
			// todo: fix fastpb
		} else {
			_, err := fastpb.ReadMessage(actualMsgBuf, fastpb.SkipTypeCheck, msg)
			if err != nil {
				return remote.NewTransErrorWithMsg(remote.ProtocolError, err.Error())
			}
			return nil
		}
	}

	// JSONPB Generic Case
	if msg, ok := data.(MessageReaderWithMethodWithContext); ok {
		err := msg.ReadPb(ctx, methodName, actualMsgBuf)
		if err != nil {
			return err
		}
		return nil
	}

	msg, ok := data.(ProtobufMsgCodec)
	if !ok {
		return remote.NewTransErrorWithMsg(remote.InvalidProtocol, "decode failed, codec msg type not match with protobufCodec")
	}
	if err = msg.Unmarshal(actualMsgBuf); err != nil {
		return remote.NewTransErrorWithMsg(remote.ProtocolError, err.Error())
	}
	return err
}

func (c protobufCodec) Name() string {
	return "protobuf"
}

// MessageWriterWithContext  writes to output bytebuffer
type MessageWriterWithContext interface {
	WritePb(ctx context.Context) (interface{}, error)
}

// MessageReaderWithMethodWithContext read from ActualMsgBuf with method
type MessageReaderWithMethodWithContext interface {
	ReadPb(ctx context.Context, method string, in []byte) error
}

type ProtobufMsgCodec interface {
	Marshal(out []byte) ([]byte, error)
	Unmarshal(in []byte) error
}

func getValidData(methodName string, message remote.Message) (interface{}, error) {
	if err := codec.NewDataIfNeeded(methodName, message); err != nil {
		return nil, err
	}
	data := message.Data()
	if message.MessageType() != remote.Exception {
		return data, nil
	}
	transErr, isTransErr := data.(*remote.TransError)
	if !isTransErr {
		if err, isError := data.(error); isError {
			encodeErr := NewPbError(remote.InternalError, err.Error())
			return encodeErr, nil
		}
		return nil, errors.New("exception relay need error type data")
	}
	encodeErr := NewPbError(transErr.TypeID(), transErr.Error())
	return encodeErr, nil
}