-
Notifications
You must be signed in to change notification settings - Fork 116
/
interceptor.go
75 lines (63 loc) · 2.09 KB
/
interceptor.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
/*
Copyright 2021 Portworx
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package correlation
import (
"context"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
// ContextInterceptor represents a correlation interceptor
type ContextInterceptor struct {
Origin Component
}
// ContextUnaryServerInterceptor creates a gRPC interceptor for adding
// correlation ID to each request
func (ci *ContextInterceptor) ContextUnaryServerInterceptor(
ctx context.Context,
req interface{},
info *grpc.UnaryServerInfo,
handler grpc.UnaryHandler,
) (interface{}, error) {
md, ok := metadata.FromIncomingContext(ctx)
if ok {
// Get request context from gRPC metadata
rc := RequestContextFromContextMetadata(md)
// Only add to context if an ID exists
if len(rc.ID) > 0 {
ctx = context.WithValue(ctx, ContextKey, rc)
}
}
ctx = WithCorrelationContext(ctx, ci.Origin)
return handler(ctx, req)
}
// ContextUnaryClientInterceptor creates a gRPC interceptor for adding
// correlation ID to each request
func ContextUnaryClientInterceptor(
ctx context.Context,
method string,
req, reply interface{},
cc *grpc.ClientConn,
invoker grpc.UnaryInvoker,
opts ...grpc.CallOption,
) error {
// Create new metadata from request context in context
newContextMap := RequestContextFromContextValue(ctx).AsMap()
newContextMD := metadata.New(newContextMap)
// Get existing metadata if exists and join with new one
existingMD, ok := metadata.FromOutgoingContext(ctx)
if ok {
newContextMD = metadata.Join(newContextMD, existingMD)
}
ctx = metadata.NewOutgoingContext(ctx, newContextMD)
return invoker(ctx, method, req, reply, cc, opts...)
}