Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add a Token struct to allow the CLI to differentiate between Bearer and APIKey tokens #66

Merged
merged 1 commit into from
Nov 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading