blob: a277bee3f7bc4941fd9df9de58024c912e6b0b0f [file] [log] [blame]
Takahiro Suzukid7bf8202020-12-17 20:21:59 +09001// Copyright 2016 Michal Witkowski. All Rights Reserved.
2// See LICENSE for licensing terms.
3
4package metautils
5
6import (
7 "strings"
8
9 "golang.org/x/net/context"
10 "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 {
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.
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 {
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
100func (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.
111func (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.
122func (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}