From e709f83253a45b8c4e4902f8a817f624a9a6d0aa Mon Sep 17 00:00:00 2001
From: Hitoshi Mitake <mitake.hitoshi@gmail.com>
Date: Mon, 24 Jul 2017 18:09:37 +0900
Subject: [PATCH] etcdmain, proxy: support authed RPCs with grpcproxy

This commit lets grpcproxy support authed RPCs. Auth tokens supplied
by clients are now forwarded to etcdserver by grpcproxy.
---
 etcdmain/grpc_proxy.go  |  2 ++
 proxy/grpcproxy/util.go | 56 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 proxy/grpcproxy/util.go

diff --git a/etcdmain/grpc_proxy.go b/etcdmain/grpc_proxy.go
index 27dcac375..e0fc331d2 100644
--- a/etcdmain/grpc_proxy.go
+++ b/etcdmain/grpc_proxy.go
@@ -204,6 +204,8 @@ func mustNewClient() *clientv3.Client {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
 	}
+	cfg.DialOptions = append(cfg.DialOptions,
+		grpc.WithUnaryInterceptor(grpcproxy.AuthUnaryClientInterceptor))
 	client, err := clientv3.New(*cfg)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
diff --git a/proxy/grpcproxy/util.go b/proxy/grpcproxy/util.go
new file mode 100644
index 000000000..28e3c111f
--- /dev/null
+++ b/proxy/grpcproxy/util.go
@@ -0,0 +1,56 @@
+// Copyright 2017 The etcd Authors
+//
+// 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 grpcproxy
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func getAuthTokenFromClient(ctx context.Context) string {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if ok {
+		ts, ok := md["token"]
+		if ok {
+			return ts[0]
+		}
+	}
+	return ""
+}
+
+type proxyTokenCredential struct {
+	token string
+}
+
+func (cred *proxyTokenCredential) RequireTransportSecurity() bool {
+	return false
+}
+
+func (cred *proxyTokenCredential) GetRequestMetadata(ctx context.Context, s ...string) (map[string]string, error) {
+	return map[string]string{
+		"token": cred.token,
+	}, nil
+}
+
+func AuthUnaryClientInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	token := getAuthTokenFromClient(ctx)
+	if token != "" {
+		tokenCred := &proxyTokenCredential{token}
+		opts = append(opts, grpc.PerRPCCredentials(tokenCred))
+	}
+	return invoker(ctx, method, req, reply, cc, opts...)
+}