Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a5bd8ce

Browse files
crosbymichaelcrosbymichael
authored
Merge pull request #40 from mxpv/headers
Support headers
2 parents f821483 + 04523b9 commit a5bd8ce

File tree

6 files changed

+226
-5
lines changed

6 files changed

+226
-5
lines changed

‎client.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,10 @@ func (c *Client) Call(ctx context.Context, service, method string, req, resp int
9999
cresp = &Response{}
100100
)
101101

102+
if metadata, ok := GetMetadata(ctx); ok {
103+
creq.Metadata = metadata
104+
}
105+
102106
if dl, ok := ctx.Deadline(); ok {
103107
creq.TimeoutNano = dl.Sub(time.Now()).Nanoseconds()
104108
}

‎metadata.go‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*
2+
Copyright The containerd Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package ttrpc
18+
19+
import "context"
20+
21+
// Metadata represents the key-value pairs (similar to http.Header) to be passed to ttrpc server from a client.
22+
type Metadata map[string]StringList
23+
24+
// Get returns the metadata for a given key when they exist.
25+
// If there is no metadata, a nil slice and false are returned.
26+
func (m Metadata) Get(key string) ([]string, bool) {
27+
list, ok := m[key]
28+
if !ok || len(list.List) == 0 {
29+
return nil, false
30+
}
31+
32+
return list.List, true
33+
}
34+
35+
// Set sets the provided values for a given key.
36+
// The values will overwrite any existing values.
37+
// If no values provided, a key will be deleted.
38+
func (m Metadata) Set(key string, values ...string) {
39+
if len(values) == 0 {
40+
delete(m, key)
41+
return
42+
}
43+
44+
m[key] = StringList{List: values}
45+
}
46+
47+
// Append appends additional values to the given key.
48+
func (m Metadata) Append(key string, values ...string) {
49+
if len(values) == 0 {
50+
return
51+
}
52+
53+
list, ok := m[key]
54+
if ok {
55+
m.Set(key, append(list.List, values...)...)
56+
} else {
57+
m.Set(key, values...)
58+
}
59+
}
60+
61+
type metadataKey struct{}
62+
63+
// GetMetadata retrieves metadata from context.Context (previously attached with WithMetadata)
64+
func GetMetadata(ctx context.Context) (Metadata, bool) {
65+
metadata, ok := ctx.Value(metadataKey{}).(Metadata)
66+
return metadata, ok
67+
}
68+
69+
// GetMetadataValue gets a specific metadata value by name from context.Context
70+
func GetMetadataValue(ctx context.Context, name string) (string, bool) {
71+
metadata, ok := GetMetadata(ctx)
72+
if !ok {
73+
return "", false
74+
}
75+
76+
if list, ok := metadata.Get(name); ok {
77+
return list[0], true
78+
}
79+
80+
return "", false
81+
}
82+
83+
// WithMetadata attaches metadata map to a context.Context
84+
func WithMetadata(ctx context.Context, headers Metadata) context.Context {
85+
return context.WithValue(ctx, metadataKey{}, headers)
86+
}

‎metadata_test.go‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
/*
2+
Copyright The containerd Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package ttrpc
18+
19+
import (
20+
"context"
21+
"testing"
22+
)
23+
24+
func TestMetadata_Get(t *testing.T) {
25+
metadata := make(Metadata)
26+
metadata.Set("foo", "1", "2")
27+
28+
if list, ok := metadata.Get("foo"); !ok {
29+
t.Error("key not found")
30+
} else if len(list) != 2 {
31+
t.Errorf("unexpected number of values: %d", len(list))
32+
} else if list[0] != "1" {
33+
t.Errorf("invalid metadata value at 0: %s", list[0])
34+
} else if list[1] != "2" {
35+
t.Errorf("invalid metadata value at 1: %s", list[1])
36+
}
37+
}
38+
39+
func TestMetadata_GetInvalidKey(t *testing.T) {
40+
metadata := make(Metadata)
41+
metadata.Set("foo", "1", "2")
42+
43+
if _, ok := metadata.Get("invalid"); ok {
44+
t.Error("found invalid key")
45+
}
46+
}
47+
48+
func TestMetadata_Unset(t *testing.T) {
49+
metadata := make(Metadata)
50+
metadata.Set("foo", "1", "2")
51+
metadata.Set("foo")
52+
53+
if _, ok := metadata.Get("foo"); ok {
54+
t.Error("key not deleted")
55+
}
56+
}
57+
58+
func TestMetadata_Replace(t *testing.T) {
59+
metadata := make(Metadata)
60+
metadata.Set("foo", "1", "2")
61+
metadata.Set("foo", "3", "4")
62+
63+
if list, ok := metadata.Get("foo"); !ok {
64+
t.Error("key not found")
65+
} else if len(list) != 2 {
66+
t.Errorf("unexpected number of values: %d", len(list))
67+
} else if list[0] != "3" {
68+
t.Errorf("invalid metadata value at 0: %s", list[0])
69+
} else if list[1] != "4" {
70+
t.Errorf("invalid metadata value at 1: %s", list[1])
71+
}
72+
}
73+
74+
func TestMetadata_Append(t *testing.T) {
75+
metadata := make(Metadata)
76+
metadata.Set("foo", "1")
77+
metadata.Append("foo", "2")
78+
metadata.Append("bar", "3")
79+
80+
if list, ok := metadata.Get("foo"); !ok {
81+
t.Error("key not found")
82+
} else if len(list) != 2 {
83+
t.Errorf("unexpected number of values: %d", len(list))
84+
} else if list[0] != "1" {
85+
t.Errorf("invalid metadata value at 0: %s", list[0])
86+
} else if list[1] != "2" {
87+
t.Errorf("invalid metadata value at 1: %s", list[1])
88+
}
89+
90+
if list, ok := metadata.Get("bar"); !ok {
91+
t.Error("key not found")
92+
} else if list[0] != "3" {
93+
t.Errorf("invalid value: %s", list[0])
94+
}
95+
}
96+
97+
func TestMetadata_Context(t *testing.T) {
98+
metadata := make(Metadata)
99+
metadata.Set("foo", "bar")
100+
101+
ctx := WithMetadata(context.Background(), metadata)
102+
103+
if bar, ok := GetMetadataValue(ctx, "foo"); !ok {
104+
t.Error("metadata not found")
105+
} else if bar != "bar" {
106+
t.Errorf("invalid metadata value: %q", bar)
107+
}
108+
}

‎server.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -466,6 +466,10 @@ func (c *serverConn) run(sctx context.Context) {
466466
var noopFunc = func() {}
467467

468468
func getRequestContext(ctx context.Context, req *Request) (retCtx context.Context, cancel func()) {
469+
if req.Metadata != nil {
470+
ctx = WithMetadata(ctx, req.Metadata)
471+
}
472+
469473
cancel = noopFunc
470474
if req.TimeoutNano == 0 {
471475
return ctx, cancel

‎server_test.go‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ func (tc *testingClient) Test(ctx context.Context, req *testPayload) (*testPaylo
6161
type testPayload struct {
6262
Foo string `protobuf:"bytes,1,opt,name=foo,proto3"`
6363
Deadline int64 `protobuf:"varint,2,opt,name=deadline,proto3"`
64+
Metadata string `protobuf:"bytes,3,opt,name=metadata,proto3"`
6465
}
6566

6667
func (r *testPayload) Reset() { *r = testPayload{} }
@@ -75,6 +76,11 @@ func (s *testingServer) Test(ctx context.Context, req *testPayload) (*testPayloa
7576
if dl, ok := ctx.Deadline(); ok {
7677
tp.Deadline = dl.UnixNano()
7778
}
79+
80+
if v, ok := GetMetadataValue(ctx, "foo"); ok {
81+
tp.Metadata = v
82+
}
83+
7884
return tp, nil
7985
}
8086

@@ -540,14 +546,16 @@ func roundTrip(ctx context.Context, t *testing.T, client *testingClient, value s
540546
}
541547
)
542548

549+
ctx = WithMetadata(ctx, Metadata{"foo": makeStringList("bar")})
550+
543551
resp, err := client.Test(ctx, tp)
544552
if err != nil {
545553
t.Fatal(err)
546554
}
547555

548556
results <- callResult{
549557
input: tp,
550-
expected: &testPayload{Foo: strings.Repeat(tp.Foo, 2)},
558+
expected: &testPayload{Foo: strings.Repeat(tp.Foo, 2), Metadata: "bar"},
551559
received: resp,
552560
}
553561
}

‎types.go‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,11 @@ import (
2323
)
2424

2525
type Request struct {
26-
Service string `protobuf:"bytes,1,opt,name=service,proto3"`
27-
Method string `protobuf:"bytes,2,opt,name=method,proto3"`
28-
Payload []byte `protobuf:"bytes,3,opt,name=payload,proto3"`
29-
TimeoutNano int64 `protobuf:"varint,4,opt,name=timeout_nano,proto3"`
26+
Service string `protobuf:"bytes,1,opt,name=service,proto3"`
27+
Method string `protobuf:"bytes,2,opt,name=method,proto3"`
28+
Payload []byte `protobuf:"bytes,3,opt,name=payload,proto3"`
29+
TimeoutNano int64 `protobuf:"varint,4,opt,name=timeout_nano,proto3"`
30+
Metadata Metadata `protobuf:"bytes,5,opt,name=metadata,proto3" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
3031
}
3132

3233
func (r *Request) Reset() { *r = Request{} }
@@ -41,3 +42,13 @@ type Response struct {
4142
func (r *Response) Reset() { *r = Response{} }
4243
func (r *Response) String() string { return fmt.Sprintf("%+#v", r) }
4344
func (r *Response) ProtoMessage() {}
45+
46+
type StringList struct {
47+
List []string `protobuf:"bytes,1,rep,name=list,proto3"`
48+
}
49+
50+
func (r *StringList) Reset() { *r = StringList{} }
51+
func (r *StringList) String() string { return fmt.Sprintf("%+#v", r) }
52+
func (r *StringList) ProtoMessage() {}
53+
54+
func makeStringList(item ...string) StringList { return StringList{List: item} }

0 commit comments

Comments
 (0)