github.com/cloudwego/kitex@v0.9.0/pkg/remote/codec/header_codec_test.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 codec

import (
	"context"
	"encoding/binary"
	"net"
	"testing"

	"github.com/bytedance/gopkg/cloud/metainfo"

	"github.com/cloudwego/kitex/internal/mocks"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/discovery"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
	"github.com/cloudwego/kitex/pkg/remote/transmeta"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	tm "github.com/cloudwego/kitex/pkg/transmeta"
	"github.com/cloudwego/kitex/transport"
)

var mockPayloadLen = 100

func TestTTHeaderCodec(t *testing.T) {
	ctx := context.Background()
	sendMsg := initClientSendMsg(transport.TTHeader)

	// encode
	out := remote.NewWriterBuffer(256)
	totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
	binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
	test.Assert(t, err == nil, err)

	// decode
	recvMsg := initServerRecvMsg()
	buf, err := out.Bytes()
	test.Assert(t, err == nil, err)
	in := remote.NewReaderBuffer(buf)
	err = ttHeaderCodec.decode(ctx, recvMsg, in)
	test.Assert(t, err == nil, err)
	test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())
}

func TestTTHeaderCodecWithTransInfo(t *testing.T) {
	ctx := context.Background()
	intKVInfo := prepareIntKVInfo()
	strKVInfo := prepareStrKVInfo()
	sendMsg := initClientSendMsg(transport.TTHeader)
	sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
	sendMsg.TransInfo().PutTransStrInfo(strKVInfo)
	sendMsg.Tags()[HeaderFlagsKey] = HeaderFlagSupportOutOfOrder

	// encode
	out := remote.NewWriterBuffer(256)
	totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
	binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
	test.Assert(t, err == nil, err)

	// decode
	recvMsg := initServerRecvMsg()
	buf, err := out.Bytes()
	test.Assert(t, err == nil, err)
	in := remote.NewReaderBuffer(buf)
	err = ttHeaderCodec.decode(ctx, recvMsg, in)
	test.Assert(t, err == nil, err)
	test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

	intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
	strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
	test.DeepEqual(t, intKVInfoRecv, intKVInfo)
	test.DeepEqual(t, strKVInfoRecv, strKVInfo)
	flag := recvMsg.Tags()[HeaderFlagsKey]
	test.Assert(t, flag != nil)
	test.Assert(t, flag == HeaderFlagSupportOutOfOrder)
}

func TestTTHeaderCodecWithTransInfoWithGDPRToken(t *testing.T) {
	ctx := context.Background()
	intKVInfo := prepareIntKVInfo()
	strKVInfo := prepareStrKVInfoWithGDPRToken()
	sendMsg := initClientSendMsg(transport.TTHeader)
	sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
	sendMsg.TransInfo().PutTransStrInfo(strKVInfo)
	sendMsg.Tags()[HeaderFlagsKey] = HeaderFlagSupportOutOfOrder

	// encode
	out := remote.NewWriterBuffer(256)
	totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
	binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
	test.Assert(t, err == nil, err)

	// decode
	recvMsg := initServerRecvMsg()
	buf, err := out.Bytes()
	test.Assert(t, err == nil, err)
	in := remote.NewReaderBuffer(buf)
	err = ttHeaderCodec.decode(ctx, recvMsg, in)
	test.Assert(t, err == nil, err)
	test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

	intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
	strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
	test.DeepEqual(t, intKVInfoRecv, intKVInfo)
	test.DeepEqual(t, strKVInfoRecv, strKVInfo)
	flag := recvMsg.Tags()[HeaderFlagsKey]
	test.Assert(t, flag != nil)
	test.Assert(t, flag == HeaderFlagSupportOutOfOrder)
}

func TestTTHeaderCodecWithTransInfoFromMetaInfoGDPRToken(t *testing.T) {
	ctx := context.Background()
	intKVInfo := prepareIntKVInfo()
	ctx = metainfo.WithValue(ctx, "gdpr-token", "test token")
	sendMsg := initClientSendMsg(transport.TTHeader)
	sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
	ctx, err := tm.MetainfoClientHandler.WriteMeta(ctx, sendMsg)
	test.Assert(t, err == nil)
	sendMsg.Tags()[HeaderFlagsKey] = HeaderFlagSupportOutOfOrder

	// encode
	out := remote.NewWriterBuffer(256)
	totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
	binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
	test.Assert(t, err == nil, err)

	// decode
	recvMsg := initServerRecvMsg()
	buf, err := out.Bytes()
	test.Assert(t, err == nil, err)
	in := remote.NewReaderBuffer(buf)
	err = ttHeaderCodec.decode(ctx, recvMsg, in)
	test.Assert(t, err == nil, err)
	test.Assert(t, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

	intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
	strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
	test.DeepEqual(t, intKVInfoRecv, intKVInfo)
	test.DeepEqual(t, strKVInfoRecv, map[string]string{transmeta.GDPRToken: "test token"})
	flag := recvMsg.Tags()[HeaderFlagsKey]
	test.Assert(t, flag != nil)
	test.Assert(t, flag == HeaderFlagSupportOutOfOrder)
}

func TestFillBasicInfoOfTTHeader(t *testing.T) {
	ctx := context.Background()
	mockAddr := "mock address"
	// 1. server side fill from address
	// encode
	sendMsg := initClientSendMsg(transport.TTHeader)
	sendMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] = mockAddr
	sendMsg.TransInfo().TransIntInfo()[transmeta.FromService] = mockServiceName
	out := remote.NewWriterBuffer(256)
	totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
	binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
	test.Assert(t, err == nil, err)
	// decode
	recvMsg := initServerRecvMsg()
	buf, err := out.Bytes()
	test.Assert(t, err == nil, err)
	in := remote.NewReaderBuffer(buf)
	err = ttHeaderCodec.decode(ctx, recvMsg, in)
	test.Assert(t, err == nil, err)
	test.Assert(t, recvMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] == mockAddr)
	test.Assert(t, recvMsg.RPCInfo().From().Address().String() == mockAddr)
	test.Assert(t, recvMsg.RPCInfo().From().ServiceName() == mockServiceName)

	// 2. client side fill to address
	// encode
	sendMsg = initServerSendMsg(transport.TTHeader)
	sendMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] = mockAddr
	out = remote.NewWriterBuffer(256)
	totalLenField, err = ttHeaderCodec.encode(ctx, sendMsg, out)
	binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
	test.Assert(t, err == nil, err)
	// decode
	recvMsg = initClientRecvMsg()
	toInfo := remoteinfo.AsRemoteInfo(recvMsg.RPCInfo().To())
	toInfo.SetInstance(&mockInst{})
	buf, err = out.Bytes()
	test.Assert(t, err == nil, err)
	in = remote.NewReaderBuffer(buf)
	err = ttHeaderCodec.decode(ctx, recvMsg, in)
	test.Assert(t, err == nil, err)
	test.Assert(t, recvMsg.TransInfo().TransStrInfo()[transmeta.HeaderTransRemoteAddr] == mockAddr)
	test.Assert(t, toInfo.Address().String() == mockAddr, toInfo.Address())
}

func BenchmarkTTHeaderCodec(b *testing.B) {
	ctx := context.Background()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		sendMsg := initClientSendMsg(transport.TTHeader)
		// encode
		out := remote.NewWriterBuffer(256)
		totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
		binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
		test.Assert(b, err == nil, err)

		// decode
		recvMsg := initServerRecvMsg()
		buf, err := out.Bytes()
		test.Assert(b, err == nil, err)
		in := remote.NewReaderBuffer(buf)
		err = ttHeaderCodec.decode(ctx, recvMsg, in)
		test.Assert(b, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())
		test.Assert(b, err == nil, err)
	}
}

func BenchmarkTTHeaderWithTransInfoParallel(b *testing.B) {
	ctx := context.Background()
	intKVInfo := prepareIntKVInfo()
	strKVInfo := prepareStrKVInfo()

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			sendMsg := initClientSendMsg(transport.TTHeader)
			sendMsg.TransInfo().PutTransIntInfo(intKVInfo)
			sendMsg.TransInfo().PutTransStrInfo(strKVInfo)
			sendMsg.Tags()[HeaderFlagsKey] = HeaderFlagSupportOutOfOrder

			// encode
			out := remote.NewWriterBuffer(256)
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
			binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
			test.Assert(b, err == nil, err)

			// decode
			recvMsg := initServerRecvMsg()
			buf, err := out.Bytes()
			test.Assert(b, err == nil, err)
			in := remote.NewReaderBuffer(buf)
			err = ttHeaderCodec.decode(ctx, recvMsg, in)
			test.Assert(b, err == nil, err)
			test.Assert(b, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())

			intKVInfoRecv := recvMsg.TransInfo().TransIntInfo()
			strKVInfoRecv := recvMsg.TransInfo().TransStrInfo()
			test.DeepEqual(b, intKVInfoRecv, intKVInfo)
			test.DeepEqual(b, strKVInfoRecv, strKVInfo)
			flag := recvMsg.Tags()[HeaderFlagsKey]
			test.Assert(b, flag != nil)
			test.Assert(b, flag == HeaderFlagSupportOutOfOrder)
		}
	})
}

func BenchmarkTTHeaderCodecParallel(b *testing.B) {
	ctx := context.Background()

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			sendMsg := initClientSendMsg(transport.TTHeader)
			// encode
			out := remote.NewWriterBuffer(256)
			totalLenField, err := ttHeaderCodec.encode(ctx, sendMsg, out)
			binary.BigEndian.PutUint32(totalLenField, uint32(out.MallocLen()-Size32+mockPayloadLen))
			test.Assert(b, err == nil, err)

			// decode
			recvMsg := initServerRecvMsg()
			buf, err := out.Bytes()
			test.Assert(b, err == nil, err)
			in := remote.NewReaderBuffer(buf)
			err = ttHeaderCodec.decode(ctx, recvMsg, in)
			test.Assert(b, err == nil, err)
			test.Assert(b, recvMsg.PayloadLen() == mockPayloadLen, recvMsg.PayloadLen())
		}
	})
}

var (
	mockServiceName = "mock service"
	mockMethod      = "mock"

	mockCliRPCInfo = rpcinfo.NewRPCInfo(
		rpcinfo.EmptyEndpointInfo(),
		remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{ServiceName: mockServiceName}, mockMethod).ImmutableView(),
		rpcinfo.NewInvocation("", mockMethod),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())

	mockSvrRPCInfo = rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(),
		rpcinfo.FromBasicInfo(&rpcinfo.EndpointBasicInfo{ServiceName: mockServiceName}),
		rpcinfo.NewServerInvocation(),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
)

func initServerRecvMsg() remote.Message {
	svcInfo := mocks.ServiceInfo()
	svcSearchMap := map[string]*serviceinfo.ServiceInfo{
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockMethod):          svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockExceptionMethod): svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockErrorMethod):     svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockOnewayMethod):    svcInfo,
		mocks.MockMethod:          svcInfo,
		mocks.MockExceptionMethod: svcInfo,
		mocks.MockErrorMethod:     svcInfo,
		mocks.MockOnewayMethod:    svcInfo,
	}
	msg := remote.NewMessageWithNewer(svcInfo, svcSearchMap, mockSvrRPCInfo, remote.Call, remote.Server, false)
	return msg
}

func initClientSendMsg(tp transport.Protocol) remote.Message {
	var req interface{}
	svcInfo := mocks.ServiceInfo()
	msg := remote.NewMessage(req, svcInfo, mockCliRPCInfo, remote.Call, remote.Client)
	msg.SetProtocolInfo(remote.NewProtocolInfo(tp, svcInfo.PayloadCodec))
	return msg
}

func initServerSendMsg(tp transport.Protocol) remote.Message {
	var resp interface{}
	msg := remote.NewMessage(resp, mocks.ServiceInfo(), mockSvrRPCInfo, remote.Reply, remote.Server)
	msg.SetProtocolInfo(remote.NewProtocolInfo(tp, mocks.ServiceInfo().PayloadCodec))
	return msg
}

func initClientRecvMsg() remote.Message {
	var resp interface{}
	svcInfo := mocks.ServiceInfo()
	msg := remote.NewMessage(resp, svcInfo, mockCliRPCInfo, remote.Reply, remote.Client)
	return msg
}

var _ discovery.Instance = &mockInst{}

type mockInst struct {
	addr net.Addr
}

func (m *mockInst) Address() net.Addr {
	return m.addr
}

func (m *mockInst) RefreshInstanceWithAddr(addr net.Addr) discovery.Instance {
	m.addr = addr
	return m
}

func (m *mockInst) Weight() int {
	return 10
}

func (m *mockInst) Tag(key string) (value string, exist bool) {
	return
}

func prepareIntKVInfo() map[uint16]string {
	kvInfo := map[uint16]string{
		transmeta.FromService: "mockFromService",
		transmeta.FromMethod:  "mockFromMethod",
		transmeta.ToService:   "mockToService",
		transmeta.ToMethod:    "mockToMethod",
	}
	return kvInfo
}

func prepareStrKVInfo() map[string]string {
	kvInfo := map[string]string{transmeta.HeaderIDLServiceName: mocks.MockServiceName}
	return kvInfo
}

func prepareStrKVInfoWithGDPRToken() map[string]string {
	kvInfo := map[string]string{
		transmeta.GDPRToken:             "mockToken",
		transmeta.HeaderTransRemoteAddr: "mockRemoteAddr",
		transmeta.HeaderIDLServiceName:  mocks.MockServiceName,
	}
	return kvInfo
}

// // TODO 是否提供buf.writeInt8/16/32方法，否则得先计算，然后malloc，最后write，待确认频繁malloc是否有影响
// 暂时不删除，测试一次malloc, 和多次malloc差异
//
//lint:ignore U1000 until encode2 is used
func (t ttHeader) encode2(ctx context.Context, message remote.Message, payloadBuf, out remote.ByteBuffer) error {
	tm := message.TransInfo()

	// 1. header meta
	headerMeta, err := out.Malloc(TTHeaderMetaSize)
	if err != nil {
		return perrors.NewProtocolError(err)
	}

	totalLenField := headerMeta[0:4]
	headerSizeField := headerMeta[12:14]

	binary.BigEndian.PutUint32(headerMeta[4:8], TTHeaderMagic+uint32(getFlags(message)))
	binary.BigEndian.PutUint32(headerMeta[8:12], uint32(message.RPCInfo().Invocation().SeqID()))

	var transformIDs []uint8
	// 2. header info, calculate size
	var headerInfo []byte
	var headerInfoSize int
	// PROTOCOL ID(u8) + NUM TRANSFORMS(always 0)(u8) + TRANSFORM IDs([]u8)
	headerInfoSize = 1 + 1 + len(transformIDs)
	// str kv info
	if len(tm.TransStrInfo()) > 0 {
		headerInfoSize += 3 // INFO ID TYPE(u8) + NUM HEADERS(u16)
		for k, v := range tm.TransStrInfo() {
			headerInfoSize += 4 + len(k) + len(v)
		}
	}
	// int kv info
	if len(tm.TransIntInfo()) > 0 {
		headerInfoSize += 3 // INFO ID TYPE(u8) + NUM HEADERS(u16)
		for _, v := range tm.TransIntInfo() {
			headerInfoSize += 4 + len(v) // key is uint16
		}
	}
	// padding blank
	padding := (4 - headerInfoSize%4) % 4
	headerInfoSize += padding
	binary.BigEndian.PutUint16(headerSizeField, uint16(headerInfoSize/4))
	totalLen := TTHeaderMetaSize - Size32 + headerInfoSize + payloadBuf.MallocLen()
	binary.BigEndian.PutUint32(totalLenField, uint32(totalLen))

	// 3.  header info, malloc and write
	if headerInfo, err = out.Malloc(headerInfoSize); err != nil {
		return perrors.NewProtocolError(err)
	}
	headerInfo[0] = byte(getProtocolID(message.ProtocolInfo()))
	headerInfo[1] = byte(len(transformIDs))
	hdIdx := 2
	for tid := range transformIDs {
		// TODO 需确认下transformId的编码格式
		headerInfo[hdIdx] = byte(tid)
		hdIdx++
	}
	// str kv info
	if len(tm.TransStrInfo()) > 0 {
		// INFO ID TYPE(u8) + NUM HEADERS(u16)
		headerInfo[hdIdx] = byte(InfoIDKeyValue)
		binary.BigEndian.PutUint16(headerInfo[hdIdx+1:hdIdx+3], uint16(len(tm.TransStrInfo())))
		hdIdx += 3
		for key, value := range tm.TransStrInfo() {
			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], uint16(len(key)))
			hdIdx += 2
			hdIdx += copy(headerInfo[hdIdx:], key)

			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], uint16(len(value)))
			hdIdx += 2
			hdIdx += copy(headerInfo[hdIdx:], value)
		}
	}
	// int kv info
	if len(tm.TransIntInfo()) > 0 {
		// INFO ID TYPE(u8) + NUM HEADERS(u16)
		headerInfo[hdIdx] = byte(InfoIDIntKeyValue)
		binary.BigEndian.PutUint16(headerInfo[hdIdx+1:hdIdx+3], uint16(len(tm.TransIntInfo())))
		hdIdx += 3
		for key, value := range tm.TransIntInfo() {
			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], key)
			hdIdx += 2

			binary.BigEndian.PutUint16(headerInfo[hdIdx:hdIdx+2], uint16(len(value)))
			hdIdx += 2
			hdIdx += copy(headerInfo[hdIdx:], value)
		}
	}
	// padding = (4 - headerInfoSize%4) % 4
	for padding := (4 - hdIdx%4) % 4; padding > 0; padding-- {
		headerInfo[hdIdx] = byte(0)
		hdIdx++
	}
	if hdIdx != headerInfoSize {
		return perrors.NewProtocolErrorWithMsg("the size that header info write is not equal with malloc size")
	}

	if err := out.AppendBuffer(payloadBuf); err != nil {
		return perrors.NewProtocolError(err)
	}
	return err
}

func TestHeaderFlags(t *testing.T) {
	// case 1: correct HeaderFlags
	msg := initClientSendMsg(transport.TTHeader)
	msg.Tags()[HeaderFlagsKey] = HeaderFlagSASL | HeaderFlagSupportOutOfOrder
	hfs := getFlags(msg)
	test.Assert(t, hfs == HeaderFlagSASL|HeaderFlagSupportOutOfOrder)

	// case 2: invalid type for HeaderFlagsKey
	msg = initClientSendMsg(transport.TTHeader)
	msg.Tags()[HeaderFlagsKey] = 1
	hfs = getFlags(msg)
	test.Assert(t, hfs == 0)

	// case 3: setFlags then get Flags
	msg = initClientSendMsg(transport.TTHeader)
	setFlags(uint16(HeaderFlagSupportOutOfOrder), msg)
	hfs = getFlags(msg)
	test.Assert(t, hfs == HeaderFlagSupportOutOfOrder, hfs)
}