Skip to content

Commit

Permalink
feat: Add a Token struct to allow the CLI to differentiate between …
Browse files Browse the repository at this point in the history
…`Bearer` and `APIKey` tokens (#66)

This PR adds a `Token` struct to allow the CLI to differentiate between the auth methods.

The CLI will have a different download flow depending on if the auth is using a bearer token vs an API key. For the bearer token the `team_name` is configured using the `cloudquery switch` command. For the API key, the key is associated with a team name already.
  • Loading branch information
mnorbury authored Nov 14, 2023
1 parent 46a8d27 commit bf70ed2
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 14 deletions.
33 changes: 24 additions & 9 deletions auth/token.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,21 @@ type tokenResponse struct {
ProjectID string `json:"project_id"`
}

type TokenType int

const (
Undefined TokenType = iota
BearerToken
APIKey
)

var UndefinedToken = Token{Type: Undefined, Value: ""}

type Token struct {
Type TokenType
Value string
}

type TokenClient struct {
url string
apiKey string
Expand All @@ -45,37 +60,37 @@ func NewTokenClient() *TokenClient {

// GetToken returns the ID token
// If CLOUDQUERY_API_KEY is set, it returns that value, otherwise it returns an ID token generated from the refresh token.
func (tc *TokenClient) GetToken() (string, error) {
func (tc *TokenClient) GetToken() (Token, error) {
if token := os.Getenv(EnvVarCloudQueryAPIKey); token != "" {
return token, nil
return Token{Type: APIKey, Value: token}, nil
}

// If the token is not expired, return it
if !tc.expiresAt.IsZero() && tc.expiresAt.Sub(time.Now().UTC()) > ExpiryBuffer {
return tc.idToken, nil
return Token{Type: BearerToken, Value: tc.idToken}, nil
}

refreshToken, err := ReadRefreshToken()
if err != nil {
return "", fmt.Errorf("failed to read refresh token: %w. Hint: You may need to run `cloudquery login` or set %s", err, EnvVarCloudQueryAPIKey)
return UndefinedToken, fmt.Errorf("failed to read refresh token: %w. Hint: You may need to run `cloudquery login` or set %s", err, EnvVarCloudQueryAPIKey)
}
if refreshToken == "" {
return "", fmt.Errorf("authentication token not found. Hint: You may need to run `cloudquery login` or set %s", EnvVarCloudQueryAPIKey)
return UndefinedToken, fmt.Errorf("authentication token not found. Hint: You may need to run `cloudquery login` or set %s", EnvVarCloudQueryAPIKey)
}
tokenResponse, err := tc.generateToken(refreshToken)
if err != nil {
return "", fmt.Errorf("failed to sign in with custom token: %w", err)
return UndefinedToken, fmt.Errorf("failed to sign in with custom token: %w", err)
}

if err := SaveRefreshToken(tokenResponse.RefreshToken); err != nil {
return "", fmt.Errorf("failed to save refresh token: %w", err)
return UndefinedToken, fmt.Errorf("failed to save refresh token: %w", err)
}

if err := tc.updateIDToken(tokenResponse); err != nil {
return "", fmt.Errorf("failed to update ID token: %w", err)
return UndefinedToken, fmt.Errorf("failed to update ID token: %w", err)
}

return tc.idToken, nil
return Token{Type: BearerToken, Value: tc.idToken}, nil
}

func (tc *TokenClient) generateToken(refreshToken string) (*tokenResponse, error) {
Expand Down
10 changes: 5 additions & 5 deletions auth/token_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ func TestTokenClient_EnvironmentVariable(t *testing.T) {
token, err := NewTokenClient().GetToken()
require.NoError(t, err)

require.Equal(t, "my_token", token)
require.Equal(t, Token{Type: APIKey, Value: "my_token"}, token)
}

func TestTokenClient_GetToken_ShortExpiry(t *testing.T) {
Expand All @@ -66,13 +66,13 @@ func TestTokenClient_GetToken_ShortExpiry(t *testing.T) {

token, err := tc.GetToken()
require.NoError(t, err)
require.Equal(t, "my_id_token_0", token, "first token")
require.Equal(t, Token{Type: BearerToken, Value: "my_id_token_0"}, token, "first token")

tc.expiresAt = t0

token, err = tc.GetToken()
require.NoError(t, err)
require.Equal(t, "my_id_token_1", token, "expected to issue new token")
require.Equal(t, Token{Type: BearerToken, Value: "my_id_token_1"}, token, "expected to issue new token")
}

func TestTokenClient_GetToken_LongExpiry(t *testing.T) {
Expand All @@ -89,11 +89,11 @@ func TestTokenClient_GetToken_LongExpiry(t *testing.T) {

token, err := tc.GetToken()
require.NoError(t, err)
require.Equal(t, "my_id_token_0", token, "first token")
require.Equal(t, Token{Type: BearerToken, Value: "my_id_token_0"}, token, "first token")

token, err = tc.GetToken()
require.NoError(t, err)
require.Equal(t, "my_id_token_0", token, "expected to reuse token")
require.Equal(t, Token{Type: BearerToken, Value: "my_id_token_0"}, token, "expected to reuse token")
}

func overrideEnvironmentVariable(t *testing.T, key, value string) func() {
Expand Down

0 comments on commit bf70ed2

Please sign in to comment.