From 39d96b2557d4e1863c43880106bf1ce815ad5674 Mon Sep 17 00:00:00 2001
From: Benjamin Wang <wachao@vmware.com>
Date: Tue, 20 Dec 2022 10:41:16 +0800
Subject: [PATCH] clientv3: add protection code to prevent SnapshotWithVersion
 from panicking

Signed-off-by: Benjamin Wang <wachao@vmware.com>
---
 client/v3/maintenance.go          |  6 ++++--
 client/v3/snapshot/v3_snapshot.go | 15 +--------------
 2 files changed, 5 insertions(+), 16 deletions(-)

diff --git a/client/v3/maintenance.go b/client/v3/maintenance.go
index 388106305..f47808d5f 100644
--- a/client/v3/maintenance.go
+++ b/client/v3/maintenance.go
@@ -239,6 +239,7 @@ func (m *maintenance) SnapshotWithVersion(ctx context.Context) (*SnapshotRespons
 	resp, err := ss.Recv()
 	if err != nil {
 		m.logAndCloseWithError(err, pw)
+		return nil, err
 	}
 	go func() {
 		// Saving response is blocking
@@ -260,10 +261,11 @@ func (m *maintenance) SnapshotWithVersion(ctx context.Context) (*SnapshotRespons
 			}
 		}
 	}()
+
 	return &SnapshotResponse{
-		Header:   resp.Header,
+		Header:   resp.GetHeader(),
 		Snapshot: &snapshotReadCloser{ctx: ctx, ReadCloser: pr},
-		Version:  resp.Version,
+		Version:  resp.GetVersion(),
 	}, err
 }
 
diff --git a/client/v3/snapshot/v3_snapshot.go b/client/v3/snapshot/v3_snapshot.go
index cc27f643a..10b51885e 100644
--- a/client/v3/snapshot/v3_snapshot.go
+++ b/client/v3/snapshot/v3_snapshot.go
@@ -68,7 +68,7 @@ func SaveWithVersion(ctx context.Context, lg *zap.Logger, cfg clientv3.Config, d
 	start := time.Now()
 	resp, err := cli.SnapshotWithVersion(ctx)
 	if err != nil {
-		return resp.Version, err
+		return "", err
 	}
 	defer resp.Snapshot.Close()
 	lg.Info("fetching snapshot", zap.String("endpoint", cfg.Endpoints[0]))
@@ -99,16 +99,3 @@ func SaveWithVersion(ctx context.Context, lg *zap.Logger, cfg clientv3.Config, d
 	lg.Info("saved", zap.String("path", dbPath))
 	return resp.Version, nil
 }
-
-// Save fetches snapshot from remote etcd server and saves data
-// to target path. If the context "ctx" is canceled or timed out,
-// snapshot save stream will error out (e.g. context.Canceled,
-// context.DeadlineExceeded). Make sure to specify only one endpoint
-// in client configuration. Snapshot API must be requested to a
-// selected node, and saved snapshot is the point-in-time state of
-// the selected node.
-// Deprecated: Use SaveWithVersion instead.
-func Save(ctx context.Context, lg *zap.Logger, cfg clientv3.Config, dbPath string) error {
-	_, err := SaveWithVersion(ctx, lg, cfg, dbPath)
-	return err
-}