create token struct from TokenResponse

- Token converts expiresIn to a timestamp
- Store the token with the endpoint suffix

Author: burmudar

internal/oauthdevice/device_flow.go

@@ -26,7 +26,7 @@ const (
 	wellKnownPath = "/.well-known/openid-configuration"
 
 	// Key used to store the token in the store
-	KeyOAuth = "oauth"
+	KeyOAuth = "Sourcegraph CLI key storage"
 
 	GrantTypeDeviceCode string = "urn:ietf:params:oauth:grant-type:device_code"
 
@@ -55,16 +55,20 @@ type DeviceAuthResponse struct {
 	ExpiresIn               int    `json:"expires_in"`
 	Interval                int    `json:"interval"`
 }
-type Token struct {
+
+type TokenResponse struct {
 	AccessToken  string `json:"access_token"`
 	RefreshToken string `json:"refresh_token,omitempty"`
 	ExpiresIn    int    `json:"expires_in,omitempty"`
+	TokenType    string `json:"token_type"`
+	Scope        string `json:"scope,omitempty"`
 }
 
-type TokenResponse struct {
-	Token
-	TokenType string `json:"token_type"`
-	Scope     string `json:"scope,omitempty"`
+type Token struct {
+	Endpoint     string    `json:"endpoint"`
+	AccessToken  string    `json:"access_token"`
+	RefreshToken string    `json:"refresh_token,omitempty"`
+	ExpiresAt    time.Time `json:"expires_at"`
 }
 
 type ErrorResponse struct {
@@ -76,7 +80,7 @@ type Client interface {
 	Discover(ctx context.Context, endpoint string) (*OIDCConfiguration, error)
 	Start(ctx context.Context, endpoint string, scopes []string) (*DeviceAuthResponse, error)
 	Poll(ctx context.Context, endpoint, deviceCode string, interval time.Duration, expiresIn int) (*TokenResponse, error)
-	Refresh(ctx context.Context, endpoint, refreshToken string) (*TokenResponse, error)
+	Refresh(ctx context.Context, token *Token) (*TokenResponse, error)
 }
 
 type httpClient struct {
@@ -318,22 +322,20 @@ func (c *httpClient) pollOnce(ctx context.Context, tokenEndpoint, deviceCode str
 }
 
 // Refresh exchanges a refresh token for a new access token.
-func (c *httpClient) Refresh(ctx context.Context, endpoint, refreshToken string) (*TokenResponse, error) {
-	endpoint = strings.TrimRight(endpoint, "/")
-
-	config, err := c.Discover(ctx, endpoint)
+func (c *httpClient) Refresh(ctx context.Context, token *Token) (*TokenResponse, error) {
+	config, err := c.Discover(ctx, token.Endpoint)
 	if err != nil {
-		return nil, errors.Wrap(err, "OIDC discovery failed")
+		errors.Wrap(err, "failed to discover OIDC configuration")
 	}
 
 	if config.TokenEndpoint == "" {
-		return nil, errors.New("token endpoint not found in OIDC configuration")
+		errors.New("OIDC configuration has no token endpoint")
 	}
 
 	data := url.Values{}
 	data.Set("client_id", c.clientID)
 	data.Set("grant_type", "refresh_token")
-	data.Set("refresh_token", refreshToken)
+	data.Set("refresh_token", token.RefreshToken)
 
 	req, err := http.NewRequestWithContext(ctx, "POST", config.TokenEndpoint, strings.NewReader(data.Encode()))
 	if err != nil {
@@ -366,7 +368,16 @@ func (c *httpClient) Refresh(ctx context.Context, endpoint, refreshToken string)
 		return nil, errors.Wrap(err, "parsing refresh token response")
 	}
 
-	return &tokenResp, nil
+	return &tokenResp, err
+}
+
+func (t *TokenResponse) Token(endpoint string) *Token {
+	return &Token{
+		Endpoint:     strings.TrimRight(endpoint, "/"),
+		RefreshToken: t.RefreshToken,
+		AccessToken:  t.AccessToken,
+		ExpiresAt:    time.Now().Add(time.Second * time.Duration(t.ExpiresIn)),
+	}
 }
 
 func StoreToken(store *keyring.Store, token *Token) error {
@@ -375,13 +386,18 @@ func StoreToken(store *keyring.Store, token *Token) error {
 		return errors.Wrap(err, "failed to marshal token")
 	}
 
-	// TODO(burmudar): do we need a suffix that is the endpoint? ex. oauth-sourcegraph.com
-	return store.Set(KeyOAuth, data)
+	if token.Endpoint == "" {
+		return errors.New("token endpoint cannot be empty when storing the token")
+	}
+
+	key := fmt.Sprintf("%s <%s>", KeyOAuth, token.Endpoint)
+	return store.Set(key, data)
 }
 
-func LoadToken(store *keyring.Store) (*Token, error) {
+func LoadToken(store *keyring.Store, endpoint string) (*Token, error) {
+	key := fmt.Sprintf("%s <%s>", KeyOAuth, endpoint)
 	var t Token
-	data, err := store.Get(KeyOAuth)
+	data, err := store.Get(key)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get token from store")
 	}

internal/oauthdevice/device_flow_test.go

@@ -266,12 +266,10 @@ func TestStart_NoDeviceEndpoint(t *testing.T) {
 
 func TestPoll_Success(t *testing.T) {
 	wantToken := TokenResponse{
-		Token: Token{
-			AccessToken: "test-access-token",
-			ExpiresIn:   3600,
-		},
-		Scope:     "read write",
-		TokenType: "Bearer",
+		AccessToken: "test-access-token",
+		ExpiresIn:   3600,
+		Scope:       "read write",
+		TokenType:   "Bearer",
 	}
 
 	server := newTestServer(t, testServerOptions{
@@ -315,6 +313,7 @@ func TestPoll_Success(t *testing.T) {
 	if resp.TokenType != wantToken.TokenType {
 		t.Errorf("TokenType = %q, want %q", resp.TokenType, wantToken.TokenType)
 	}
+
 }
 
 func TestPoll_AuthorizationPending(t *testing.T) {
@@ -337,8 +336,8 @@ func TestPoll_AuthorizationPending(t *testing.T) {
 				}
 
 				json.NewEncoder(w).Encode(TokenResponse{
-					Token:     Token{AccessToken: "test-access-token"},
-					TokenType: "Bearer",
+					AccessToken: "test-access-token",
+					TokenType:   "Bearer",
 				})
 			},
 		},
@@ -379,8 +378,8 @@ func TestPoll_SlowDown(t *testing.T) {
 				}
 
 				json.NewEncoder(w).Encode(TokenResponse{
-					Token:     Token{AccessToken: "test-access-token"},
-					TokenType: "Bearer",
+					AccessToken: "test-access-token",
+					TokenType:   "Bearer",
 				})
 			},
 		},
@@ -527,20 +526,24 @@ func TestRefresh_Success(t *testing.T) {
 
 				w.Header().Set("Content-Type", "application/json")
 				json.NewEncoder(w).Encode(TokenResponse{
-					Token: Token{
-						AccessToken:  "new-access-token",
-						RefreshToken: "new-refresh-token",
-						ExpiresIn:    3600,
-					},
-					TokenType: "Bearer",
+					AccessToken:  "new-access-token",
+					RefreshToken: "new-refresh-token",
+					ExpiresIn:    3600,
+					TokenType:    "Bearer",
 				})
 			},
 		},
 	})
 	defer server.Close()
 
 	client := NewClient(DefaultClientID)
-	resp, err := client.Refresh(context.Background(), server.URL, "test-refresh-token")
+	token := &Token{
+		Endpoint:     server.URL,
+		AccessToken:  "new-access-token",
+		RefreshToken: "test-refresh-token",
+		ExpiresAt:    time.Now().Add(time.Second * time.Duration(3600)),
+	}
+	resp, err := client.Refresh(context.Background(), token)
 	if err != nil {
 		t.Fatalf("Refresh() error = %v", err)
 	}