From 93594006dff269452b52cbd59080cd863a9fc0fd Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Mon, 10 Apr 2017 13:57:12 -0700
Subject: [PATCH] embed: register grpc-gateway endpoints for v3lock and
 v3election

---
 embed/serve.go | 37 +++++++++++++++----------------------
 1 file changed, 15 insertions(+), 22 deletions(-)

diff --git a/embed/serve.go b/embed/serve.go
index 16f0b97ce..46634b7c5 100644
--- a/embed/serve.go
+++ b/embed/serve.go
@@ -165,34 +165,27 @@ func servePeerHTTP(l net.Listener, handler http.Handler) error {
 	return srv.Serve(l)
 }
 
+type registerHandlerFunc func(context.Context, *gw.ServeMux, string, []grpc.DialOption) error
+
 func (sctx *serveCtx) registerGateway(opts []grpc.DialOption) (*gw.ServeMux, error) {
 	ctx := sctx.ctx
 	addr := sctx.l.Addr().String()
 	gwmux := gw.NewServeMux()
 
-	err := pb.RegisterKVHandlerFromEndpoint(ctx, gwmux, addr, opts)
-	if err != nil {
-		return nil, err
+	handlers := []registerHandlerFunc{
+		pb.RegisterKVHandlerFromEndpoint,
+		pb.RegisterWatchHandlerFromEndpoint,
+		pb.RegisterLeaseHandlerFromEndpoint,
+		pb.RegisterClusterHandlerFromEndpoint,
+		pb.RegisterMaintenanceHandlerFromEndpoint,
+		pb.RegisterAuthHandlerFromEndpoint,
+		v3lockpb.RegisterLockHandlerFromEndpoint,
+		v3electionpb.RegisterElectionHandlerFromEndpoint,
 	}
-	err = pb.RegisterWatchHandlerFromEndpoint(ctx, gwmux, addr, opts)
-	if err != nil {
-		return nil, err
-	}
-	err = pb.RegisterLeaseHandlerFromEndpoint(ctx, gwmux, addr, opts)
-	if err != nil {
-		return nil, err
-	}
-	err = pb.RegisterClusterHandlerFromEndpoint(ctx, gwmux, addr, opts)
-	if err != nil {
-		return nil, err
-	}
-	err = pb.RegisterMaintenanceHandlerFromEndpoint(ctx, gwmux, addr, opts)
-	if err != nil {
-		return nil, err
-	}
-	err = pb.RegisterAuthHandlerFromEndpoint(ctx, gwmux, addr, opts)
-	if err != nil {
-		return nil, err
+	for _, h := range handlers {
+		if err := h(ctx, gwmux, addr, opts); err != nil {
+			return nil, err
+		}
 	}
 	return gwmux, nil
 }