github.com/cloudwego/kitex@v0.9.0/pkg/transmeta/http2_test.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  // Package transmeta metadata handler for translation
    18  package transmeta
    19  
    20  import (
    21  	"context"
    22  	"testing"
    23  
    24  	"github.com/cloudwego/kitex/internal/test"
    25  	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"
    26  	"github.com/cloudwego/kitex/pkg/remote/transmeta"
    27  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    28  
    29  	"github.com/cloudwego/kitex/transport"
    30  )
    31  
    32  func TestIsGRPC(t *testing.T) {
    33  	type args struct {
    34  		ri rpcinfo.RPCInfo
    35  	}
    36  	tests := []struct {
    37  		name string
    38  		args args
    39  		want bool
    40  	}{
    41  		{"with ttheader", args{ri: rpcinfo.NewRPCInfo(
    42  			nil,
    43  			nil,
    44  			nil,
    45  			func() rpcinfo.RPCConfig {
    46  				cfg := rpcinfo.NewRPCConfig()
    47  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader)
    48  				return cfg
    49  			}(),
    50  			nil,
    51  		)}, false},
    52  
    53  		{"with ttheader framed", args{ri: rpcinfo.NewRPCInfo(
    54  			nil,
    55  			nil,
    56  			nil,
    57  			func() rpcinfo.RPCConfig {
    58  				cfg := rpcinfo.NewRPCConfig()
    59  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeaderFramed)
    60  				return cfg
    61  			}(),
    62  			nil,
    63  		)}, false},
    64  
    65  		{"with grpc", args{ri: rpcinfo.NewRPCInfo(
    66  			nil,
    67  			nil,
    68  			nil,
    69  			func() rpcinfo.RPCConfig {
    70  				cfg := rpcinfo.NewRPCConfig()
    71  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
    72  				return cfg
    73  			}(),
    74  			nil,
    75  		)}, true},
    76  
    77  		{"with http", args{ri: rpcinfo.NewRPCInfo(
    78  			nil,
    79  			nil,
    80  			nil,
    81  			func() rpcinfo.RPCConfig {
    82  				cfg := rpcinfo.NewRPCConfig()
    83  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
    84  				return cfg
    85  			}(),
    86  			nil,
    87  		)}, false},
    88  
    89  		{"with ttheader and http", args{ri: rpcinfo.NewRPCInfo(
    90  			nil,
    91  			nil,
    92  			nil,
    93  			func() rpcinfo.RPCConfig {
    94  				cfg := rpcinfo.NewRPCConfig()
    95  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader)
    96  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
    97  				return cfg
    98  			}(),
    99  			nil,
   100  		)}, false},
   101  
   102  		{"with ttheader framed and http", args{ri: rpcinfo.NewRPCInfo(
   103  			nil,
   104  			nil,
   105  			nil,
   106  			func() rpcinfo.RPCConfig {
   107  				cfg := rpcinfo.NewRPCConfig()
   108  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeaderFramed)
   109  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
   110  				return cfg
   111  			}(),
   112  			nil,
   113  		)}, false},
   114  
   115  		{"with ttheader and ttheader framed", args{ri: rpcinfo.NewRPCInfo(
   116  			nil,
   117  			nil,
   118  			nil,
   119  			func() rpcinfo.RPCConfig {
   120  				cfg := rpcinfo.NewRPCConfig()
   121  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader)
   122  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeaderFramed)
   123  				return cfg
   124  			}(),
   125  			nil,
   126  		)}, false},
   127  
   128  		{"with http and grpc", args{ri: rpcinfo.NewRPCInfo(
   129  			nil,
   130  			nil,
   131  			nil,
   132  			func() rpcinfo.RPCConfig {
   133  				cfg := rpcinfo.NewRPCConfig()
   134  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
   135  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   136  				return cfg
   137  			}(),
   138  			nil,
   139  		)}, true},
   140  
   141  		{"with ttheader and grpc", args{ri: rpcinfo.NewRPCInfo(
   142  			nil,
   143  			nil,
   144  			nil,
   145  			func() rpcinfo.RPCConfig {
   146  				cfg := rpcinfo.NewRPCConfig()
   147  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader)
   148  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   149  				return cfg
   150  			}(),
   151  			nil,
   152  		)}, true},
   153  
   154  		{"with ttheader framed and grpc", args{ri: rpcinfo.NewRPCInfo(
   155  			nil,
   156  			nil,
   157  			nil,
   158  			func() rpcinfo.RPCConfig {
   159  				cfg := rpcinfo.NewRPCConfig()
   160  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeaderFramed)
   161  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   162  				return cfg
   163  			}(),
   164  			nil,
   165  		)}, true},
   166  
   167  		{"with ttheader,http and grpc", args{ri: rpcinfo.NewRPCInfo(
   168  			nil,
   169  			nil,
   170  			nil,
   171  			func() rpcinfo.RPCConfig {
   172  				cfg := rpcinfo.NewRPCConfig()
   173  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader)
   174  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
   175  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   176  				return cfg
   177  			}(),
   178  			nil,
   179  		)}, true},
   180  
   181  		{"with ttheader framed,http and grpc", args{ri: rpcinfo.NewRPCInfo(
   182  			nil,
   183  			nil,
   184  			nil,
   185  			func() rpcinfo.RPCConfig {
   186  				cfg := rpcinfo.NewRPCConfig()
   187  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeaderFramed)
   188  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
   189  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   190  				return cfg
   191  			}(),
   192  			nil,
   193  		)}, true},
   194  
   195  		{"with ttheader ,ttheader framed and http", args{ri: rpcinfo.NewRPCInfo(
   196  			nil,
   197  			nil,
   198  			nil,
   199  			func() rpcinfo.RPCConfig {
   200  				cfg := rpcinfo.NewRPCConfig()
   201  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader)
   202  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeaderFramed)
   203  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
   204  				return cfg
   205  			}(),
   206  			nil,
   207  		)}, false},
   208  
   209  		{"with ttheader ,ttheader framed,http and grpc", args{ri: rpcinfo.NewRPCInfo(
   210  			nil,
   211  			nil,
   212  			nil,
   213  			func() rpcinfo.RPCConfig {
   214  				cfg := rpcinfo.NewRPCConfig()
   215  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeader)
   216  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.TTHeaderFramed)
   217  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.HTTP)
   218  				rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   219  				return cfg
   220  			}(),
   221  			nil,
   222  		)}, true},
   223  	}
   224  	for _, tt := range tests {
   225  		t.Run(tt.name, func(t *testing.T) {
   226  			if got := isGRPC(tt.args.ri); got != tt.want {
   227  				t.Errorf("isGRPC() = %v, want %v", got, tt.want)
   228  			}
   229  		})
   230  	}
   231  }
   232  
   233  func TestHTTP2ClientOnConnectStream(t *testing.T) {
   234  	cfg := rpcinfo.NewRPCConfig()
   235  	fromInfo := rpcinfo.NewEndpointInfo("fromServiceName", "fromMethod", nil, nil)
   236  	toInfo := rpcinfo.NewEndpointInfo("toServiceName", "toMethod", nil, nil)
   237  	ri := rpcinfo.NewRPCInfo(fromInfo, toInfo, rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
   238  	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
   239  
   240  	// not grpc
   241  	ctx, err := ClientHTTP2Handler.OnConnectStream(ctx)
   242  	test.Assert(t, err == nil)
   243  	md, exist := metadata.FromOutgoingContext(ctx)
   244  	test.Assert(t, !exist)
   245  	test.Assert(t, len(md) == 0)
   246  
   247  	// grpc
   248  	rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   249  	ctx, err = ClientHTTP2Handler.OnConnectStream(ctx)
   250  	test.Assert(t, err == nil)
   251  
   252  	md, exist = metadata.FromOutgoingContext(ctx)
   253  	test.Assert(t, exist)
   254  	test.Assert(t, md.Get(transmeta.HTTPSourceService)[0] == fromInfo.ServiceName())
   255  	test.Assert(t, md.Get(transmeta.HTTPSourceMethod)[0] == fromInfo.Method())
   256  	test.Assert(t, md.Get(transmeta.HTTPDestService)[0] == toInfo.ServiceName())
   257  	test.Assert(t, md.Get(transmeta.HTTPDestMethod)[0] == toInfo.Method())
   258  }
   259  
   260  func TestHTTP2ServerOnReadStream(t *testing.T) {
   261  	cfg := rpcinfo.NewRPCConfig()
   262  	ri := rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(), rpcinfo.EmptyEndpointInfo(), rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
   263  	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
   264  
   265  	var (
   266  		toServiceName   = "toServiceName"
   267  		toMethod        = "toMethod"
   268  		fromServiceName = "fromServiceName"
   269  		fromMethod      = "fromMethod"
   270  	)
   271  
   272  	// set incoming metadata
   273  	md := metadata.MD{}
   274  	md.Append(transmeta.HTTPDestService, toServiceName)
   275  	md.Append(transmeta.HTTPDestMethod, toMethod)
   276  	md.Append(transmeta.HTTPSourceService, fromServiceName)
   277  	md.Append(transmeta.HTTPSourceMethod, fromMethod)
   278  
   279  	// not grpc
   280  	ctx, err := ServerHTTP2Handler.OnReadStream(ctx)
   281  	test.Assert(t, err == nil)
   282  	test.Assert(t, rpcinfo.GetRPCInfo(ctx).From().ServiceName() == "")
   283  	test.Assert(t, rpcinfo.GetRPCInfo(ctx).From().Method() == "")
   284  
   285  	// grpc, no incoming context
   286  	rpcinfo.AsMutableRPCConfig(cfg).SetTransportProtocol(transport.GRPC)
   287  	ctx, err = ServerHTTP2Handler.OnReadStream(ctx)
   288  	test.Assert(t, err == nil)
   289  	test.Assert(t, rpcinfo.GetRPCInfo(ctx).From().ServiceName() == "")
   290  	test.Assert(t, rpcinfo.GetRPCInfo(ctx).From().Method() == "")
   291  
   292  	// grpc with incoming context
   293  	ctx = metadata.NewIncomingContext(ctx, md)
   294  	ctx, err = ServerHTTP2Handler.OnReadStream(ctx)
   295  	test.Assert(t, err == nil)
   296  
   297  	newRPCInfo := rpcinfo.GetRPCInfo(ctx)
   298  	test.Assert(t, newRPCInfo.From().ServiceName() == fromServiceName)
   299  	test.Assert(t, newRPCInfo.From().Method() == fromMethod)
   300  }