From 5111ac92ed721c8b08de077bb36864a04d5f069b Mon Sep 17 00:00:00 2001
From: Hyper <hyper.jiang@huolala.cn>
Date: Sun, 22 Aug 2021 17:53:35 +0800
Subject: [PATCH] add WithAccessKeySecret to support using access key

---
 README.md      | 12 ++++++++++++
 client.go      | 19 ++++++++++++-------
 client_test.go |  7 ++++++-
 options.go     |  9 +++++++++
 sign.go        | 35 +++++++++++++++++++++++++++++++++++
 5 files changed, 74 insertions(+), 8 deletions(-)
 create mode 100644 sign.go

diff --git a/README.md b/README.md
index 62d27d1..5b3cfc5 100644
--- a/README.md
+++ b/README.md
@@ -84,3 +84,15 @@ app := lunar.New("myAppID")
 
 app.UseCache(lunar.NewFileCache("myAppID", "/tmp"))
 ```
+
+## Enable Access Key
+
+You can use `WithAccessKeySecret` to enable access key feature:
+
+```
+app := lunar.New(
+	"myAppID",
+	lunar.WithServer("localhost:8080"),
+	lunar.WithAccessKeySecret("mySecret"),
+)
+```
diff --git a/client.go b/client.go
index 85bc4c8..82599be 100644
--- a/client.go
+++ b/client.go
@@ -43,7 +43,8 @@ func NewApolloClient(appID string, opts ...Option) *ApolloClient {
 	return c
 }
 
-func (c *ApolloClient) get(url string, result interface{}) error {
+func (c *ApolloClient) get(pathWithQuery string, result interface{}) error {
+	url := c.Server + pathWithQuery
 	c.Logger.Printf("%s", url)
 
 	req, err := http.NewRequest(http.MethodGet, url, nil)
@@ -51,6 +52,13 @@ func (c *ApolloClient) get(url string, result interface{}) error {
 		return err
 	}
 
+	if len(c.AccessKeySecret) > 0 {
+		headers := buildHeaders(pathWithQuery, c.AppID, c.AccessKeySecret)
+		for k, v := range headers {
+			req.Header.Add(k, v)
+		}
+	}
+
 	resp, err := c.Client.Do(req)
 	if err != nil {
 		return err
@@ -73,8 +81,7 @@ func (c *ApolloClient) get(url string, result interface{}) error {
 
 // GetCachedItems gets cached configs from apollo
 func (c *ApolloClient) GetCachedItems(namespace string) (Items, error) {
-	url := fmt.Sprintf("%s/configfiles/json/%s/%s/%s?ip=%s",
-		c.Server,
+	url := fmt.Sprintf("/configfiles/json/%s/%s/%s?ip=%s",
 		url.QueryEscape(c.AppID),
 		url.QueryEscape(c.Cluster),
 		url.QueryEscape(namespace),
@@ -102,8 +109,7 @@ func (c *ApolloClient) GetNamespace(namespace string, releaseKey string) (*Names
 		namespace = defaultNamespace
 	}
 
-	url := fmt.Sprintf("%s/configs/%s/%s/%s?releaseKey=%s&ip=%s",
-		c.Server,
+	url := fmt.Sprintf("/configs/%s/%s/%s?releaseKey=%s&ip=%s",
 		url.QueryEscape(c.AppID),
 		url.QueryEscape(c.Cluster),
 		url.QueryEscape(namespace),
@@ -138,8 +144,7 @@ func (c *ApolloClient) GetNotifications(ns Notifications) (Notifications, error)
 		ns = append(ns, Notification{Namespace: defaultNamespace, NotificationID: defaultNotificationID})
 	}
 
-	url := fmt.Sprintf("%s/notifications/v2?appId=%s&cluster=%s&notifications=%s",
-		c.Server,
+	url := fmt.Sprintf("/notifications/v2?appId=%s&cluster=%s&notifications=%s",
 		url.QueryEscape(c.AppID),
 		url.QueryEscape(c.Cluster),
 		url.QueryEscape(ns.String()),
diff --git a/client_test.go b/client_test.go
index ff3eaa5..cc32565 100644
--- a/client_test.go
+++ b/client_test.go
@@ -24,7 +24,12 @@ func TestApolloClientTestSuite(t *testing.T) {
 
 // SetupSuite run once at the very start of the testing suite, before any tests are run.
 func (ts *ApolloClientTestSuite) SetupSuite() {
-	ts.client = NewApolloClient("SampleApp", WithServer("localhost:8080"), WithLogger(Printf))
+	ts.client = NewApolloClient(
+		"SampleApp",
+		WithServer("localhost:8080"),
+		WithLogger(Printf),
+		WithAccessKeySecret("12848b38781e4daf9d05054580282a8e"),
+	)
 }
 
 // TearDownSuite run once at the very end of the testing suite, after all tests have been run.
diff --git a/options.go b/options.go
index eefe3ad..40a7027 100644
--- a/options.go
+++ b/options.go
@@ -1,6 +1,7 @@
 package lunar
 
 import (
+	"strings"
 	"time"
 )
 
@@ -18,6 +19,7 @@ const (
 type Options struct {
 	Server           string
 	Cluster          string
+	AccessKeySecret  string
 	Logger           Logger
 	ClientTimeout    time.Duration
 	LongPollInterval time.Duration
@@ -76,3 +78,10 @@ func WithLongPollInterval(interval time.Duration) Option {
 		o.LongPollInterval = interval
 	}
 }
+
+// WithAccessKeySecret sets access key secret
+func WithAccessKeySecret(secret string) Option {
+	return func(o *Options) {
+		o.AccessKeySecret = strings.TrimSpace(secret)
+	}
+}
diff --git a/sign.go b/sign.go
new file mode 100644
index 0000000..ebaba39
--- /dev/null
+++ b/sign.go
@@ -0,0 +1,35 @@
+package lunar
+
+import (
+	"crypto/hmac"
+	"crypto/sha1"
+	"encoding/base64"
+	"fmt"
+	"time"
+)
+
+const (
+	AuthorizationFormat = "Apollo %s:%s"
+	Delimiter           = "\n"
+)
+
+func sign(timestamp, pathWithQuery, secret string) string {
+	stringToSign := timestamp + Delimiter + pathWithQuery
+
+	h := hmac.New(sha1.New, []byte(secret))
+	h.Write([]byte(stringToSign))
+
+	return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}
+
+func buildHeaders(pathWithQuery, appID, secret string) map[string]string {
+	ms := time.Now().UnixNano() / int64(time.Millisecond)
+	timestamp := fmt.Sprintf("%d", ms)
+	signature := sign(timestamp, pathWithQuery, secret)
+
+	m := make(map[string]string)
+	m["Authorization"] = fmt.Sprintf(AuthorizationFormat, appID, signature)
+	m["Timestamp"] = timestamp
+
+	return m
+}