github.com/erda-project/erda-infra@v1.0.10-0.20240327085753-f3a249292aeb/pkg/transport/transport.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package transport
    16  
    17  import (
    18  	"context"
    19  	"net/http"
    20  	"strings"
    21  
    22  	transgrpc "github.com/erda-project/erda-infra/pkg/transport/grpc"
    23  	transhttp "github.com/erda-project/erda-infra/pkg/transport/http"
    24  	"github.com/erda-project/erda-infra/pkg/transport/interceptor"
    25  	"google.golang.org/grpc/metadata"
    26  )
    27  
    28  type (
    29  	// Register .
    30  	Register interface {
    31  		transhttp.Router
    32  		transgrpc.ServiceRegistrar
    33  	}
    34  	// ServiceOption .
    35  	ServiceOption func(*ServiceOptions)
    36  	// ServiceOptions .
    37  	ServiceOptions struct {
    38  		HTTP []transhttp.HandleOption
    39  		GRPC []transgrpc.HandleOption
    40  	}
    41  )
    42  
    43  // WithInterceptors .
    44  func WithInterceptors(o ...interceptor.Interceptor) ServiceOption {
    45  	return func(opts *ServiceOptions) {
    46  		if len(o) <= 0 {
    47  			return
    48  		}
    49  		inter := interceptor.Chain(o[0], o[1:]...)
    50  		opts.HTTP = append(opts.HTTP, transhttp.WithInterceptor(inter))
    51  		opts.GRPC = append(opts.GRPC, transgrpc.WithInterceptor(inter))
    52  	}
    53  }
    54  
    55  // WithHTTPOptions .
    56  func WithHTTPOptions(o ...transhttp.HandleOption) ServiceOption {
    57  	return func(opts *ServiceOptions) {
    58  		opts.HTTP = append(opts.HTTP, o...)
    59  	}
    60  }
    61  
    62  // DefaultServiceOptions .
    63  func DefaultServiceOptions() *ServiceOptions {
    64  	return &ServiceOptions{}
    65  }
    66  
    67  // ServiceInfo .
    68  type ServiceInfo interface {
    69  	Service() string
    70  	Method() string
    71  	Instance() interface{}
    72  }
    73  
    74  // NewServiceInfo .
    75  func NewServiceInfo(service string, method string, instance interface{}) ServiceInfo {
    76  	return &serviceInfo{
    77  		service:  service,
    78  		method:   method,
    79  		instance: instance,
    80  	}
    81  }
    82  
    83  type serviceInfoContextKey int8
    84  
    85  // ServiceInfoContextKey .
    86  const ServiceInfoContextKey = serviceInfoContextKey(0)
    87  
    88  // WithServiceInfo .
    89  func WithServiceInfo(ctx context.Context, info ServiceInfo) context.Context {
    90  	return context.WithValue(ctx, ServiceInfoContextKey, info)
    91  }
    92  
    93  // ContextServiceInfo .
    94  func ContextServiceInfo(ctx context.Context) ServiceInfo {
    95  	info, _ := ctx.Value(ServiceInfoContextKey).(ServiceInfo)
    96  	return info
    97  }
    98  
    99  // GetFullMethodName .
   100  func GetFullMethodName(ctx context.Context) string {
   101  	info, _ := ctx.Value(ServiceInfoContextKey).(ServiceInfo)
   102  	if info != nil {
   103  		return info.Service() + "/" + info.Method()
   104  	}
   105  	return ""
   106  }
   107  
   108  type serviceInfo struct {
   109  	service  string
   110  	method   string
   111  	instance interface{}
   112  }
   113  
   114  func (si *serviceInfo) Service() string       { return si.service }
   115  func (si *serviceInfo) Method() string        { return si.method }
   116  func (si *serviceInfo) Instance() interface{} { return si.instance }
   117  
   118  // Header .
   119  type Header = metadata.MD
   120  
   121  // WithHeader setup header for caller
   122  func WithHeader(ctx context.Context, header Header) context.Context {
   123  	if len(header) <= 0 {
   124  		return ctx
   125  	}
   126  	return metadata.NewOutgoingContext(ctx, header)
   127  }
   128  
   129  // ContextHeader get header in server
   130  func ContextHeader(ctx context.Context) Header {
   131  	md, _ := metadata.FromIncomingContext(ctx)
   132  	if md == nil {
   133  		md, _ = metadata.FromOutgoingContext(ctx)
   134  	}
   135  	return md
   136  }
   137  
   138  // WithHTTPHeaderForServer setup header for http server
   139  func WithHTTPHeaderForServer(ctx context.Context, header http.Header) context.Context {
   140  	if len(header) <= 0 {
   141  		return ctx
   142  	}
   143  	md := metadata.MD{}
   144  	for key, values := range header {
   145  		key = strings.ToLower(key)
   146  		md[key] = values
   147  	}
   148  	return metadata.NewIncomingContext(ctx, md)
   149  }