github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/detection/server_handler_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 detection

import (
	"context"
	"errors"
	"fmt"
	"net"
	"testing"

	"github.com/golang/mock/gomock"

	"github.com/cloudwego/kitex/internal/mocks"
	mocksklog "github.com/cloudwego/kitex/internal/mocks/klog"
	npmocks "github.com/cloudwego/kitex/internal/mocks/netpoll"
	remote_mocks "github.com/cloudwego/kitex/internal/mocks/remote"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec"
	"github.com/cloudwego/kitex/pkg/remote/trans/netpoll"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/grpc"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/utils"
)

var (
	prefaceReadAtMost = func() int {
		// min(len(ClientPreface), len(flagBuf))
		// len(flagBuf) = 2 * codec.Size32
		if 2*codec.Size32 < grpc.ClientPrefaceLen {
			return 2 * codec.Size32
		}
		return grpc.ClientPrefaceLen
	}()
	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,
	}
)

func TestServerHandlerCall(t *testing.T) {
	transHdler, _ := NewSvrTransHandlerFactory(netpoll.NewSvrTransHandlerFactory(), nphttp2.NewSvrTransHandlerFactory()).NewTransHandler(&remote.ServerOption{
		SvcSearchMap:  svcSearchMap,
		TargetSvcInfo: svcInfo,
	})

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	npConn := npmocks.NewMockConnection(ctrl)
	npReader := npmocks.NewMockReader(ctrl)
	hdl := remote_mocks.NewMockServerTransHandler(ctrl)

	errOnActive := errors.New("mock on active error")
	errOnRead := errors.New("mock on read error")

	triggerReadErr := false
	triggerActiveErr := false

	hdl.EXPECT().OnActive(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, conn net.Conn) (context.Context, error) {
		if triggerActiveErr {
			return ctx, errOnActive
		}
		return ctx, nil
	}).AnyTimes()
	hdl.EXPECT().OnRead(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, conn net.Conn) error {
		if triggerReadErr {
			return errOnRead
		}
		return nil
	}).AnyTimes()
	hdl.EXPECT().OnInactive(gomock.Any(), gomock.Any()).AnyTimes()
	hdl.EXPECT().OnError(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, err error, conn net.Conn) {
		if conn != nil {
			klog.CtxErrorf(ctx, "KITEX: processing error, remoteAddr=%v, error=%s", conn.RemoteAddr(), err.Error())
		} else {
			klog.CtxErrorf(ctx, "KITEX: processing error, error=%s", err.Error())
		}
	}).AnyTimes()

	npReader.EXPECT().Peek(prefaceReadAtMost).Return([]byte("connection prefix bytes"), nil).AnyTimes()
	npConn.EXPECT().Reader().Return(npReader).AnyTimes()
	npConn.EXPECT().RemoteAddr().Return(nil).AnyTimes()

	transHdler.(*svrTransHandler).defaultHandler = hdl

	// case1 successful call: onActive() and onRead() all success
	triggerActiveErr = false
	triggerReadErr = false
	err := mockCall(transHdler, npConn)
	test.Assert(t, err == nil, err)

	// case2 onActive failed: onActive() err and close conn
	triggerActiveErr = true
	triggerReadErr = false
	err = mockCall(transHdler, npConn)
	test.Assert(t, err == errOnActive, err)

	// case3 onRead failed: onRead() err and close conn
	triggerActiveErr = false
	triggerReadErr = true
	err = mockCall(transHdler, npConn)
	test.Assert(t, err == errOnRead, err)
}

func TestOnError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer func() {
		klog.SetLogger(klog.DefaultLogger())
		ctrl.Finish()
	}()
	transHdler, err := NewSvrTransHandlerFactory(netpoll.NewSvrTransHandlerFactory(), nphttp2.NewSvrTransHandlerFactory()).NewTransHandler(&remote.ServerOption{
		SvcSearchMap:  svcSearchMap,
		TargetSvcInfo: svcInfo,
	})
	test.Assert(t, err == nil)

	mocklogger := mocksklog.NewMockFullLogger(ctrl)
	klog.SetLogger(mocklogger)

	var errMsg string
	mocklogger.EXPECT().CtxErrorf(gomock.Any(), gomock.Any(), gomock.Any()).Do(func(ctx context.Context, format string, v ...interface{}) {
		errMsg = fmt.Sprintf(format, v...)
	})
	transHdler.OnError(context.Background(), errors.New("mock error"), nil)
	test.Assert(t, errMsg == "KITEX: processing error, error=mock error", errMsg)

	conn := &mocks.Conn{
		RemoteAddrFunc: func() (r net.Addr) {
			return utils.NewNetAddr("mock", "mock")
		},
	}
	mocklogger.EXPECT().CtxErrorf(gomock.Any(), gomock.Any(), gomock.Any()).Do(func(ctx context.Context, format string, v ...interface{}) {
		errMsg = fmt.Sprintf(format, v...)
	})
	transHdler.OnError(context.Background(), errors.New("mock error"), conn)
	test.Assert(t, errMsg == "KITEX: processing error, remoteAddr=mock, error=mock error", errMsg)
}

// TestOnInactive covers onInactive() codes to check panic
func TestOnInactive(t *testing.T) {
	transHdler, err := NewSvrTransHandlerFactory(netpoll.NewSvrTransHandlerFactory(), nphttp2.NewSvrTransHandlerFactory()).NewTransHandler(&remote.ServerOption{
		SvcSearchMap:  svcSearchMap,
		TargetSvcInfo: svcInfo,
	})
	test.Assert(t, err == nil)

	conn := &mocks.Conn{
		RemoteAddrFunc: func() (r net.Addr) {
			return utils.NewNetAddr("mock", "mock")
		},
	}

	// case1 test noopHandler onInactive()
	transHdler.OnInactive(context.Background(), conn)

	// mock a ctx and set handlerKey
	subHandler := &mocks.MockSvrTransHandler{}
	subHandlerCtx := context.WithValue(
		context.Background(),
		handlerKey{},
		&handlerWrapper{
			handler: subHandler,
		},
	)

	ctx := context.WithValue(
		context.Background(),
		handlerKey{},
		&handlerWrapper{
			ctx: subHandlerCtx,
		},
	)
	// case2 test subHandler onInactive()
	transHdler.OnInactive(ctx, conn)
}

// mockCall mocks how detection transHdlr processing with incoming requests
func mockCall(transHdlr remote.ServerTransHandler, conn net.Conn) (err error) {
	ctx := context.Background()
	// do onConnActive
	ctxWithHandler, err := transHdlr.OnActive(ctx, conn)
	// onActive failed, such as connections limitation
	if err != nil {
		transHdlr.OnError(ctx, err, conn)
		transHdlr.OnInactive(ctx, conn)
		return
	}
	// do onConnRead
	err = transHdlr.OnRead(ctxWithHandler, conn)
	if err != nil {
		transHdlr.OnError(ctxWithHandler, err, conn)
		transHdlr.OnInactive(ctxWithHandler, conn)
		return
	}
	// do onConnInactive
	transHdlr.OnInactive(ctxWithHandler, conn)
	return
}