Implement a transfer flow to restrict ops

packages/core/data-transfer/src/strapi/providers/remote-destination/index.ts

@@ -1,6 +1,7 @@
 import { WebSocket } from 'ws';
 import { v4 } from 'uuid';
 import { Writable } from 'stream';
+import { once } from 'lodash/fp';
 
 import { createDispatcher } from './utils';
 
@@ -72,6 +73,46 @@ class RemoteStrapiDestinationProvider implements IDestinationProvider {
     });
   }
 
+  #startStepOnce(stage: client.TransferPushStep) {
+    return once(() => this.#startStep(stage));
+  }
+
+  async #startStep<T extends client.TransferPushStep>(step: T) {
+    try {
+      await this.dispatcher?.dispatchTransferStep({ action: 'start', step });
+    } catch (e) {
+      if (e instanceof Error) {
+        return e;
+      }
+
+      if (typeof e === 'string') {
+        return new ProviderTransferError(e);
+      }
+
+      return new ProviderTransferError('Unexpected error');
+    }
+
+    return null;
+  }
+
+  async #endStep<T extends client.TransferPushStep>(step: T) {
+    try {
+      await this.dispatcher?.dispatchTransferStep({ action: 'end', step });
+    } catch (e) {
+      if (e instanceof Error) {
+        return e;
+      }
+
+      if (typeof e === 'string') {
+        return new ProviderTransferError(e);
+      }
+
+      return new ProviderTransferError('Unexpected error');
+    }
+
+    return null;
+  }
+
   async #streamStep<T extends client.TransferPushStep>(
     step: T,
     data: client.GetTransferPushStreamData<T>
@@ -174,48 +215,106 @@ class RemoteStrapiDestinationProvider implements IDestinationProvider {
   }
 
   createEntitiesWriteStream(): Writable {
+    const startEntitiesTransferOnce = this.#startStepOnce('entities');
+
     return new Writable({
       objectMode: true,
-      write: async (entity: IEntity, _encoding, callback) => {
-        const e = await this.#streamStep('entities', entity);
+      final: async (callback) => {
+        const e = await this.#endStep('entities');
 
         callback(e);
       },
+      write: async (entity: IEntity, _encoding, callback) => {
+        const startError = await startEntitiesTransferOnce();
+
+        if (startError) {
+          return callback(startError);
+        }
+
+        const streamError = await this.#streamStep('entities', entity);
+
+        callback(streamError);
+      },
     });
   }
 
   createLinksWriteStream(): Writable {
+    const startLinksTransferOnce = this.#startStepOnce('links');
+
     return new Writable({
       objectMode: true,
-      write: async (link: ILink, _encoding, callback) => {
-        const e = await this.#streamStep('links', link);
+      final: async (callback) => {
+        const e = await this.#endStep('links');
 
         callback(e);
       },
+      write: async (link: ILink, _encoding, callback) => {
+        const startError = await startLinksTransferOnce();
+
+        if (startError) {
+          return callback(startError);
+        }
+
+        const streamError = await this.#streamStep('links', link);
+
+        callback(streamError);
+      },
     });
   }
 
   createConfigurationWriteStream(): Writable {
+    const startConfigurationTransferOnce = this.#startStepOnce('configuration');
+
     return new Writable({
       objectMode: true,
-      write: async (configuration: IConfiguration, _encoding, callback) => {
-        const e = await this.#streamStep('configuration', configuration);
+      final: async (callback) => {
+        const e = await this.#endStep('configuration');
 
         callback(e);
       },
+      write: async (configuration: IConfiguration, _encoding, callback) => {
+        const startError = await startConfigurationTransferOnce();
+
+        if (startError) {
+          return callback(startError);
+        }
+
+        const streamError = await this.#streamStep('configuration', configuration);
+
+        callback(streamError);
+      },
     });
   }
 
   createAssetsWriteStream(): Writable | Promise<Writable> {
+    const startAssetsTransferOnce = this.#startStepOnce('assets');
+
     return new Writable({
       objectMode: true,
       final: async (callback) => {
         // TODO: replace this stream call by an end call
-        const e = await this.#streamStep('assets', null);
+        const endError = await this.#streamStep('assets', null);
 
-        callback(e);
+        if (endError) {
+          return callback(endError);
+        }
+
+        const endStepError = await this.#endStep('assets');
+
+        if (endStepError) {
+          return callback(endStepError);
+        }
+
+        return callback(null);
       },
+
       write: async (asset: IAsset, _encoding, callback) => {
+        const startError = await startAssetsTransferOnce();
+
+        if (startError) {
+          return callback(startError);
+        }
+
         const { filename, filepath, stats, stream } = asset;
         const assetID = v4();
 

packages/core/data-transfer/src/strapi/remote/flows/default.ts

@@ -0,0 +1,40 @@
+import type { Step } from '.';
+
+export default [
+  {
+    kind: 'action',
+    action: 'bootstrap',
+  },
+  {
+    kind: 'action',
+    action: 'init',
+  },
+  {
+    kind: 'action',
+    action: 'beforeTransfer',
+  },
+  {
+    kind: 'transfer',
+    stage: 'schemas',
+  },
+  {
+    kind: 'transfer',
+    stage: 'entities',
+  },
+  {
+    kind: 'transfer',
+    stage: 'assets',
+  },
+  {
+    kind: 'transfer',
+    stage: 'links',
+  },
+  {
+    kind: 'transfer',
+    stage: 'configuration',
+  },
+  {
+    kind: 'action',
+    action: 'close',
+  },
+] as readonly Step[];

packages/core/data-transfer/src/strapi/remote/flows/index.ts

@@ -0,0 +1,78 @@
+import type { TransferStage } from '../../../../types';
+
+export type Step =
+  | { kind: 'action'; action: string }
+  | { kind: 'transfer'; stage: TransferStage; locked?: boolean };
+
+export { default as DEFAULT_TRANSFER_FLOW } from './default';
+
+interface IState {
+  step: Step | null;
+}
+
+export const createFlow = (flow: readonly Step[]) => {
+  const state: IState = { step: null };
+
+  /**
+   * Equality check between two steps
+   */
+  const stepEqual = (stepA: Step, stepB: Step): boolean => {
+    if (stepA.kind === 'action' && stepB.kind === 'action') {
+      return stepA.action === stepB.action;
+    }
+
+    if (stepA.kind === 'transfer' && stepB.kind === 'transfer') {
+      return stepA.stage === stepB.stage;
+    }
+
+    return false;
+  };
+
+  /**
+   * Find the index for a given step
+   */
+  const findStepIndex = (step: Step) => flow.findIndex((flowStep) => stepEqual(step, flowStep));
+
+  return {
+    has(step: Step) {
+      return findStepIndex(step) !== -1;
+    },
+
+    can(step: Step) {
+      if (state.step === null) {
+        return true;
+      }
+
+      const delta = findStepIndex(step) - findStepIndex(state.step);
+
+      // It's possible to send multiple time the same transfer step in a row
+      if (delta === 0 && step.kind === 'transfer') {
+        return true;
+      }
+
+      return delta > 0;
+    },
+
+    cannot(step: Step) {
+      return !this.can(step);
+    },
+
+    set(step: Step) {
+      const canSwitch = this.can(step);
+
+      if (!canSwitch) {
+        throw new Error('Impossible to proceed to the given step');
+      }
+
+      state.step = step;
+
+      return this;
+    },
+
+    get() {
+      return state.step;
+    },
+  };
+};
+
+export type TransferFlow = ReturnType<typeof createFlow>;

packages/core/data-transfer/src/strapi/remote/handlers.ts

@@ -5,16 +5,23 @@ import { randomUUID } from 'crypto';
 import { WebSocket } from 'ws';
 
 import type { IPushController } from './controllers/push';
+import type { TransferFlow, Step } from './flows';
+import type { client, server } from '../../../types/remote/protocol';
 
 import createPushController from './controllers/push';
-import type { client, server } from '../../../types/remote/protocol';
 import { ProviderTransferError, ProviderInitializationError } from '../../errors/providers';
 import { TRANSFER_METHODS } from './constants';
+import { createFlow, DEFAULT_TRANSFER_FLOW } from './flows';
 
 type TransferMethod = typeof TRANSFER_METHODS[number];
 
 interface ITransferState {
-  transfer?: { id: string; kind: client.TransferKind; startedAt: number };
+  transfer?: {
+    id: string;
+    kind: client.TransferKind;
+    startedAt: number;
+    flow: TransferFlow;
+  };
   controller?: IPushController;
 }
 
@@ -44,6 +51,16 @@ export const createTransferHandler = (options: IHandlerOptions) => {
         const state: ITransferState = {};
         let uuid: string | undefined;
 
+        function assertValidTransfer(
+          transferState: ITransferState
+        ): asserts transferState is Required<ITransferState> {
+          const { transfer, controller } = transferState;
+
+          if (!controller || !transfer) {
+            throw new ProviderTransferError('Invalid transfer process');
+          }
+        }
+
         /**
          * Format error & message to follow the remote transfer protocol
          */
@@ -140,7 +157,12 @@ export const createTransferHandler = (options: IHandlerOptions) => {
             });
           }
 
-          state.transfer = { id: randomUUID(), kind: transfer, startedAt: Date.now() };
+          state.transfer = {
+            id: randomUUID(),
+            kind: transfer,
+            startedAt: Date.now(),
+            flow: createFlow(DEFAULT_TRANSFER_FLOW),
+          };
 
           return { transferID: state.transfer.id };
         };
@@ -172,6 +194,7 @@ export const createTransferHandler = (options: IHandlerOptions) => {
           }
 
           if (command === 'end') {
+            assertValidTransfer(state);
             await answer(() => end(msg));
           }
 
@@ -181,15 +204,17 @@ export const createTransferHandler = (options: IHandlerOptions) => {
         };
 
         const onTransferCommand = async (msg: client.TransferMessage) => {
-          const { transferID, kind } = msg;
-          const { controller } = state;
+          assertValidTransfer(state);
 
-          await verifyAuth(state.transfer?.kind);
+          const { transferID, kind } = msg;
+          const { controller, transfer } = state;
+
+          await verifyAuth(transfer.kind);
 
           // TODO: (re)move this check
           // It shouldn't be possible to start a pull transfer for now, so reaching
           // this code should be impossible too, but this has been added by security
-          if (state.transfer?.kind === 'pull') {
+          if (transfer.kind === 'pull') {
             return callback(new ProviderTransferError('Pull transfer not implemented'));
           }
 
@@ -214,6 +239,24 @@ export const createTransferHandler = (options: IHandlerOptions) => {
               );
             }
 
+            const step: Step = { kind: 'action', action };
+            const isStepRegistered = transfer.flow.has(step);
+
+            if (isStepRegistered) {
+              if (transfer.flow.cannot(step)) {
+                return callback(
+                  new ProviderTransferError(
+                    `Invalid action "${action}" found for the current flow `,
+                    {
+                      action,
+                    }
+                  )
+                );
+              }
+
+              transfer.flow.set(step);
+            }
+
             await answer(() => controller.actions[action as keyof typeof controller.actions]());
           }
 
@@ -222,19 +265,78 @@ export const createTransferHandler = (options: IHandlerOptions) => {
             // We can only have push transfer message for the moment
             const message = msg as client.TransferPushMessage;
 
-            // TODO: lock transfer process
+            const currentStep = transfer.flow.get();
+            const step: Step = { kind: 'transfer', stage: message.step };
+
+            // Lock the current transfer stage
             if (message.action === 'start') {
-              // console.log('Starting transfer for ', message.step);
+              if (currentStep?.kind === 'transfer' && currentStep.locked) {
+                return callback(
+                  new ProviderTransferError(
+                    `It's not possible to start a new transfer stage (${message.step}) while another one is in progress (${currentStep.stage})`
+                  )
+                );
+              }
+
+              if (transfer.flow.cannot(step)) {
+                return callback(
+                  new ProviderTransferError(
+                    `Invalid stage (${message.step}) provided for the current flow`,
+                    { step }
+                  )
+                );
+              }
+
+              transfer?.flow.set({ ...step, locked: true });
+
+              return callback(null, { ok: true });
             }
 
-            // Stream step
-            else if (message.action === 'stream') {
+            // Stream operation on the current transfer stage
+            if (message.action === 'stream') {
+              if (currentStep?.kind === 'transfer' && !currentStep.locked) {
+                return callback(
+                  new ProviderTransferError(
+                    `You need to initialize the transfer stage (${message.step}) before starting to stream data`
+                  )
+                );
+              }
+              if (transfer?.flow.cannot(step)) {
+                return callback(
+                  new ProviderTransferError(
+                    `Invalid stage (${message.step}) provided for the current flow`,
+                    { step }
+                  )
+                );
+              }
+
               await answer(() => controller.transfer[message.step]?.(message.data as never));
             }
 
-            // TODO: unlock transfer process
-            else if (message.action === 'end') {
-              // console.log('Ending transfer for ', message.step);
+            // Unlock the current transfer stage
+            if (message.action === 'end') {
+              // Cannot unlock if not locked (aka: started)
+              if (currentStep?.kind === 'transfer' && !currentStep.locked) {
+                return callback(
+                  new ProviderTransferError(
+                    `You need to initialize the transfer stage (${message.step}) before ending it`
+                  )
+                );
+              }
+
+              // Cannot unlock if invalid step provided
+              if (transfer?.flow.cannot(step)) {
+                return callback(
+                  new ProviderTransferError(
+                    `Invalid stage (${message.step}) provided for the current flow`,
+                    { step }
+                  )
+                );
+              }
+
+              transfer?.flow.set({ ...step, locked: false });
+
+              return callback(null, { ok: true });
             }
           }
         };