metadata.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. /*
  2. *
  3. * Copyright 2014, Google Inc.
  4. * All rights reserved.
  5. *
  6. * Redistribution and use in source and binary forms, with or without
  7. * modification, are permitted provided that the following conditions are
  8. * met:
  9. *
  10. * * Redistributions of source code must retain the above copyright
  11. * notice, this list of conditions and the following disclaimer.
  12. * * Redistributions in binary form must reproduce the above
  13. * copyright notice, this list of conditions and the following disclaimer
  14. * in the documentation and/or other materials provided with the
  15. * distribution.
  16. * * Neither the name of Google Inc. nor the names of its
  17. * contributors may be used to endorse or promote products derived from
  18. * this software without specific prior written permission.
  19. *
  20. * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21. * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22. * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23. * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24. * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25. * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26. * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27. * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28. * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29. * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30. * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31. *
  32. */
  33. // Package metadata define the structure of the metadata supported by gRPC library.
  34. package metadata
  35. import (
  36. "encoding/base64"
  37. "fmt"
  38. "strings"
  39. "golang.org/x/net/context"
  40. )
  41. const (
  42. binHdrSuffix = "-bin"
  43. )
  44. // encodeKeyValue encodes key and value qualified for transmission via gRPC.
  45. // Transmitting binary headers violates HTTP/2 spec.
  46. // TODO(zhaoq): Maybe check if k is ASCII also.
  47. func encodeKeyValue(k, v string) (string, string) {
  48. k = strings.ToLower(k)
  49. if strings.HasSuffix(k, binHdrSuffix) {
  50. val := base64.StdEncoding.EncodeToString([]byte(v))
  51. v = string(val)
  52. }
  53. return k, v
  54. }
  55. // DecodeKeyValue returns the original key and value corresponding to the
  56. // encoded data in k, v.
  57. // If k is a binary header and v contains comma, v is split on comma before decoded,
  58. // and the decoded v will be joined with comma before returned.
  59. func DecodeKeyValue(k, v string) (string, string, error) {
  60. if !strings.HasSuffix(k, binHdrSuffix) {
  61. return k, v, nil
  62. }
  63. vvs := strings.Split(v, ",")
  64. for i, vv := range vvs {
  65. val, err := base64.StdEncoding.DecodeString(vv)
  66. if err != nil {
  67. return "", "", err
  68. }
  69. vvs[i] = string(val)
  70. }
  71. return k, strings.Join(vvs, ","), nil
  72. }
  73. // MD is a mapping from metadata keys to values. Users should use the following
  74. // two convenience functions New and Pairs to generate MD.
  75. type MD map[string][]string
  76. // New creates a MD from given key-value map.
  77. func New(m map[string]string) MD {
  78. md := MD{}
  79. for k, v := range m {
  80. key, val := encodeKeyValue(k, v)
  81. md[key] = append(md[key], val)
  82. }
  83. return md
  84. }
  85. // Pairs returns an MD formed by the mapping of key, value ...
  86. // Pairs panics if len(kv) is odd.
  87. func Pairs(kv ...string) MD {
  88. if len(kv)%2 == 1 {
  89. panic(fmt.Sprintf("metadata: Pairs got the odd number of input pairs for metadata: %d", len(kv)))
  90. }
  91. md := MD{}
  92. var k string
  93. for i, s := range kv {
  94. if i%2 == 0 {
  95. k = s
  96. continue
  97. }
  98. key, val := encodeKeyValue(k, s)
  99. md[key] = append(md[key], val)
  100. }
  101. return md
  102. }
  103. // Len returns the number of items in md.
  104. func (md MD) Len() int {
  105. return len(md)
  106. }
  107. // Copy returns a copy of md.
  108. func (md MD) Copy() MD {
  109. out := MD{}
  110. for k, v := range md {
  111. for _, i := range v {
  112. out[k] = append(out[k], i)
  113. }
  114. }
  115. return out
  116. }
  117. type mdKey struct{}
  118. // NewContext creates a new context with md attached.
  119. func NewContext(ctx context.Context, md MD) context.Context {
  120. return context.WithValue(ctx, mdKey{}, md)
  121. }
  122. // FromContext returns the MD in ctx if it exists.
  123. func FromContext(ctx context.Context) (md MD, ok bool) {
  124. md, ok = ctx.Value(mdKey{}).(MD)
  125. return
  126. }