@uppy/companion: implement refresh for authentication tokens

packages/@uppy/aws-s3-multipart/src/index.js

@@ -1,4 +1,4 @@
-import BasePlugin from '@uppy/core/lib/BasePlugin.js'
+import UploaderPlugin from '@uppy/core/lib/UploaderPlugin.js'
 import { Socket, Provider, RequestClient } from '@uppy/companion-client'
 import EventTracker from '@uppy/utils/lib/EventTracker'
 import emitSocketProgress from '@uppy/utils/lib/emitSocketProgress'
@@ -313,10 +313,10 @@ class HTTPCommunicationQueue {
   }
 }
 
-export default class AwsS3Multipart extends BasePlugin {
+export default class AwsS3Multipart extends UploaderPlugin {
   static VERSION = packageJson.version
 
-  #queueRequestSocketToken
+  queueRequestSocketToken
 
   #companionCommunicationQueue
 
@@ -369,7 +369,7 @@ export default class AwsS3Multipart extends BasePlugin {
     this.uploaderEvents = Object.create(null)
     this.uploaderSockets = Object.create(null)
 
-    this.#queueRequestSocketToken = this.requests.wrapPromiseFunction(this.#requestSocketToken, { priority: -1 })
+    this.queueRequestSocketToken = this.requests.wrapPromiseFunction(this.#requestSocketToken, { priority: -1 })
   }
 
   [Symbol.for('uppy test: getClient')] () { return this.#client }
@@ -683,28 +683,6 @@ export default class AwsS3Multipart extends BasePlugin {
     return res.token
   }
 
-  // NOTE! Keep this duplicated code in sync with other plugins
-  // TODO we should probably abstract this into a common function
-  async #uploadRemote (file) {
-    this.resetUploaderReferences(file.id)
-
-    try {
-      if (file.serverToken) {
-        return await this.connectToServerSocket(file)
-      }
-      const serverToken = await this.#queueRequestSocketToken(file)
-
-      if (!this.uppy.getState().files[file.id]) return undefined
-
-      this.uppy.setFileState(file.id, { serverToken })
-      return await this.connectToServerSocket(this.uppy.getFile(file.id))
-    } catch (err) {
-      this.uppy.setFileState(file.id, { serverToken: undefined })
-      this.uppy.emit('upload-error', file, err)
-      throw err
-    }
-  }
-
   async connectToServerSocket (file) {
     return new Promise((resolve, reject) => {
       let queuedRequest
@@ -827,7 +805,8 @@ export default class AwsS3Multipart extends BasePlugin {
 
     const promises = filesFiltered.map((file) => {
       if (file.isRemote) {
-        return this.#uploadRemote(file)
+        this.resetUploaderReferences(file.id)
+        return this.uploadRemoteFile(file)
       }
       return this.#uploadFile(file)
     })

packages/@uppy/aws-s3/src/MiniXHRUpload.js

@@ -1,5 +1,5 @@
 import { nanoid } from 'nanoid/non-secure'
-import { Provider, RequestClient, Socket } from '@uppy/companion-client'
+import { Socket } from '@uppy/companion-client'
 import emitSocketProgress from '@uppy/utils/lib/emitSocketProgress'
 import getSocketHost from '@uppy/utils/lib/getSocketHost'
 import EventTracker from '@uppy/utils/lib/EventTracker'
@@ -53,8 +53,6 @@ function createFormDataUpload (file, opts) {
 const createBareUpload = file => file.data
 
 export default class MiniXHRUpload {
-  #queueRequestSocketToken
-
   constructor (uppy, opts) {
     this.uppy = uppy
     this.opts = {
@@ -67,11 +65,9 @@ export default class MiniXHRUpload {
     this.requests = opts[internalRateLimitedQueue]
     this.uploaderEvents = Object.create(null)
     this.i18n = opts.i18n
-
-    this.#queueRequestSocketToken = this.requests.wrapPromiseFunction(this.#requestSocketToken, { priority: -1 })
   }
 
-  #getOptions (file) {
+  getOptions (file) {
     const { uppy } = this
 
     const overrides = uppy.getState().xhrUpload
@@ -89,16 +85,6 @@ export default class MiniXHRUpload {
     return opts
   }
 
-  uploadFile (id, current, total) {
-    const file = this.uppy.getFile(id)
-    if (file.error) {
-      throw new Error(file.error)
-    } else if (file.isRemote) {
-      return this.#uploadRemoteFile(file, current, total)
-    }
-    return this.#uploadLocalFile(file, current, total)
-  }
-
   #addEventHandlerForFile (eventName, fileID, eventHandler) {
     this.uploaderEvents[fileID].on(eventName, (fileOrID) => {
       // TODO (major): refactor Uppy events to consistently send file objects (or consistently IDs)
@@ -115,8 +101,8 @@ export default class MiniXHRUpload {
     })
   }
 
-  #uploadLocalFile (file, current, total) {
-    const opts = this.#getOptions(file)
+  uploadLocalFile (file, current, total) {
+    const opts = this.getOptions(file)
 
     this.uppy.log(`uploading ${current} of ${total}`)
     return new Promise((resolve, reject) => {
@@ -251,58 +237,9 @@ export default class MiniXHRUpload {
     })
   }
 
-  #requestSocketToken = async (file) => {
-    const opts = this.#getOptions(file)
-    const Client = file.remote.providerOptions.provider ? Provider : RequestClient
-    const client = new Client(this.uppy, file.remote.providerOptions)
-    const allowedMetaFields = Array.isArray(opts.allowedMetaFields)
-      ? opts.allowedMetaFields
-      // Send along all fields by default.
-      : Object.keys(file.meta)
-
-    if (file.tus) {
-      // Install file-specific upload overrides.
-      Object.assign(opts, file.tus)
-    }
-
-    const res = await client.post(file.remote.url, {
-      ...file.remote.body,
-      protocol: 'multipart',
-      endpoint: opts.endpoint,
-      size: file.data.size,
-      fieldname: opts.fieldName,
-      metadata: Object.fromEntries(allowedMetaFields.map(name => [name, file.meta[name]])),
-      httpMethod: opts.method,
-      useFormData: opts.formData,
-      headers: opts.headers,
-    })
-    return res.token
-  }
-
-  // NOTE! Keep this duplicated code in sync with other plugins
-  // TODO we should probably abstract this into a common function
-  async #uploadRemoteFile (file) {
-    // TODO: we could rewrite this to use server-sent events instead of creating WebSockets.
-    try {
-      if (file.serverToken) {
-        return await this.connectToServerSocket(file)
-      }
-      const serverToken = await this.#queueRequestSocketToken(file)
-
-      if (!this.uppy.getState().files[file.id]) return undefined
-
-      this.uppy.setFileState(file.id, { serverToken })
-      return await this.connectToServerSocket(this.uppy.getFile(file.id))
-    } catch (err) {
-      this.uppy.setFileState(file.id, { serverToken: undefined })
-      this.uppy.emit('upload-error', file, err)
-      throw err
-    }
-  }
-
   async connectToServerSocket (file) {
     return new Promise((resolve, reject) => {
-      const opts = this.#getOptions(file)
+      const opts = this.getOptions(file)
       const token = file.serverToken
       const host = getSocketHost(file.remote.companionUrl)
       let socket

packages/@uppy/aws-s3/src/index.js

@@ -25,9 +25,9 @@
  * the XHRUpload code, but at least it's not horrifically broken :)
  */
 
-import BasePlugin from '@uppy/core/lib/BasePlugin.js'
+import UploaderPlugin from '@uppy/core/lib/UploaderPlugin.js'
 import { RateLimitedQueue, internalRateLimitedQueue } from '@uppy/utils/lib/RateLimitedQueue'
-import { RequestClient } from '@uppy/companion-client'
+import { RequestClient, Provider } from '@uppy/companion-client'
 import { filterNonFailedFiles, filterFilesToEmitUploadStarted } from '@uppy/utils/lib/fileFilters'
 
 import packageJson from '../package.json'
@@ -102,9 +102,11 @@ function defaultGetResponseError (content, xhr) {
 let warnedSuccessActionStatus = false
 
 // TODO deprecate this, will use s3-multipart instead
-export default class AwsS3 extends BasePlugin {
+export default class AwsS3 extends UploaderPlugin {
   static VERSION = packageJson.version
 
+  queueRequestSocketToken
+
   #client
 
   #requests
@@ -138,6 +140,8 @@ export default class AwsS3 extends BasePlugin {
 
     this.#client = new RequestClient(uppy, opts)
     this.#requests = new RateLimitedQueue(this.opts.limit)
+
+    this.queueRequestSocketToken = this.#requests.wrapPromiseFunction(this.#requestSocketToken, { priority: -1 })
   }
 
   [Symbol.for('uppy test: getClient')] () { return this.#client }
@@ -241,6 +245,48 @@ export default class AwsS3 extends BasePlugin {
     return Promise.resolve()
   }
 
+  connectToServerSocket (file) {
+    return this.#uploader.connectToServerSocket(file)
+  }
+
+  #requestSocketToken = async (file) => {
+    const opts = this.#uploader.getOptions(file)
+    const Client = file.remote.providerOptions.provider ? Provider : RequestClient
+    const client = new Client(this.uppy, file.remote.providerOptions)
+    const allowedMetaFields = Array.isArray(opts.allowedMetaFields)
+      ? opts.allowedMetaFields
+      // Send along all fields by default.
+      : Object.keys(file.meta)
+
+    if (file.tus) {
+      // Install file-specific upload overrides.
+      Object.assign(opts, file.tus)
+    }
+
+    const res = await client.post(file.remote.url, {
+      ...file.remote.body,
+      protocol: 'multipart',
+      endpoint: opts.endpoint,
+      size: file.data.size,
+      fieldname: opts.fieldName,
+      metadata: Object.fromEntries(allowedMetaFields.map(name => [name, file.meta[name]])),
+      httpMethod: opts.method,
+      useFormData: opts.formData,
+      headers: opts.headers,
+    })
+    return res.token
+  }
+
+  uploadFile (id, current, total) {
+    const file = this.uppy.getFile(id)
+    if (file.error) {
+      throw new Error(file.error)
+    } else if (file.isRemote) {
+      return this.uploadRemoteFile(file, current, total)
+    }
+    return this.#uploader.uploadLocalFile(file, current, total)
+  }
+
   install () {
     const { uppy } = this
     uppy.addPreProcessor(this.#setCompanionHeaders)

packages/@uppy/companion-client/src/Provider.js

@@ -8,6 +8,8 @@ const getName = (id) => {
 }
 
 export default class Provider extends RequestClient {
+  #refreshingTokenPromise
+
   constructor (uppy, opts) {
     super(uppy, opts)
     this.provider = opts.provider
@@ -20,7 +22,7 @@ export default class Provider extends RequestClient {
   }
 
   async headers () {
-    const [headers, token] = await Promise.all([super.headers(), this.getAuthToken()])
+    const [headers, token] = await Promise.all([super.headers(), this.#getAuthToken()])
     const authHeaders = {}
     if (token) {
       authHeaders['uppy-auth-token'] = token
@@ -43,14 +45,18 @@ export default class Provider extends RequestClient {
     return response
   }
 
-  setAuthToken (token) {
+  async setAuthToken (token) {
     return this.uppy.getPlugin(this.pluginId).storage.setItem(this.tokenKey, token)
   }
 
-  getAuthToken () {
+  async #getAuthToken () {
     return this.uppy.getPlugin(this.pluginId).storage.getItem(this.tokenKey)
   }
 
+  async #removeAuthToken () {
+    return this.uppy.getPlugin(this.pluginId).storage.removeItem(this.tokenKey)
+  }
+
   /**
    * Ensure we have a preauth token if necessary. Attempts to fetch one if we don't,
    * or rejects if loading one fails.
@@ -74,10 +80,43 @@ export default class Provider extends RequestClient {
     return `${this.hostname}/${this.id}/connect?${params}`
   }
 
+  refreshTokenUrl () {
+    return `${this.hostname}/${this.id}/refresh-token`
+  }
+
   fileUrl (id) {
     return `${this.hostname}/${this.id}/get/${id}`
   }
 
+  async request (...args) {
+    await this.#refreshingTokenPromise
+
+    try {
+      // throw Object.assign(new Error(), { isAuthError: true }) // testing simulate access token expired (to refresh token)
+      return await super.request(...args)
+    } catch (err) {
+      if (!err.isAuthError) throw err // only handle auth errors (401 from provider)
+
+      await this.#refreshingTokenPromise
+
+      // Many provider requests may be starting at once, however refresh token should only be called once.
+      // Once a refresh token operation has started, we need all other request to wait for this operation (atomically)
+      this.#refreshingTokenPromise = (async () => {
+        try {
+          const response = await super.request({ path: this.refreshTokenUrl(), method: 'POST' })
+          await this.setAuthToken(response.uppyAuthToken)
+        } finally {
+          this.#refreshingTokenPromise = undefined
+        }
+      })()
+
+      await this.#refreshingTokenPromise
+
+      // now retry the request with our new refresh token
+      return super.request(...args)
+    }
+  }
+
   async fetchPreAuthToken () {
     if (!this.companionKeysParams) {
       return
@@ -95,12 +134,10 @@ export default class Provider extends RequestClient {
     return this.get(`${this.id}/list/${directory || ''}`)
   }
 
-  logout () {
-    return this.get(`${this.id}/logout`)
-      .then((response) => Promise.all([
-        response,
-        this.uppy.getPlugin(this.pluginId).storage.removeItem(this.tokenKey),
-      ])).then(([response]) => response)
+  async logout () {
+    const response = await this.get(`${this.id}/logout`)
+    await this.#removeAuthToken()
+    return response
   }
 
   static initPlugin (plugin, opts, defaultOpts) {

packages/@uppy/companion-client/src/RequestClient.js

@@ -150,7 +150,7 @@ export default class RequestClient {
     }))
   }
 
-  async #request ({ path, method = 'GET', data, skipPostResponse, signal }) {
+  async request ({ path, method = 'GET', data, skipPostResponse, signal }) {
     try {
       const headers = await this.preflightAndHeaders(path)
       const response = await fetchWithNetworkError(this.#getUrl(path), {
@@ -172,20 +172,20 @@ export default class RequestClient {
     // TODO: remove boolean support for options that was added for backward compatibility.
     // eslint-disable-next-line no-param-reassign
     if (typeof options === 'boolean') options = { skipPostResponse: options }
-    return this.#request({ ...options, path })
+    return this.request({ ...options, path })
   }
 
   async post (path, data, options = undefined) {
     // TODO: remove boolean support for options that was added for backward compatibility.
     // eslint-disable-next-line no-param-reassign
     if (typeof options === 'boolean') options = { skipPostResponse: options }
-    return this.#request({ ...options, path, method: 'POST', data })
+    return this.request({ ...options, path, method: 'POST', data })
   }
 
   async delete (path, data = undefined, options) {
     // TODO: remove boolean support for options that was added for backward compatibility.
     // eslint-disable-next-line no-param-reassign
     if (typeof options === 'boolean') options = { skipPostResponse: options }
-    return this.#request({ ...options, path, method: 'DELETE', data })
+    return this.request({ ...options, path, method: 'DELETE', data })
   }
 }