blob: 1c60585dd3565b755633ca737cbf67046f80a4a0 [file] [log] [blame]
Girish Kumar74240652020-07-10 11:54:28 +00001// Copyright 2016 Michal Witkowski. All Rights Reserved.
2// See LICENSE for licensing terms.
3
4package metautils
5
6import (
khenaidoo26721882021-08-11 17:42:52 -04007 "context"
Girish Kumar74240652020-07-10 11:54:28 +00008 "strings"
9
Girish Kumar74240652020-07-10 11:54:28 +000010 "google.golang.org/grpc/metadata"
11)
12
13// NiceMD is a convenience wrapper definiting extra functions on the metadata.
14type 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.
20func 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.
32func 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.
44func (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 {
khenaidoo26721882021-08-11 17:42:52 -040052 if strings.EqualFold(allowedKey, k) {
Girish Kumar74240652020-07-10 11:54:28 +000053 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.
68func (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..
75func (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.
85func (m NiceMD) Get(key string) string {
khenaidoo26721882021-08-11 17:42:52 -040086 k := strings.ToLower(key)
Girish Kumar74240652020-07-10 11:54:28 +000087 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
100func (m NiceMD) Del(key string) NiceMD {
khenaidoo26721882021-08-11 17:42:52 -0400101 k := strings.ToLower(key)
Girish Kumar74240652020-07-10 11:54:28 +0000102 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.
111func (m NiceMD) Set(key string, value string) NiceMD {
khenaidoo26721882021-08-11 17:42:52 -0400112 k := strings.ToLower(key)
113 m[k] = []string{value}
Girish Kumar74240652020-07-10 11:54:28 +0000114 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.
122func (m NiceMD) Add(key string, value string) NiceMD {
khenaidoo26721882021-08-11 17:42:52 -0400123 k := strings.ToLower(key)
124 m[k] = append(m[k], value)
Girish Kumar74240652020-07-10 11:54:28 +0000125 return m
126}