request_response_funcs.go

Documentation: github.com/go-kit/kit/transport/grpc

     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  		val = base64.StdEncoding.EncodeToString([]byte(val))
    73  	}
    74  	return key, val
    75  }
    76  
    77  type contextKey int
    78  
    79  const (
    80  	ContextKeyRequestMethod contextKey = iota
    81  )
    82
View as plain text