blob: ac9bd4882befb55b8c76132c5fda10089cac1893 [file] [log] [blame]
khenaidoo26721882021-08-11 17:42:52 -04001/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20package thrift
21
22import (
23 "context"
24)
25
26// See https://godoc.org/context#WithValue on why do we need the unexported typedefs.
27type (
28 headerKey string
29 headerKeyList int
30)
31
32// Values for headerKeyList.
33const (
34 headerKeyListRead headerKeyList = iota
35 headerKeyListWrite
36)
37
38// SetHeader sets a header in the context.
39func SetHeader(ctx context.Context, key, value string) context.Context {
40 return context.WithValue(
41 ctx,
42 headerKey(key),
43 value,
44 )
45}
46
47// UnsetHeader unsets a previously set header in the context.
48func UnsetHeader(ctx context.Context, key string) context.Context {
49 return context.WithValue(
50 ctx,
51 headerKey(key),
52 nil,
53 )
54}
55
56// GetHeader returns a value of the given header from the context.
57func GetHeader(ctx context.Context, key string) (value string, ok bool) {
58 if v := ctx.Value(headerKey(key)); v != nil {
59 value, ok = v.(string)
60 }
61 return
62}
63
64// SetReadHeaderList sets the key list of read THeaders in the context.
65func SetReadHeaderList(ctx context.Context, keys []string) context.Context {
66 return context.WithValue(
67 ctx,
68 headerKeyListRead,
69 keys,
70 )
71}
72
73// GetReadHeaderList returns the key list of read THeaders from the context.
74func GetReadHeaderList(ctx context.Context) []string {
75 if v := ctx.Value(headerKeyListRead); v != nil {
76 if value, ok := v.([]string); ok {
77 return value
78 }
79 }
80 return nil
81}
82
83// SetWriteHeaderList sets the key list of THeaders to write in the context.
84func SetWriteHeaderList(ctx context.Context, keys []string) context.Context {
85 return context.WithValue(
86 ctx,
87 headerKeyListWrite,
88 keys,
89 )
90}
91
92// GetWriteHeaderList returns the key list of THeaders to write from the context.
93func GetWriteHeaderList(ctx context.Context) []string {
94 if v := ctx.Value(headerKeyListWrite); v != nil {
95 if value, ok := v.([]string); ok {
96 return value
97 }
98 }
99 return nil
100}
101
102// AddReadTHeaderToContext adds the whole THeader headers into context.
103func AddReadTHeaderToContext(ctx context.Context, headers THeaderMap) context.Context {
104 keys := make([]string, 0, len(headers))
105 for key, value := range headers {
106 ctx = SetHeader(ctx, key, value)
107 keys = append(keys, key)
108 }
109 return SetReadHeaderList(ctx, keys)
110}