nicemd.go

Documentation: github.com/grpc-ecosystem/go-grpc-middleware/util/metautils

     1  // Copyright 2016 Michal Witkowski. All Rights Reserved.
     2  // See LICENSE for licensing terms.
     3  
     4  package metautils
     5  
     6  import (
     7  	"context"
     8  	"strings"
     9  
    10  	"google.golang.org/grpc/metadata"
    11  )
    12  
    13  // NiceMD is a convenience wrapper defining 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 allow-listed keys. If no keys are explicitly allow-listed
    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.EqualFold(allowedKey, 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 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 := strings.ToLower(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 := strings.ToLower(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 := strings.ToLower(key)
   113  	m[k] = []string{value}
   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 := strings.ToLower(key)
   124  	m[k] = append(m[k], value)
   125  	return m
   126  }
   127
View as plain text