gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/go-grpc-middleware/util/metautils/nicemd.go (about)

     1  // Copyright 2016 Michal Witkowski. All Rights Reserved.
     2  // See LICENSE for licensing terms.
     3  
     4  package metautils
     5  
     6  import (
     7  	"strings"
     8  
     9  	"gitee.com/ks-custle/core-gm/grpc/metadata"
    10  	"gitee.com/ks-custle/core-gm/net/context"
    11  )
    12  
    13  // NiceMD is a convenience wrapper definiting extra functions on the metadata.
    14  type NiceMD metadata.MD
    15  
    16  // ExtractIncoming extracts an inbound metadata from the server-side context.
    17  //
    18  // This function always returns a NiceMD wrapper of the metadata.MD, in case the context doesn't have metadata it returns
    19  // a new empty NiceMD.
    20  func ExtractIncoming(ctx context.Context) NiceMD {
    21  	md, ok := metadata.FromIncomingContext(ctx)
    22  	if !ok {
    23  		return NiceMD(metadata.Pairs())
    24  	}
    25  	return NiceMD(md)
    26  }
    27  
    28  // ExtractOutgoing extracts an outbound metadata from the client-side context.
    29  //
    30  // This function always returns a NiceMD wrapper of the metadata.MD, in case the context doesn't have metadata it returns
    31  // a new empty NiceMD.
    32  func ExtractOutgoing(ctx context.Context) NiceMD {
    33  	md, ok := metadata.FromOutgoingContext(ctx)
    34  	if !ok {
    35  		return NiceMD(metadata.Pairs())
    36  	}
    37  	return NiceMD(md)
    38  }
    39  
    40  // Clone performs a *deep* copy of the metadata.MD.
    41  //
    42  // You can specify the lower-case copiedKeys to only copy certain whitelisted keys. If no keys are explicitly whitelisted
    43  // all keys get copied.
    44  func (m NiceMD) Clone(copiedKeys ...string) NiceMD {
    45  	newMd := NiceMD(metadata.Pairs())
    46  	for k, vv := range m {
    47  		found := false
    48  		if len(copiedKeys) == 0 {
    49  			found = true
    50  		} else {
    51  			for _, allowedKey := range copiedKeys {
    52  				if strings.ToLower(allowedKey) == strings.ToLower(k) {
    53  					found = true
    54  					break
    55  				}
    56  			}
    57  		}
    58  		if !found {
    59  			continue
    60  		}
    61  		newMd[k] = make([]string, len(vv))
    62  		copy(newMd[k], vv)
    63  	}
    64  	return NiceMD(newMd)
    65  }
    66  
    67  // ToOutgoing sets the given NiceMD as a client-side context for dispatching.
    68  func (m NiceMD) ToOutgoing(ctx context.Context) context.Context {
    69  	return metadata.NewOutgoingContext(ctx, metadata.MD(m))
    70  }
    71  
    72  // ToIncoming sets the given NiceMD as a server-side context for dispatching.
    73  //
    74  // This is mostly useful in ServerInterceptors..
    75  func (m NiceMD) ToIncoming(ctx context.Context) context.Context {
    76  	return metadata.NewIncomingContext(ctx, metadata.MD(m))
    77  }
    78  
    79  // Get retrieves a single value from the metadata.
    80  //
    81  // It works analogously to http.Header.Get, returning the first value if there are many set. If the value is not set,
    82  // an empty string is returned.
    83  //
    84  // The function is binary-key safe.
    85  func (m NiceMD) Get(key string) string {
    86  	k, _ := encodeKeyValue(key, "")
    87  	vv, ok := m[k]
    88  	if !ok {
    89  		return ""
    90  	}
    91  	return vv[0]
    92  }
    93  
    94  // Del retrieves a single value from the metadata.
    95  //
    96  // It works analogously to http.Header.Del, deleting all values if they exist.
    97  //
    98  // The function is binary-key safe.
    99  
   100  func (m NiceMD) Del(key string) NiceMD {
   101  	k, _ := encodeKeyValue(key, "")
   102  	delete(m, k)
   103  	return m
   104  }
   105  
   106  // Set sets the given value in a metadata.
   107  //
   108  // It works analogously to http.Header.Set, overwriting all previous metadata values.
   109  //
   110  // The function is binary-key safe.
   111  func (m NiceMD) Set(key string, value string) NiceMD {
   112  	k, v := encodeKeyValue(key, value)
   113  	m[k] = []string{v}
   114  	return m
   115  }
   116  
   117  // Add retrieves a single value from the metadata.
   118  //
   119  // It works analogously to http.Header.Add, as it appends to any existing values associated with key.
   120  //
   121  // The function is binary-key safe.
   122  func (m NiceMD) Add(key string, value string) NiceMD {
   123  	k, v := encodeKeyValue(key, value)
   124  	m[k] = append(m[k], v)
   125  	return m
   126  }