fix tenant switching

aserto-dev · Aug 15, 2024 · 29518f3 · 29518f3
1 parent 8b2f036
commit 29518f3
Show file tree

Hide file tree

Showing 5 changed files with 106 additions and 51 deletions.
diff --git a/pkg/cc/cc.go b/pkg/cc/cc.go
@@ -88,6 +88,14 @@ func newAuthSettings(auth *config.Auth) *auth0.Settings {
 	return auth.GetSettings()
 }
 
+func (ctx *CommonCtx) TenantID() string {
+	tkn, err := ctx.Token()
+	if err != nil {
+		return ""
+	}
+	return tkn.TenantID
+}
+
 func (ctx *CommonCtx) AccessToken() (string, error) {
 	tkn, err := ctx.Token()
 	if err != nil {

diff --git a/pkg/cc/clients/factory.go b/pkg/cc/clients/factory.go
@@ -18,8 +18,6 @@ import (
 )
 
 type Factory interface {
-	TenantID() string
-
 	TenantClient(ctx context.Context) (*tenant.Client, error)
 	DecisionLogsClient(ctx context.Context) (*dl.Client, error)
 	ControlPlaneClient(ctx context.Context) (*cp.Client, error)
@@ -62,10 +60,6 @@ func NewClientFactory(
 	}, nil
 }
 
-func (c *AsertoFactory) TenantID() string {
-	return c.tenantID
-}
-
 func (c *AsertoFactory) TenantClient(ctx context.Context) (*tenant.Client, error) {
 	options, err := c.options(x.TenantService)
 	if err != nil {

diff --git a/pkg/handlers/config/config.go b/pkg/handlers/config/config.go
@@ -10,6 +10,7 @@ import (
 	"github.com/aserto-dev/aserto/pkg/cc"
 	"github.com/aserto-dev/aserto/pkg/cc/config"
 	errs "github.com/aserto-dev/aserto/pkg/cc/errors"
+	"github.com/aserto-dev/aserto/pkg/handlers/user"
 	"github.com/aserto-dev/go-grpc/aserto/api/v1"
 	account "github.com/aserto-dev/go-grpc/aserto/tenant/account/v1"
 	topazConfig "github.com/aserto-dev/topaz/pkg/cc/config"
@@ -143,7 +144,17 @@ func (cmd *UseConfigCmd) Run(c *cc.CommonCtx) error {
 			return errors.Wrapf(errs.ResolveTenantErr, tenantName)
 		}
 
+		token, err := c.Token()
+		if err != nil {
+			return err
+		}
+
 		c.Config.TenantID = tenant[0].Id
+		token.TenantID = tenant[0].Id
+
+		if err := user.SwitchKeyRing(c, token, tenant[0].Id); err != nil {
+			return err
+		}
 	}
 
 	return c.SaveContextConfig(config.DefaultConfigFilePath)

diff --git a/pkg/handlers/user/login.go b/pkg/handlers/user/login.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/aserto-dev/aserto/pkg/auth0/api"
 	"github.com/aserto-dev/aserto/pkg/auth0/device"
 	"github.com/aserto-dev/aserto/pkg/cc"
 	"github.com/aserto-dev/aserto/pkg/clients/tenant"
@@ -47,61 +48,98 @@ func (d *LoginCmd) Run(c *cc.CommonCtx) error {
 		fmt.Printf("Open browser %s\n", flow.GetVerificationURI())
 	}
 
-	{ // intentionally scoped.
-		ctx, cancel := context.WithTimeout(c.Context, flow.ExpiresIn())
-		defer cancel()
-
-		for {
-			if ok, err := flow.RequestAccessToken(ctx); ok {
-				fmt.Fprintln(c.StdOut(), ".")
-				break
-			} else if err != nil {
-				return err
-			}
-
-			select {
-			case <-time.After(flow.Interval()):
-				fmt.Fprint(c.StdOut(), ".")
-			case <-ctx.Done():
-				return errors.New("canceled")
-			}
+	ctx, cancel := context.WithTimeout(c.Context, flow.ExpiresIn())
+	defer cancel()
+
+	for {
+		if ok, err := flow.RequestAccessToken(ctx); ok {
+			fmt.Fprintln(c.StdOut(), ".")
+			break
+		} else if err != nil {
+			return err
+		}
+
+		select {
+		case <-time.After(flow.Interval()):
+			fmt.Fprint(c.StdOut(), ".")
+		case <-ctx.Done():
+			return errors.New("canceled")
 		}
 	}
 
-	token := flow.AccessToken()
+	if err := UpdateKeyRing(c, flow.AccessToken()); err != nil {
+		return err
+	}
 
-	{ // intentionally scoped.
-		ctx, cancel := context.WithTimeout(c.Context, time.Second*10)
-		defer cancel()
+	c.Con().Info().Msg("Login successful")
 
-		conn, err := tenant.NewClient(
-			ctx,
-			client.WithAddr(c.Environment.TenantService.Address),
-			client.WithTokenAuth(token.Access),
-		)
-		if err != nil {
-			return err
-		}
+	return nil
+}
 
-		if err = getTenantID(ctx, conn, token); err != nil {
-			return errors.Wrapf(err, "get tenant id")
-		}
+func UpdateKeyRing(c *cc.CommonCtx, token *api.Token) error {
+	ctx, cancel := context.WithTimeout(c.Context, time.Second*10)
+	defer cancel()
 
-		if err = GetConnectionKeys(ctx, conn, token); err != nil {
-			return errors.Wrapf(err, "get connection keys")
-		}
+	conn, err := tenant.NewClient(
+		ctx,
+		client.WithAddr(c.Environment.TenantService.Address),
+		client.WithTokenAuth(token.Access),
+	)
+	if err != nil {
+		return err
+	}
 
-		kr, err := keyring.NewKeyRing(c.Auth.Issuer)
-		if err != nil {
-			return err
-		}
+	if err = getTenantID(ctx, conn, token); err != nil {
+		return errors.Wrapf(err, "get tenant id")
+	}
 
-		if err := kr.SetToken(token); err != nil {
-			return err
-		}
+	if err = GetConnectionKeys(ctx, conn, token); err != nil {
+		return errors.Wrapf(err, "get connection keys")
+	}
+
+	kr, err := keyring.NewKeyRing(c.Auth.Issuer)
+	if err != nil {
+		return err
+	}
+
+	if err := kr.SetToken(token); err != nil {
+		return err
+	}
+
+	c.Con().Info().Msg("Switched to tenant-id %q", c.TenantID())
+
+	return nil
+}
+
+func SwitchKeyRing(c *cc.CommonCtx, token *api.Token, tenantID string) error {
+	ctx, cancel := context.WithTimeout(c.Context, time.Second*10)
+	defer cancel()
+
+	conn, err := tenant.NewClient(
+		ctx,
+		client.WithAddr(c.Environment.TenantService.Address),
+		client.WithTokenAuth(token.Access),
+	)
+	if err != nil {
+		return err
+	}
 
-		fmt.Fprintln(c.StdOut(), "Login successful")
+	token.TenantID = tenantID
+
+	if err = GetConnectionKeys(ctx, conn, token); err != nil {
+		return errors.Wrapf(err, "get connection keys")
 	}
 
+	kr, err := keyring.NewKeyRing(c.Auth.Issuer)
+	if err != nil {
+		return err
+	}
+
+	if err := kr.SetToken(token); err != nil {
+		return err
+	}
+
+	c.Con().Info().Msg("Switched to tenant-id %q", c.TenantID())
+
 	return nil
 }
diff --git a/pkg/handlers/user/props.go b/pkg/handlers/user/props.go
@@ -29,7 +29,11 @@ func (cmd *GetCmd) Run(c *cc.CommonCtx) error {
 	case "access-token":
 		propValue, err = c.AccessToken()
 	case "tenant-id":
-		propValue = c.TenantID()
+		token, tokenErr := c.Token()
+		if tokenErr != nil {
+			return tokenErr
+		}
+		propValue = token.TenantID
 	case "authorizer-key":
 		propValue, err = c.AuthorizerAPIKey()
 	case "directory-read-key":