...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 package grpcproxy
16
17 import (
18 "context"
19
20 "go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
21
22 "google.golang.org/grpc"
23 "google.golang.org/grpc/metadata"
24 )
25
26 func getAuthTokenFromClient(ctx context.Context) string {
27 md, ok := metadata.FromIncomingContext(ctx)
28 if ok {
29 ts, ok := md[rpctypes.TokenFieldNameGRPC]
30 if ok {
31 return ts[0]
32 }
33 }
34 return ""
35 }
36
37 func withClientAuthToken(ctx, ctxWithToken context.Context) context.Context {
38 token := getAuthTokenFromClient(ctxWithToken)
39 if token != "" {
40 ctx = context.WithValue(ctx, rpctypes.TokenFieldNameGRPC, token)
41 }
42 return ctx
43 }
44
45 type proxyTokenCredential struct {
46 token string
47 }
48
49 func (cred *proxyTokenCredential) RequireTransportSecurity() bool {
50 return false
51 }
52
53 func (cred *proxyTokenCredential) GetRequestMetadata(ctx context.Context, s ...string) (map[string]string, error) {
54 return map[string]string{
55 rpctypes.TokenFieldNameGRPC: cred.token,
56 }, nil
57 }
58
59 func AuthUnaryClientInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
60 token := getAuthTokenFromClient(ctx)
61 if token != "" {
62 tokenCred := &proxyTokenCredential{token}
63 opts = append(opts, grpc.PerRPCCredentials(tokenCred))
64 }
65 return invoker(ctx, method, req, reply, cc, opts...)
66 }
67
68 func AuthStreamClientInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
69 tokenif := ctx.Value(rpctypes.TokenFieldNameGRPC)
70 if tokenif != nil {
71 tokenCred := &proxyTokenCredential{tokenif.(string)}
72 opts = append(opts, grpc.PerRPCCredentials(tokenCred))
73 }
74 return streamer(ctx, desc, cc, method, opts...)
75 }
76
View as plain text