From 757a8e8f5b05d7459800398f66811b9f13d59892 Mon Sep 17 00:00:00 2001
From: Hitoshi Mitake <h.mitake@gmail.com>
Date: Sat, 28 Aug 2021 23:47:14 +0900
Subject: [PATCH] *: implement a retry logic for auth old revision in the
 client

---
 clientv3/retry_interceptor.go          |   4 +-
 clientv3/retry_interceptor_test.go     | 124 +++++++++++++++++++++++++
 etcdserver/api/v3rpc/rpctypes/error.go |   3 +
 etcdserver/api/v3rpc/util.go           |   1 +
 4 files changed, 131 insertions(+), 1 deletion(-)
 create mode 100644 clientv3/retry_interceptor_test.go

diff --git a/clientv3/retry_interceptor.go b/clientv3/retry_interceptor.go
index 21d20300e..4284c9a70 100644
--- a/clientv3/retry_interceptor.go
+++ b/clientv3/retry_interceptor.go
@@ -156,7 +156,9 @@ func (c *Client) shouldRefreshToken(err error, callOpts *options) bool {
 		// which is possible when the client token is cleared somehow
 		return c.authTokenBundle != nil // equal to c.Username != "" && c.Password != ""
 	}
-	return callOpts.retryAuth && rpctypes.Error(err) == rpctypes.ErrInvalidAuthToken
+
+	return callOpts.retryAuth &&
+		(rpctypes.Error(err) == rpctypes.ErrInvalidAuthToken || rpctypes.Error(err) == rpctypes.ErrAuthOldRevision)
 }
 
 // type serverStreamingRetryingStream is the implementation of grpc.ClientStream that acts as a
diff --git a/clientv3/retry_interceptor_test.go b/clientv3/retry_interceptor_test.go
new file mode 100644
index 000000000..b850b56e0
--- /dev/null
+++ b/clientv3/retry_interceptor_test.go
@@ -0,0 +1,124 @@
+package clientv3
+
+import (
+	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
+	"go.etcd.io/etcd/client/v3/credentials"
+	grpccredentials "google.golang.org/grpc/credentials"
+	"testing"
+)
+
+type dummyAuthTokenBundle struct{}
+
+func (d dummyAuthTokenBundle) TransportCredentials() grpccredentials.TransportCredentials {
+	return nil
+}
+
+func (d dummyAuthTokenBundle) PerRPCCredentials() grpccredentials.PerRPCCredentials {
+	return nil
+}
+
+func (d dummyAuthTokenBundle) NewWithMode(mode string) (grpccredentials.Bundle, error) {
+	return nil, nil
+}
+
+func (d dummyAuthTokenBundle) UpdateAuthToken(token string) {
+}
+
+func TestClientShouldRefreshToken(t *testing.T) {
+	type fields struct {
+		authTokenBundle credentials.Bundle
+	}
+	type args struct {
+		err      error
+		callOpts *options
+	}
+
+	optsWithTrue := &options{
+		retryAuth: true,
+	}
+	optsWithFalse := &options{
+		retryAuth: false,
+	}
+
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   bool
+	}{
+		{
+			name: "ErrUserEmpty and non nil authTokenBundle",
+			fields: fields{
+				authTokenBundle: &dummyAuthTokenBundle{},
+			},
+			args: args{rpctypes.ErrGRPCUserEmpty, optsWithTrue},
+			want: true,
+		},
+		{
+			name: "ErrUserEmpty and nil authTokenBundle",
+			fields: fields{
+				authTokenBundle: nil,
+			},
+			args: args{rpctypes.ErrGRPCUserEmpty, optsWithTrue},
+			want: false,
+		},
+		{
+			name: "ErrGRPCInvalidAuthToken and retryAuth",
+			fields: fields{
+				authTokenBundle: nil,
+			},
+			args: args{rpctypes.ErrGRPCInvalidAuthToken, optsWithTrue},
+			want: true,
+		},
+		{
+			name: "ErrGRPCInvalidAuthToken and !retryAuth",
+			fields: fields{
+				authTokenBundle: nil,
+			},
+			args: args{rpctypes.ErrGRPCInvalidAuthToken, optsWithFalse},
+			want: false,
+		},
+		{
+			name: "ErrGRPCAuthOldRevision and retryAuth",
+			fields: fields{
+				authTokenBundle: nil,
+			},
+			args: args{rpctypes.ErrGRPCAuthOldRevision, optsWithTrue},
+			want: true,
+		},
+		{
+			name: "ErrGRPCAuthOldRevision and !retryAuth",
+			fields: fields{
+				authTokenBundle: nil,
+			},
+			args: args{rpctypes.ErrGRPCAuthOldRevision, optsWithFalse},
+			want: false,
+		},
+		{
+			name: "Other error and retryAuth",
+			fields: fields{
+				authTokenBundle: nil,
+			},
+			args: args{rpctypes.ErrGRPCAuthFailed, optsWithTrue},
+			want: false,
+		},
+		{
+			name: "Other error and !retryAuth",
+			fields: fields{
+				authTokenBundle: nil,
+			},
+			args: args{rpctypes.ErrGRPCAuthFailed, optsWithFalse},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &Client{
+				authTokenBundle: tt.fields.authTokenBundle,
+			}
+			if got := c.shouldRefreshToken(tt.args.err, tt.args.callOpts); got != tt.want {
+				t.Errorf("shouldRefreshToken() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/etcdserver/api/v3rpc/rpctypes/error.go b/etcdserver/api/v3rpc/rpctypes/error.go
index d00890850..b7570ff0b 100644
--- a/etcdserver/api/v3rpc/rpctypes/error.go
+++ b/etcdserver/api/v3rpc/rpctypes/error.go
@@ -64,6 +64,7 @@ var (
 	ErrGRPCAuthNotEnabled       = status.New(codes.FailedPrecondition, "etcdserver: authentication is not enabled").Err()
 	ErrGRPCInvalidAuthToken     = status.New(codes.Unauthenticated, "etcdserver: invalid auth token").Err()
 	ErrGRPCInvalidAuthMgmt      = status.New(codes.InvalidArgument, "etcdserver: invalid auth management").Err()
+	ErrGRPCAuthOldRevision      = status.New(codes.InvalidArgument, "etcdserver: revision of auth store is old").Err()
 
 	ErrGRPCNoLeader                   = status.New(codes.Unavailable, "etcdserver: no leader").Err()
 	ErrGRPCNotLeader                  = status.New(codes.FailedPrecondition, "etcdserver: not leader").Err()
@@ -121,6 +122,7 @@ var (
 		ErrorDesc(ErrGRPCAuthNotEnabled):       ErrGRPCAuthNotEnabled,
 		ErrorDesc(ErrGRPCInvalidAuthToken):     ErrGRPCInvalidAuthToken,
 		ErrorDesc(ErrGRPCInvalidAuthMgmt):      ErrGRPCInvalidAuthMgmt,
+		ErrorDesc(ErrGRPCAuthOldRevision):      ErrGRPCAuthOldRevision,
 
 		ErrorDesc(ErrGRPCNoLeader):                   ErrGRPCNoLeader,
 		ErrorDesc(ErrGRPCNotLeader):                  ErrGRPCNotLeader,
@@ -179,6 +181,7 @@ var (
 	ErrPermissionNotGranted = Error(ErrGRPCPermissionNotGranted)
 	ErrAuthNotEnabled       = Error(ErrGRPCAuthNotEnabled)
 	ErrInvalidAuthToken     = Error(ErrGRPCInvalidAuthToken)
+	ErrAuthOldRevision      = Error(ErrGRPCAuthOldRevision)
 	ErrInvalidAuthMgmt      = Error(ErrGRPCInvalidAuthMgmt)
 
 	ErrNoLeader                   = Error(ErrGRPCNoLeader)
diff --git a/etcdserver/api/v3rpc/util.go b/etcdserver/api/v3rpc/util.go
index 281ddc7a0..cae348edb 100644
--- a/etcdserver/api/v3rpc/util.go
+++ b/etcdserver/api/v3rpc/util.go
@@ -77,6 +77,7 @@ var toGRPCErrorMap = map[error]error{
 	auth.ErrAuthNotEnabled:       rpctypes.ErrGRPCAuthNotEnabled,
 	auth.ErrInvalidAuthToken:     rpctypes.ErrGRPCInvalidAuthToken,
 	auth.ErrInvalidAuthMgmt:      rpctypes.ErrGRPCInvalidAuthMgmt,
+	auth.ErrAuthOldRevision:      rpctypes.ErrGRPCAuthOldRevision,
 }
 
 func togRPCError(err error) error {