From 9169633580a2cb116e5ba73a1a848c90221121e2 Mon Sep 17 00:00:00 2001
From: Michael Benford <michael.benford@corp.globo.com>
Date: Tue, 9 Jul 2019 16:00:10 -0300
Subject: [PATCH] Add support for individual passwords for ring shards

---
 ring.go      | 17 ++++++++++++++---
 ring_test.go | 23 +++++++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/ring.go b/ring.go
index b9dc3e8..f7bda31 100644
--- a/ring.go
+++ b/ring.go
@@ -27,6 +27,10 @@ type RingOptions struct {
 	// Map of name => host:port addresses of ring shards.
 	Addrs map[string]string
 
+	// Map of name => password of ring shards, to allow different shards to have
+	// different passwords. It will be ignored if the Password field is set.
+	Passwords map[string]string
+
 	// Frequency of PING commands sent to check shards availability.
 	// Shard is considered down after 3 subsequent failed checks.
 	HeartbeatFrequency time.Duration
@@ -98,12 +102,12 @@ func (opt *RingOptions) init() {
 	}
 }
 
-func (opt *RingOptions) clientOptions() *Options {
+func (opt *RingOptions) clientOptions(shard string) *Options {
 	return &Options{
 		OnConnect: opt.OnConnect,
 
 		DB:       opt.DB,
-		Password: opt.Password,
+		Password: opt.getPassword(shard),
 
 		DialTimeout:  opt.DialTimeout,
 		ReadTimeout:  opt.ReadTimeout,
@@ -118,6 +122,13 @@ func (opt *RingOptions) clientOptions() *Options {
 	}
 }
 
+func (opt *RingOptions) getPassword(shard string) string {
+	if opt.Password == "" {
+		return opt.Passwords[shard]
+	}
+	return opt.Password
+}
+
 //------------------------------------------------------------------------------
 
 type ringShard struct {
@@ -365,7 +376,7 @@ func NewRing(opt *RingOptions) *Ring {
 	ring.cmdsInfoCache = newCmdsInfoCache(ring.cmdsInfo)
 
 	for name, addr := range opt.Addrs {
-		clopt := opt.clientOptions()
+		clopt := opt.clientOptions(name)
 		clopt.Addr = addr
 		ring.shards.Add(name, NewClient(clopt))
 	}
diff --git a/ring_test.go b/ring_test.go
index 3c94dea..ff4b3fb 100644
--- a/ring_test.go
+++ b/ring_test.go
@@ -172,6 +172,29 @@ var _ = Describe("Redis Ring", func() {
 			Expect(ringShard2.Info().Val()).To(ContainSubstring("keys=100"))
 		})
 	})
+
+	Describe("shard passwords", func() {
+		It("can be initialized with a single password, used for all shards", func() {
+			opts := redisRingOptions()
+			opts.Password = "password"
+			ring = redis.NewRing(opts)
+
+			err := ring.Ping().Err()
+			Expect(err).To(MatchError("ERR Client sent AUTH, but no password is set"))
+		})
+
+		It("can be initialized with a passwords map, one for each shard", func() {
+			opts := redisRingOptions()
+			opts.Passwords = map[string]string{
+				"ringShardOne": "password1",
+				"ringShardTwo": "password2",
+			}
+			ring = redis.NewRing(opts)
+
+			err := ring.Ping().Err()
+			Expect(err).To(MatchError("ERR Client sent AUTH, but no password is set"))
+		})
+	})
 })
 
 var _ = Describe("empty Redis Ring", func() {