github.com/go-graphite/carbonapi@v0.17.0/util/ctx/ctx.go (about)

     1  package ctx
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  )
     7  
     8  type key int
     9  
    10  const (
    11  	HeaderUUIDAPI    = "X-CTX-CarbonAPI-UUID"
    12  	HeaderUUIDZipper = "X-CTX-CarbonZipper-UUID"
    13  
    14  	uuidKey key = iota
    15  	headersToPassKey
    16  	headersToLogKey
    17  	maxDataPoints
    18  )
    19  
    20  func ifaceToString(v interface{}) string {
    21  	if v != nil {
    22  		return v.(string)
    23  	}
    24  	return ""
    25  }
    26  
    27  func getCtxString(ctx context.Context, k key) string {
    28  	return ifaceToString(ctx.Value(k))
    29  }
    30  
    31  func getCtxInt64(ctx context.Context, k key) int64 {
    32  	v := ctx.Value(k)
    33  	if v != nil {
    34  		return v.(int64)
    35  	}
    36  	return 0
    37  }
    38  
    39  func getCtxMapString(ctx context.Context, k key) map[string]string {
    40  	v := ctx.Value(k)
    41  	if v != nil {
    42  		vv, ok := v.(map[string]string)
    43  		if !ok {
    44  			return map[string]string{}
    45  		}
    46  		return vv
    47  	}
    48  	return map[string]string{}
    49  }
    50  
    51  func GetPassHeaders(ctx context.Context) map[string]string {
    52  	return getCtxMapString(ctx, headersToPassKey)
    53  }
    54  
    55  func SetPassHeaders(ctx context.Context, h map[string]string) context.Context {
    56  	return context.WithValue(ctx, headersToPassKey, h)
    57  }
    58  
    59  func GetLogHeaders(ctx context.Context) map[string]string {
    60  	return getCtxMapString(ctx, headersToLogKey)
    61  }
    62  
    63  func SetLogHeaders(ctx context.Context, h map[string]string) context.Context {
    64  	return context.WithValue(ctx, headersToLogKey, h)
    65  }
    66  
    67  func GetUUID(ctx context.Context) string {
    68  	return getCtxString(ctx, uuidKey)
    69  }
    70  
    71  func SetUUID(ctx context.Context, v string) context.Context {
    72  	return context.WithValue(ctx, uuidKey, v)
    73  }
    74  
    75  func SetMaxDatapoints(ctx context.Context, h int64) context.Context {
    76  	return context.WithValue(ctx, maxDataPoints, h)
    77  }
    78  
    79  func GetMaxDatapoints(ctx context.Context) int64 {
    80  	return getCtxInt64(ctx, maxDataPoints)
    81  }
    82  
    83  func ParseCtx(h http.HandlerFunc, uuidKey string) http.HandlerFunc {
    84  	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
    85  		uuid := req.Header.Get(uuidKey)
    86  
    87  		ctx := req.Context()
    88  		ctx = SetUUID(ctx, uuid)
    89  
    90  		h.ServeHTTP(rw, req.WithContext(ctx))
    91  	})
    92  }
    93  
    94  func MarshalPassHeaders(ctx context.Context, response *http.Request) *http.Request {
    95  	headers := GetPassHeaders(ctx)
    96  	for name, value := range headers {
    97  		response.Header.Add(name, value)
    98  	}
    99  
   100  	return response
   101  }
   102  
   103  func MarshalCtx(ctx context.Context, response *http.Request, uuidKey string) *http.Request {
   104  	response.Header.Add(uuidKey, GetUUID(ctx))
   105  
   106  	return response
   107  }