vendor/github.com/go-kit/kit/transport/grpc/request_response_funcs.go

   1 package grpc
   2
   3 import (
   4         "context"
   5         "encoding/base64"
   6         "strings"
   7
   8         "google.golang.org/grpc/metadata"
   9 )
  10
  11 const (
  12         binHdrSuffix = "-bin"
  13 )
  14
  15 // ClientRequestFunc may take information from context and use it to construct
  16 // metadata headers to be transported to the server. ClientRequestFuncs are
  17 // executed after creating the request but prior to sending the gRPC request to
  18 // the server.
  19 type ClientRequestFunc func(context.Context, *metadata.MD) context.Context
  20
  21 // ServerRequestFunc may take information from the received metadata header and
  22 // use it to place items in the request scoped context. ServerRequestFuncs are
  23 // executed prior to invoking the endpoint.
  24 type ServerRequestFunc func(context.Context, metadata.MD) context.Context
  25
  26 // ServerResponseFunc may take information from a request context and use it to
  27 // manipulate the gRPC response metadata headers and trailers. ResponseFuncs are
  28 // only executed in servers, after invoking the endpoint but prior to writing a
  29 // response.
  30 type ServerResponseFunc func(ctx context.Context, header *metadata.MD, trailer *metadata.MD) context.Context
  31
  32 // ClientResponseFunc may take information from a gRPC metadata header and/or
  33 // trailer and make the responses available for consumption. ClientResponseFuncs
  34 // are only executed in clients, after a request has been made, but prior to it
  35 // being decoded.
  36 type ClientResponseFunc func(ctx context.Context, header metadata.MD, trailer metadata.MD) context.Context
  37
  38 // SetRequestHeader returns a ClientRequestFunc that sets the specified metadata
  39 // key-value pair.
  40 func SetRequestHeader(key, val string) ClientRequestFunc {
  41         return func(ctx context.Context, md *metadata.MD) context.Context {
  42                 key, val := EncodeKeyValue(key, val)
  43                 (*md)[key] = append((*md)[key], val)
  44                 return ctx
  45         }
  46 }
  47
  48 // SetResponseHeader returns a ResponseFunc that sets the specified metadata
  49 // key-value pair.
  50 func SetResponseHeader(key, val string) ServerResponseFunc {
  51         return func(ctx context.Context, md *metadata.MD, _ *metadata.MD) context.Context {
  52                 key, val := EncodeKeyValue(key, val)
  53                 (*md)[key] = append((*md)[key], val)
  54                 return ctx
  55         }
  56 }
  57
  58 // SetResponseTrailer returns a ResponseFunc that sets the specified metadata
  59 // key-value pair.
  60 func SetResponseTrailer(key, val string) ServerResponseFunc {
  61         return func(ctx context.Context, _ *metadata.MD, md *metadata.MD) context.Context {
  62                 key, val := EncodeKeyValue(key, val)
  63                 (*md)[key] = append((*md)[key], val)
  64                 return ctx
  65         }
  66 }
  67
  68 // EncodeKeyValue sanitizes a key-value pair for use in gRPC metadata headers.
  69 func EncodeKeyValue(key, val string) (string, string) {
  70         key = strings.ToLower(key)
  71         if strings.HasSuffix(key, binHdrSuffix) {
  72                 v := base64.StdEncoding.EncodeToString([]byte(val))
  73                 val = string(v)
  74         }
  75         return key, val
  76 }