From 42640e92e19909a6e50c44676ff8f36a70dc4ee5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9E=97=E6=99=93=E7=82=9C?= <linxiaowei.me@gmail.com>
Date: Fri, 29 Apr 2022 15:49:30 +0800
Subject: [PATCH] feat(registry): contrib/registry/zookeeper add digest acl
 support (#1964)

* feat(registry): contrib/registry/zookeeper add digest acl support

* feat(registry): use WithDigestACL to put user and password together
---
 contrib/registry/zookeeper/register.go      | 23 ++++++++++++++++++++-
 contrib/registry/zookeeper/register_test.go |  2 +-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/contrib/registry/zookeeper/register.go b/contrib/registry/zookeeper/register.go
index 341058b5d..adae0050f 100644
--- a/contrib/registry/zookeeper/register.go
+++ b/contrib/registry/zookeeper/register.go
@@ -25,6 +25,8 @@ type options struct {
 	ctx      context.Context
 	rootPath string
 	timeout  time.Duration
+	user     string
+	password string
 }
 
 // WithContext with registry context.
@@ -42,6 +44,14 @@ func WithTimeout(timeout time.Duration) Option {
 	return func(o *options) { o.timeout = timeout }
 }
 
+// WithDigestACL with registry password.
+func WithDigestACL(user string, password string) Option {
+	return func(o *options) {
+		o.user = user
+		o.password = password
+	}
+}
+
 // Registry is consul registry
 type Registry struct {
 	opts     *options
@@ -63,6 +73,12 @@ func New(zkServers []string, opts ...Option) (*Registry, error) {
 	if err != nil {
 		return nil, err
 	}
+	if len(options.user) > 0 && len(options.password) > 0 {
+		err = conn.AddAuth("digest", []byte(options.user+":"+options.password))
+		if err != nil {
+			return nil, err
+		}
+	}
 	return &Registry{
 		opts:     options,
 		conn:     conn,
@@ -182,7 +198,12 @@ func (r *Registry) ensureName(path string, data []byte, flags int32) error {
 		return err
 	}
 	if !exists {
-		_, err := r.conn.Create(path, data, flags, zk.WorldACL(zk.PermAll))
+		var err error
+		if len(r.opts.user) > 0 && len(r.opts.password) > 0 {
+			_, err = r.conn.Create(path, data, flags, zk.DigestACL(zk.PermAll, r.opts.user, r.opts.password))
+		} else {
+			_, err = r.conn.Create(path, data, flags, zk.WorldACL(zk.PermAll))
+		}
 		if err != nil {
 			return err
 		}
diff --git a/contrib/registry/zookeeper/register_test.go b/contrib/registry/zookeeper/register_test.go
index f746b104a..9293e2b93 100644
--- a/contrib/registry/zookeeper/register_test.go
+++ b/contrib/registry/zookeeper/register_test.go
@@ -16,7 +16,7 @@ func TestRegistry(t *testing.T) {
 		Endpoints: []string{"http://127.0.0.1:1111"},
 	}
 
-	r, _ := New([]string{"127.0.0.1:2181"})
+	r, _ := New([]string{"127.0.0.1:2181"}, WithDigestACL("username", "password"))
 
 	w, err := r.Watch(ctx, s.Name)
 	if err != nil {