Repository: cloudwego/kitex
File: pkg/generic/binarythrift_codec.go (Go)

Code before:

type binaryReqType = []byte

type binaryThriftCodec struct {
	thriftCodec remote.PayloadCodec
}

func (c *binaryThriftCodec) Marshal(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
	data := msg.Data()
	if data == nil {
		return perrors.NewProtocolErrorWithMsg("invalid marshal data in rawThriftBinaryCodec: nil")
	}
	if msg.MessageType() == remote.Exception {
		if err := c.thriftCodec.Marshal(ctx, msg, out); err != nil {
			return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("rawThriftBinaryCodec Marshal exception failed, err: %s", err.Error()))
		}
		return nil
	}
	var transBuff []byte
	var ok bool
	if msg.RPCRole() == remote.Server {
		// handle biz error
		if msg.RPCInfo().Invocation().BizStatusErr() != nil {
			sz := gthrift.Binary.MessageBeginLength(msg.RPCInfo().Invocation().MethodName())
			sz += gthrift.Binary.FieldStopLength()
			b, err := out.Malloc(sz)
			if err != nil {
				return perrors.NewProtocolError(fmt.Errorf("binary thrift generic marshal, remote.ByteBuffer Malloc err: %w", err))
			}
			b = gthrift.Binary.AppendMessageBegin(b[:0],
				msg.RPCInfo().Invocation().MethodName(), gthrift.TMessageType(msg.MessageType()), msg.RPCInfo().Invocation().SeqID())
			b = gthrift.Binary.AppendFieldStop(b)
			_ = b
			return nil
		}
		gResult := data.(*Result)
		transBinary := gResult.Success
		if transBuff, ok = transBinary.(binaryReqType); !ok {
			return perrors.NewProtocolErrorWithMsg("invalid marshal result in rawThriftBinaryCodec: must be []byte")
		}
	} else {
		gArg := data.(*Args)
		transBinary := gArg.Request
		if transBuff, ok = transBinary.(binaryReqType); !ok {
			return perrors.NewProtocolErrorWithMsg("invalid marshal request in rawThriftBinaryCodec: must be []byte")
		}
		if err := SetSeqID(msg.RPCInfo().Invocation().SeqID(), transBuff); err != nil {
			return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("rawThriftBinaryCodec set seqID failed, err: %s", err.Error()))
		}
	}
	out.WriteBinary(transBuff)

Review comment (suggestion):
Business error only work properly when using TTHeader and HTTP2 transmission protocols. If there is a business error, `data.(*Result).Success` will be nil, and an empty payload will be constructed here to return 

Code after:
	wb                     = thrift.NewWriteBinary()
)

type binaryReqType = []byte

type binaryThriftCodec struct {
	thriftCodec remote.PayloadCodec
}

func (c *binaryThriftCodec) Marshal(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
	data := msg.Data()
	if data == nil {
		return perrors.NewProtocolErrorWithMsg("invalid marshal data in rawThriftBinaryCodec: nil")
	}
	if msg.MessageType() == remote.Exception {
		if err := c.thriftCodec.Marshal(ctx, msg, out); err != nil {
			return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("rawThriftBinaryCodec Marshal exception failed, err: %s", err.Error()))
		}
		return nil
	}
	var transBuff []byte
	var ok bool
	if msg.RPCRole() == remote.Server {
		// Business error only works properly when using TTHeader and HTTP2 transmission protocols
		// If there is a business error, data.(*Result).Success will be nil, and an empty payload will be constructed here to return
		if msg.RPCInfo().Invocation().BizStatusErr() != nil {
			msg.Data().(WithCodec).SetCodec(wb)
			return thriftCodec.Marshal(ctx, msg, out)
		}
		gResult := data.(*Result)
		transBinary := gResult.Success
		if transBuff, ok = transBinary.(binaryReqType); !ok {
			return perrors.NewProtocolErrorWithMsg("invalid marshal result in rawThriftBinaryCodec: must be []byte")
		}
	} else {
		gArg := data.(*Args)
		transBinary := gArg.Request
		if transBuff, ok = transBinary.(binaryReqType); !ok {
			return perrors.NewProtocolErrorWithMsg("invalid marshal request in rawThriftBinaryCodec: must be []byte")
		}
		if err := SetSeqID(msg.RPCInfo().Invocation().SeqID(), transBuff); err != nil {
			return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("rawThriftBinaryCodec set seqID failed, err: %s", err.Error()))
		}
	}
	out.WriteBinary(transBuff)
	return nil
}

func (c *binaryThriftCodec) Unmarshal(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
	magicAndMsgType, err := codec.PeekUint32(in)
	if err != nil {