From ab8e5a4f8eeec1c197720f736d49a6b377f9838d Mon Sep 17 00:00:00 2001
From: Andrew Lytvynov <andrew@awly.dev>
Date: Mon, 17 May 2021 17:06:38 -0700
Subject: [PATCH] client/v3: do not overwrite authTokenBundle on dial

`Client.dial` can be called multiple times. For example, from regular
`NewClient` and [from the `Maintenance`
wrapper](https://github.com/etcd-io/etcd/blob/6d451ab61d2da992622b3e7dbf85ded04ebcc515/client/v3/maintenance.go#L81-L84).

This ends up creating two (if `Client.dial` is called twice) independent
`credentials.Bundle` instances:
- the first one is wired into `PerRPCCredentials` callback in gRPC
  client.
- the second one is in `Client.authTokenBundle` and receives the latest
  auth token updates.

When using username/password authentication and the server is configured
to issue JWTs, token rotation logic breaks because of the above
`credentials.Bundle` confusion.
---
 client/v3/client.go      |  4 +--
 client/v3/client_test.go | 53 +++++++++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/client/v3/client.go b/client/v3/client.go
index 75e3c97e2..290e7b68e 100644
--- a/client/v3/client.go
+++ b/client/v3/client.go
@@ -283,8 +283,7 @@ func (c *Client) dial(creds grpccredentials.TransportCredentials, dopts ...grpc.
 	if err != nil {
 		return nil, fmt.Errorf("failed to configure dialer: %v", err)
 	}
-	if c.Username != "" && c.Password != "" {
-		c.authTokenBundle = credentials.NewBundle(credentials.Config{})
+	if c.authTokenBundle != nil {
 		opts = append(opts, grpc.WithPerRPCCredentials(c.authTokenBundle.PerRPCCredentials()))
 	}
 
@@ -365,6 +364,7 @@ func newClient(cfg *Config) (*Client, error) {
 	if cfg.Username != "" && cfg.Password != "" {
 		client.Username = cfg.Username
 		client.Password = cfg.Password
+		client.authTokenBundle = credentials.NewBundle(credentials.Config{})
 	}
 	if cfg.MaxCallSendMsgSize > 0 || cfg.MaxCallRecvMsgSize > 0 {
 		if cfg.MaxCallRecvMsgSize > 0 && cfg.MaxCallSendMsgSize > cfg.MaxCallRecvMsgSize {
diff --git a/client/v3/client_test.go b/client/v3/client_test.go
index b2ff7ef17..ba31b68d8 100644
--- a/client/v3/client_test.go
+++ b/client/v3/client_test.go
@@ -17,13 +17,15 @@ package clientv3
 import (
 	"context"
 	"fmt"
-	"go.uber.org/zap"
 	"net"
+	"path/filepath"
 	"testing"
 	"time"
 
+	"go.etcd.io/etcd/api/v3/etcdserverpb"
 	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
 	"go.etcd.io/etcd/client/pkg/v3/testutil"
+	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest"
 
 	"google.golang.org/grpc"
@@ -198,3 +200,52 @@ func TestZapWithLogger(t *testing.T) {
 		t.Errorf("WithZapLogger should modify *zap.Logger")
 	}
 }
+
+func TestAuthTokenBundleNoOverwrite(t *testing.T) {
+	// Create a mock AuthServer to handle Authenticate RPCs.
+	lis, err := net.Listen("unix", filepath.Join(t.TempDir(), "etcd-auth-test:0"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer lis.Close()
+	addr := "unix://" + lis.Addr().String()
+	srv := grpc.NewServer()
+	etcdserverpb.RegisterAuthServer(srv, mockAuthServer{})
+	go srv.Serve(lis)
+	defer srv.Stop()
+
+	// Create a client, which should call Authenticate on the mock server to
+	// exchange username/password for an auth token.
+	c, err := NewClient(t, Config{
+		DialTimeout: 5 * time.Second,
+		Endpoints:   []string{addr},
+		Username:    "foo",
+		Password:    "bar",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+	oldTokenBundle := c.authTokenBundle
+
+	// Call the public Dial again, which should preserve the original
+	// authTokenBundle.
+	gc, err := c.Dial(addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer gc.Close()
+	newTokenBundle := c.authTokenBundle
+
+	if oldTokenBundle != newTokenBundle {
+		t.Error("Client.authTokenBundle has been overwritten during Client.Dial")
+	}
+}
+
+type mockAuthServer struct {
+	*etcdserverpb.UnimplementedAuthServer
+}
+
+func (mockAuthServer) Authenticate(context.Context, *etcdserverpb.AuthenticateRequest) (*etcdserverpb.AuthenticateResponse, error) {
+	return &etcdserverpb.AuthenticateResponse{Token: "mock-token"}, nil
+}