From 366fb228e509915b3f1b685b3f7266524645ebb6 Mon Sep 17 00:00:00 2001
From: Ben Toogood <bentoogood@gmail.com>
Date: Wed, 13 May 2020 17:54:47 +0100
Subject: [PATCH] Auth: Set address

---
 auth/options.go         |  9 ++++++++
 auth/service/service.go | 50 +++++++++++++++++++++++------------------
 2 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/auth/options.go b/auth/options.go
index da0982b7..bb120241 100644
--- a/auth/options.go
+++ b/auth/options.go
@@ -44,10 +44,19 @@ type Options struct {
 	Store store.Store
 	// Client to use for RPC
 	Client client.Client
+	// Addrs sets the addresses of auth
+	Addrs []string
 }
 
 type Option func(o *Options)
 
+// Addrs is the auth addresses to use
+func Addrs(addrs ...string) Option {
+	return func(o *Options) {
+		o.Addrs = addrs
+	}
+}
+
 // Namespace the service belongs to
 func Namespace(n string) Option {
 	return func(o *Options) {
diff --git a/auth/service/service.go b/auth/service/service.go
index b8a163b0..b6b886fc 100644
--- a/auth/service/service.go
+++ b/auth/service/service.go
@@ -19,7 +19,27 @@ import (
 
 // NewAuth returns a new instance of the Auth service
 func NewAuth(opts ...auth.Option) auth.Auth {
-	return &svc{options: auth.NewOptions(opts...)}
+	options := auth.NewOptions(opts...)
+
+	addrs := options.Addrs
+	if len(addrs) == 0 {
+		addrs = []string{"127.0.0.1:8010"}
+	}
+
+	svc := &svc{options: options, addrs: addrs}
+
+	// load rules periodically from the auth service
+	go func() {
+		ruleTimer := time.NewTicker(time.Second * 30)
+
+		for {
+			time.Sleep(jitter.Do(time.Second * 5))
+			svc.loadRules()
+			<-ruleTimer.C
+		}
+	}()
+
+	return svc
 }
 
 // svc is the service implementation of the Auth interface
@@ -28,6 +48,7 @@ type svc struct {
 	auth    pb.AuthService
 	rule    pb.RulesService
 	jwt     token.Provider
+	addrs   []string
 
 	rules []*pb.Rule
 	sync.Mutex
@@ -54,21 +75,6 @@ func (s *svc) Init(opts ...auth.Option) {
 	if key := s.options.PublicKey; len(key) > 0 {
 		s.jwt = jwt.NewTokenProvider(token.WithPublicKey(key))
 	}
-
-	// load rules periodically from the auth service
-	go func() {
-		ruleTimer := time.NewTicker(time.Second * 30)
-
-		for {
-			// jitter for up to 5 seconds, this stops
-			// all the services calling the auth service
-			// at the exact same time
-			time.Sleep(jitter.Do(time.Second * 5))
-			s.loadRules()
-
-			<-ruleTimer.C
-		}
-	}()
 }
 
 func (s *svc) Options() auth.Options {
@@ -89,7 +95,7 @@ func (s *svc) Generate(id string, opts ...auth.GenerateOption) (*auth.Account, e
 		Metadata:  options.Metadata,
 		Provider:  options.Provider,
 		Namespace: options.Namespace,
-	})
+	}, client.WithAddress(s.addrs...))
 	if err != nil {
 		return nil, err
 	}
@@ -108,7 +114,7 @@ func (s *svc) Grant(role string, res *auth.Resource) error {
 			Name:      res.Name,
 			Endpoint:  res.Endpoint,
 		},
-	})
+	}, client.WithAddress(s.addrs...))
 	return err
 }
 
@@ -123,7 +129,7 @@ func (s *svc) Revoke(role string, res *auth.Resource) error {
 			Name:      res.Name,
 			Endpoint:  res.Endpoint,
 		},
-	})
+	}, client.WithAddress(s.addrs...))
 	return err
 }
 
@@ -193,7 +199,7 @@ func (s *svc) Inspect(token string) (*auth.Account, error) {
 
 	// the token is not a JWT or we do not have the keys to decode it,
 	// fall back to the auth service
-	rsp, err := s.auth.Inspect(context.TODO(), &pb.InspectRequest{Token: token})
+	rsp, err := s.auth.Inspect(context.TODO(), &pb.InspectRequest{Token: token}, client.WithAddress(s.addrs...))
 	if err != nil {
 		return nil, err
 	}
@@ -209,7 +215,7 @@ func (s *svc) Token(opts ...auth.TokenOption) (*auth.Token, error) {
 		Secret:       options.Secret,
 		RefreshToken: options.RefreshToken,
 		TokenExpiry:  int64(options.Expiry.Seconds()),
-	})
+	}, client.WithAddress(s.addrs...))
 	if err != nil {
 		return nil, err
 	}
@@ -274,7 +280,7 @@ func (s *svc) listRules(filters ...string) []*pb.Rule {
 
 // loadRules retrieves the rules from the auth service
 func (s *svc) loadRules() {
-	rsp, err := s.rule.List(context.TODO(), &pb.ListRequest{})
+	rsp, err := s.rule.List(context.TODO(), &pb.ListRequest{}, client.WithAddress(s.addrs...))
 	s.Lock()
 	defer s.Unlock()