callctx.go raw

   1  // Copyright 2023, Google Inc.
   2  // All rights reserved.
   3  //
   4  // Redistribution and use in source and binary forms, with or without
   5  // modification, are permitted provided that the following conditions are
   6  // met:
   7  //
   8  //     * Redistributions of source code must retain the above copyright
   9  // notice, this list of conditions and the following disclaimer.
  10  //     * Redistributions in binary form must reproduce the above
  11  // copyright notice, this list of conditions and the following disclaimer
  12  // in the documentation and/or other materials provided with the
  13  // distribution.
  14  //     * Neither the name of Google Inc. nor the names of its
  15  // contributors may be used to endorse or promote products derived from
  16  // this software without specific prior written permission.
  17  //
  18  // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  19  // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  20  // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  21  // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  22  // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  23  // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  24  // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  25  // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  26  // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  27  // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  28  // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  29  
  30  // Package callctx provides helpers for storing and retrieving values out of
  31  // [context.Context]. These values are used by our client libraries in various
  32  // ways across the stack.
  33  package callctx
  34  
  35  import (
  36  	"context"
  37  	"fmt"
  38  )
  39  
  40  const (
  41  	// XGoogFieldMaskHeader is the canonical header key for the [System Parameter]
  42  	// that specifies the response read mask. The value(s) for this header
  43  	// must adhere to format described in [fieldmaskpb].
  44  	//
  45  	// [System Parameter]: https://cloud.google.com/apis/docs/system-parameters
  46  	// [fieldmaskpb]: https://google.golang.org/protobuf/types/known/fieldmaskpb
  47  	XGoogFieldMaskHeader = "x-goog-fieldmask"
  48  
  49  	headerKey = contextKey("header")
  50  )
  51  
  52  // contextKey is a private type used to store/retrieve context values.
  53  type contextKey string
  54  
  55  // HeadersFromContext retrieves headers set from [SetHeaders]. These headers
  56  // can then be cast to http.Header or metadata.MD to send along on requests.
  57  func HeadersFromContext(ctx context.Context) map[string][]string {
  58  	m, ok := ctx.Value(headerKey).(map[string][]string)
  59  	if !ok {
  60  		return nil
  61  	}
  62  	return m
  63  }
  64  
  65  // SetHeaders stores key value pairs in the returned context that can later
  66  // be retrieved by [HeadersFromContext]. Values stored in this manner will
  67  // automatically be retrieved by client libraries and sent as outgoing headers
  68  // on all requests. keyvals should have a corresponding value for every key
  69  // provided. If there is an odd number of keyvals this method will panic.
  70  func SetHeaders(ctx context.Context, keyvals ...string) context.Context {
  71  	if len(keyvals)%2 != 0 {
  72  		panic(fmt.Sprintf("callctx: an even number of key value pairs must be provided, got %d", len(keyvals)))
  73  	}
  74  	h, ok := ctx.Value(headerKey).(map[string][]string)
  75  	if !ok {
  76  		h = make(map[string][]string)
  77  	} else {
  78  		h = cloneHeaders(h)
  79  	}
  80  
  81  	for i := 0; i < len(keyvals); i = i + 2 {
  82  		h[keyvals[i]] = append(h[keyvals[i]], keyvals[i+1])
  83  	}
  84  	return context.WithValue(ctx, headerKey, h)
  85  }
  86  
  87  // cloneHeaders makes a new key-value map while reusing the value slices.
  88  // As such, new values should be appended to the value slice, and modifying
  89  // indexed values is not thread safe.
  90  //
  91  // TODO: Replace this with maps.Clone when Go 1.21 is the minimum version.
  92  func cloneHeaders(h map[string][]string) map[string][]string {
  93  	c := make(map[string][]string, len(h))
  94  	for k, v := range h {
  95  		vc := make([]string, len(v))
  96  		copy(vc, v)
  97  		c[k] = vc
  98  	}
  99  	return c
 100  }
 101