github.com/cloudwego/kitex@v0.9.0/pkg/utils/kitexutil/kitexutil.go (about)

     1  /*
     2   * Copyright 2023 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 kitexutil provides some util methods to get RPC information
    18  package kitexutil
    19  
    20  import (
    21  	"context"
    22  	"net"
    23  
    24  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    25  	"github.com/cloudwego/kitex/pkg/utils"
    26  )
    27  
    28  // GetCaller is used to get the Service Name of the caller.
    29  // Return false if failed to get the information.
    30  func GetCaller(ctx context.Context) (string, bool) {
    31  	defer func() { recover() }()
    32  
    33  	ri := rpcinfo.GetRPCInfo(ctx)
    34  	if ri == nil {
    35  		return "", false
    36  	}
    37  	return ri.From().ServiceName(), true
    38  }
    39  
    40  // GetMethod is used to get the current RPC Method name.
    41  // Return false if failed to get the information.
    42  func GetMethod(ctx context.Context) (string, bool) {
    43  	defer func() { recover() }()
    44  
    45  	ri := rpcinfo.GetRPCInfo(ctx)
    46  	if ri == nil {
    47  		return "", false
    48  	}
    49  	return ri.To().Method(), true
    50  }
    51  
    52  // GetCallerHandlerMethod is used to get the method name of caller.
    53  // Only the caller is a Kitex server will have this method information by default, or you can set K_METHOD into context.Context then kitex will get it.
    54  // Return false if failed to get the information.
    55  func GetCallerHandlerMethod(ctx context.Context) (string, bool) {
    56  	defer func() { recover() }()
    57  
    58  	ri := rpcinfo.GetRPCInfo(ctx)
    59  	if ri == nil {
    60  		return "", false
    61  	}
    62  	return ri.From().Method(), true
    63  }
    64  
    65  // GetIDLServiceName gets the serviceName which defined in IDL.
    66  // Return false if failed to get the information.
    67  func GetIDLServiceName(ctx context.Context) (string, bool) {
    68  	defer func() { recover() }()
    69  
    70  	ri := rpcinfo.GetRPCInfo(ctx)
    71  	if ri == nil {
    72  		return "", false
    73  	}
    74  	return ri.Invocation().ServiceName(), true
    75  }
    76  
    77  // GetCallerAddr is used for the server to get the Address of the caller.
    78  // Return false if failed to get the information.
    79  func GetCallerAddr(ctx context.Context) (net.Addr, bool) {
    80  	defer func() { recover() }()
    81  
    82  	ri := rpcinfo.GetRPCInfo(ctx)
    83  	if ri == nil {
    84  		return nil, false
    85  	}
    86  	return ri.From().Address(), true
    87  }
    88  
    89  // GetCallerIP is used for the server to get the IP of the caller.
    90  // Return false if failed to get the information.
    91  func GetCallerIP(ctx context.Context) (string, bool) {
    92  	defer func() { recover() }()
    93  
    94  	ri := rpcinfo.GetRPCInfo(ctx)
    95  	if ri == nil {
    96  		return "", false
    97  	}
    98  	addrStr := ri.From().Address().String()
    99  	if len(addrStr) == 0 {
   100  		return "", false
   101  	}
   102  
   103  	if ip, _, err := net.SplitHostPort(addrStr); err == nil {
   104  		return ip, true
   105  	}
   106  	return addrStr, true
   107  }
   108  
   109  // GetTransportProtocol gets the transport protocol of the request.
   110  // Return false if failed to get the information.
   111  func GetTransportProtocol(ctx context.Context) (string, bool) {
   112  	defer func() { recover() }()
   113  
   114  	ri := rpcinfo.GetRPCInfo(ctx)
   115  	if ri == nil {
   116  		return "", false
   117  	}
   118  	return ri.Config().TransportProtocol().String(), true
   119  }
   120  
   121  // GetRPCInfo gets the RPCInfo in ctx.
   122  // Return false if failed to get the information.
   123  func GetRPCInfo(ctx context.Context) (rpcinfo.RPCInfo, bool) {
   124  	defer func() { recover() }()
   125  	ri := rpcinfo.GetRPCInfo(ctx)
   126  	if ri == nil {
   127  		return nil, false
   128  	}
   129  	return ri, true
   130  }
   131  
   132  // GetRealReqFromKitexArgs assert the req to be KitexArgs and return the real request if succeeded, otherwise return nil.
   133  // This method should be used in the middleware.
   134  func GetRealReqFromKitexArgs(req interface{}) interface{} {
   135  	if arg, ok := req.(utils.KitexArgs); ok {
   136  		return arg.GetFirstArgument()
   137  	}
   138  	return nil
   139  }
   140  
   141  // GetRealRespFromKitexResult assert the req to be KitexResult and return the real response if succeeded, otherwise return nil.
   142  // This method should be used in the middleware.
   143  func GetRealRespFromKitexResult(resp interface{}) interface{} {
   144  	if res, ok := resp.(utils.KitexResult); ok {
   145  		return res.GetResult()
   146  	}
   147  	return nil
   148  }