浏览代码

proxy: handle authed snapshot request in grpcproxy

Like the previous commit 10f783efdd12, this commit lets grpcproxy
forward an auth token supplied by its client in an explicit
manner. snapshot is a stream RPC so this process is required like
watch.
Hitoshi Mitake 8 年之前
父节点
当前提交
e8c18e3368
共有 3 个文件被更改,包括 11 次插入5 次删除
  1. 2 0
      proxy/grpcproxy/maintenance.go
  2. 8 0
      proxy/grpcproxy/util.go
  3. 1 5
      proxy/grpcproxy/watch_broadcast.go

+ 2 - 0
proxy/grpcproxy/maintenance.go

@@ -42,6 +42,8 @@ func (mp *maintenanceProxy) Snapshot(sr *pb.SnapshotRequest, stream pb.Maintenan
 	ctx, cancel := context.WithCancel(stream.Context())
 	defer cancel()
 
+	ctx = withClientAuthToken(ctx, stream.Context())
+
 	sc, err := pb.NewMaintenanceClient(conn).Snapshot(ctx, sr)
 	if err != nil {
 		return err

+ 8 - 0
proxy/grpcproxy/util.go

@@ -32,6 +32,14 @@ func getAuthTokenFromClient(ctx context.Context) string {
 	return ""
 }
 
+func withClientAuthToken(ctx context.Context, ctxWithToken context.Context) context.Context {
+	token := getAuthTokenFromClient(ctxWithToken)
+	if token != "" {
+		ctx = context.WithValue(ctx, "token", token)
+	}
+	return ctx
+}
+
 type proxyTokenCredential struct {
 	token string
 }

+ 1 - 5
proxy/grpcproxy/watch_broadcast.go

@@ -58,11 +58,7 @@ func newWatchBroadcast(wp *watchProxy, w *watcher, update func(*watchBroadcast))
 			clientv3.WithCreatedNotify(),
 		}
 
-		// Forward a token from client to server.
-		token := getAuthTokenFromClient(w.wps.stream.Context())
-		if token != "" {
-			cctx = context.WithValue(cctx, "token", token)
-		}
+		cctx = withClientAuthToken(cctx, w.wps.stream.Context())
 
 		wch := wp.cw.Watch(cctx, w.wr.key, opts...)