fix: access token logic

src/browser/types.ts

@@ -20,9 +20,14 @@ export const featureFlagDefaults: FeatureFlags = {
     newInject: false,
 }
 
+export interface AccessToken {
+    id: string
+    token: string
+}
+
 /** A map where the key is the server URL and the value is the token. */
 export interface AccessTokens {
-    [url: string]: string
+    [url: string]: AccessToken
 }
 
 // TODO(chris) Switch to Partial<StorageItems> to eliminate bugs caused by

src/extension/scripts/background.tsx

@@ -166,6 +166,19 @@ storage.addSyncMigration((items, set, remove) => {
         set({ accessTokens: {} })
     }
 
+    if (items.accessTokens) {
+        const accessTokens = {}
+
+        for (const url of Object.keys(items.accessTokens)) {
+            const token = items.accessTokens[url]
+            if (typeof token !== 'string' && token.id && token.token) {
+                accessTokens[url] = token
+            }
+        }
+
+        set({ accessTokens })
+    }
+
     if (items.phabricatorURL) {
         remove('phabricatorURL')
 

src/shared/auth/access_token.ts

@@ -2,20 +2,21 @@ import { omit } from 'lodash'
 import { Observable } from 'rxjs'
 import { switchMap } from 'rxjs/operators'
 import storage from '../../browser/storage'
+import { AccessToken } from '../../browser/types'
 
-export const getAccessToken = (url: string): Observable<string | undefined> =>
+export const getAccessToken = (url: string): Observable<AccessToken | undefined> =>
     new Observable(observer => {
         storage.getSync(items => {
             observer.next(items.accessTokens[url])
             observer.complete()
         })
     })
 
-export const setAccessToken = (url: string) => (tokens: Observable<string>): Observable<string> =>
+export const setAccessToken = (url: string) => (tokens: Observable<AccessToken>): Observable<AccessToken> =>
     tokens.pipe(
         switchMap(
             token =>
-                new Observable<string>(observer => {
+                new Observable(observer => {
                     storage.getSync(({ accessTokens }) =>
                         storage.setSync({ accessTokens: { ...accessTokens, [url]: token } }, () => {
                             observer.next(token)

src/shared/backend/auth.ts

@@ -1,34 +1,70 @@
+import { Observable } from 'rxjs'
 import { map } from 'rxjs/operators'
+import { AccessToken } from '../../browser/types'
 import { GQL } from '../../types/gqlschema'
 import { getPlatformName } from '../util/context'
 import { memoizeObservable } from '../util/memoize'
 import { getContext } from './context'
 import { createAggregateError } from './errors'
-import { mutateGraphQL } from './graphql'
+import { mutateGraphQL, queryGraphQL } from './graphql'
 
 /**
  * Create an access token for the current user on the currently configured
  * sourcegraph instance.
  */
-export const createAccessToken = memoizeObservable((userID: GQL.ID) =>
-    mutateGraphQL({
-        ctx: getContext({ repoKey: '' }),
-        request: `
+export const createAccessToken = memoizeObservable(
+    (userID: GQL.ID): Observable<AccessToken> =>
+        mutateGraphQL({
+            ctx: getContext({ repoKey: '' }),
+            request: `
         mutation CreateAccessToken($userID: ID!, $scopes: [String!]!, $note: String!) {
             createAccessToken(user: $userID, scopes: $scopes, note: $note) {
                 id
                 token
             }
         }
         `,
-        variables: { userID, scopes: ['user:all'], note: `sourcegraph-${getPlatformName()}` },
-        useAccessToken: false,
-    }).pipe(
-        map(({ data, errors }) => {
-            if (!data || !data.createAccessToken || (errors && errors.length > 0)) {
-                throw createAggregateError(errors)
+            variables: { userID, scopes: ['user:all'], note: `sourcegraph-${getPlatformName()}` },
+            useAccessToken: false,
+        }).pipe(
+            map(({ data, errors }) => {
+                if (!data || !data.createAccessToken || (errors && errors.length > 0)) {
+                    throw createAggregateError(errors)
+                }
+                return data.createAccessToken
+            })
+        )
+)
+
+export const fetchAccessTokenIDs = memoizeObservable(
+    (userID: GQL.ID): Observable<Pick<AccessToken, 'id'>[]> =>
+        queryGraphQL({
+            ctx: getContext({ repoKey: '' }),
+            request: `
+            query AccessTokenIDs {
+                currentUser {
+                    accessTokens {
+                        nodes {
+                          id
+                        }
+                    }
+                }
             }
-            return data.createAccessToken.token
-        })
-    )
+            `,
+            variables: { userID, scopes: ['user:all'], note: `sourcegraph-${getPlatformName()}` },
+            useAccessToken: false,
+        }).pipe(
+            map(({ data, errors }) => {
+                if (
+                    !data ||
+                    !data.currentUser ||
+                    !data.currentUser.accessTokens ||
+                    !data.currentUser.accessTokens.nodes ||
+                    (errors && errors.length > 0)
+                ) {
+                    throw createAggregateError(errors)
+                }
+                return data.currentUser.accessTokens.nodes
+            })
+        )
 )

src/shared/backend/headers.tsx

@@ -1,26 +1,9 @@
 import { Observable, of } from 'rxjs'
-import { filter, map, switchMap } from 'rxjs/operators'
+import { map, switchMap } from 'rxjs/operators'
 
-import { isDefined } from '@sourcegraph/codeintellify/lib/helpers'
-import { getAccessToken, setAccessToken } from '../auth/access_token'
-import { isInPage, isPhabricator } from '../context'
+import { getAccessToken } from '../auth/access_token'
+import { isContent, isInPage, isPhabricator } from '../context'
 import { getExtensionVersionSync, getPlatformName, isSourcegraphDotCom } from '../util/context'
-import { createAccessToken } from './auth'
-import { fetchCurrentUser } from './server'
-
-const withAccessToken = (url: string) =>
-    getAccessToken(url).pipe(
-        switchMap(token => {
-            if (token) {
-                return of(token)
-            }
-
-            return fetchCurrentUser(false).pipe(
-                filter(isDefined),
-                switchMap(user => createAccessToken(user.id).pipe(setAccessToken(url)))
-            )
-        })
-    )
 
 /**
  * getHeaders emits the required headers for making requests to Sourcegraph server instances.
@@ -41,11 +24,11 @@ export function getHeaders(
     }
 
     return of(url).pipe(
-        switchMap(url => (useToken && !isSourcegraphDotCom(url) ? withAccessToken(url) : of(undefined))),
+        switchMap(url => (isContent && useToken && !isSourcegraphDotCom(url) ? getAccessToken(url) : of(undefined))),
         map(accessToken => {
             const headers = new Headers()
             if (accessToken) {
-                headers.append('Authorization', `token ${accessToken}`)
+                headers.append('Authorization', `token ${accessToken.token}`)
             }
 
             return headers

src/shared/components/options/ConnectionCard.tsx

@@ -1,3 +1,4 @@
+import { propertyIsDefined } from '@sourcegraph/codeintellify/lib/helpers'
 import * as React from 'react'
 import {
     Alert,
@@ -14,12 +15,15 @@ import {
     ListGroupItemHeading,
     Row,
 } from 'reactstrap'
-import { Subscription } from 'rxjs'
+import { of, Subscription } from 'rxjs'
+import { filter, map, switchMap } from 'rxjs/operators'
 import * as permissions from '../../../browser/permissions'
 import storage from '../../../browser/storage'
 import { StorageItems } from '../../../browser/types'
 import { GQL } from '../../../types/gqlschema'
-import { fetchSite } from '../../backend/server'
+import { getAccessToken, setAccessToken } from '../../auth/access_token'
+import { createAccessToken, fetchAccessTokenIDs } from '../../backend/auth'
+import { fetchCurrentUser, fetchSite } from '../../backend/server'
 import { DEFAULT_SOURCEGRAPH_URL, isSourcegraphDotCom, setSourcegraphUrl, sourcegraphUrl } from '../../util/context'
 
 interface Props {
@@ -232,8 +236,10 @@ export class ConnectionCard extends React.Component<Props, State> {
     }
 
     private checkConnection = (): void => {
+        const fetchingSite = fetchSite()
+
         this.subscriptions.add(
-            fetchSite().subscribe(
+            fetchingSite.subscribe(
                 site => {
                     this.setState({ site })
                 },
@@ -242,6 +248,34 @@ export class ConnectionCard extends React.Component<Props, State> {
                 }
             )
         )
+
+        this.subscriptions.add(
+            // Ensure the site is valid.
+            fetchingSite
+                .pipe(
+                    // Get the access token for this server if we have it.
+                    switchMap(() => getAccessToken(sourcegraphUrl)),
+                    switchMap(token => fetchCurrentUser(false).pipe(map(user => ({ user, token })))),
+                    filter(propertyIsDefined('user')),
+                    // Get the IDs for all access tokens for the user.
+                    switchMap(({ token, user }) =>
+                        fetchAccessTokenIDs(user.id).pipe(map(usersTokenIDs => ({ usersTokenIDs, user, token })))
+                    ),
+                    // Make sure the token still exists on the server. If it
+                    // does exits, use it, otherwise create a new one.
+                    switchMap(({ user, token, usersTokenIDs }) => {
+                        const tokenExists = token && usersTokenIDs.map(({ id }) => id).includes(token.id)
+
+                        return token && tokenExists
+                            ? of(token)
+                            : createAccessToken(user.id).pipe(setAccessToken(sourcegraphUrl))
+                    })
+                )
+                .subscribe(() => {
+                    // We don't need to do anything with the token now. We just
+                    // needed to ensure we had one saved.
+                })
+        )
     }
 
     public render(): JSX.Element | null {