[backend] Improve support-package generation using redis to handle co…

…ncurrency (#5548) Co-authored-by: Julien Richard <julien.richard@filigran.io>
OpenCTI-Platform · May 7, 2024 · bd2fdcd · bd2fdcd
1 parent bd19833
commit bd2fdcd
Show file tree

Hide file tree

Showing 15 changed files with 120 additions and 268 deletions.
diff --git a/...cti-platform/opencti-front/src/private/components/settings/support/SupportPackageLine.tsx b/...cti-platform/opencti-front/src/private/components/settings/support/SupportPackageLine.tsx
@@ -25,6 +25,7 @@ import { deleteNode } from '../../../../utils/store';
 import { hexToRGB } from '../../../../utils/Colors';
 import { DataColumns } from '../../../../components/list_lines';
 import useApiMutation from '../../../../utils/hooks/useApiMutation';
+import { minutesBetweenDates, now } from '../../../../utils/Time';
 
 const styles = {
   bodyItem: {
@@ -55,7 +56,7 @@ const styles = {
   },
 };
 
-type PackageStatus = 'IN_PROGRESS' | 'READY' | 'IN_ERROR' | '%future added value';
+type PackageStatus = 'IN_PROGRESS' | 'READY' | 'IN_ERROR' | 'TIMEOUT' | '%future added value';
 
 const SupportPackageLineForceZipMutation = graphql`
   mutation SupportPackageLineForceZipMutation(    
@@ -93,6 +94,7 @@ const packageStatusColors: { [key in PackageStatus]: string } = {
   IN_PROGRESS: '#303f9f',
   READY: '#4caf50',
   IN_ERROR: '#f44336',
+  TIMEOUT: '#f44336',
   '%future added value': '#9e9e9e',
 };
 
@@ -114,6 +116,10 @@ const SupportPackageLine: FunctionComponent<SupportPackageLineProps> = ({
   const [commitDelete] = useApiMutation(SupportPackageLineDeleteMutation);
   const [commitForceZip] = useApiMutation(SupportPackageLineForceZipMutation);
   const isProgress = data.package_status === 'IN_PROGRESS';
+  const isReady = data.package_status === 'READY';
+  const isTooLong = minutesBetweenDates(data.created_at, now()) > 1;
+  const isTimeout = !isReady && minutesBetweenDates(data.created_at, now()) > 5;
+  const finalStatus = isTimeout ? 'TIMEOUT' : data.package_status;
 
   const handleOpenDelete = () => {
     setDisplayDelete(true);
@@ -170,12 +176,17 @@ const SupportPackageLine: FunctionComponent<SupportPackageLineProps> = ({
     <>
       <ListItem divider={true} style={{ ...styles.item }}>
         <ListItemIcon>
-          {isProgress && (
+          {isProgress && !isTimeout && (
             <CircularProgress
               size={20}
               color='inherit'
             />
           )}
+          {isProgress && isTimeout && (
+          <FileOutline
+            color='inherit'
+          />
+          )}
           {!isProgress && (
             <FileOutline
               color='inherit'
@@ -191,13 +202,13 @@ const SupportPackageLine: FunctionComponent<SupportPackageLineProps> = ({
               <div style={{ width: dataColumns.package_status.width, ...styles.bodyItem }}>
                 <Chip
                   style={{
-                    color: packageStatusColors[data.package_status],
-                    borderColor: packageStatusColors[data.package_status],
-                    backgroundColor: hexToRGB(packageStatusColors[data.package_status]),
+                    color: packageStatusColors[finalStatus],
+                    borderColor: packageStatusColors[finalStatus],
+                    backgroundColor: hexToRGB(packageStatusColors[finalStatus]),
                     ...styles.chipInList,
                     ...styles.label,
                   }}
-                  label={data.package_status}
+                  label={t_i18n(finalStatus)}
                 />
               </div>
               <div
@@ -209,29 +220,29 @@ const SupportPackageLine: FunctionComponent<SupportPackageLineProps> = ({
           }
         />
         <ListItemSecondaryAction>
-          <Tooltip title={t_i18n('Force download on this support package')}>
+          {!isReady && <Tooltip title={t_i18n('Force download on this support package')}>
             <span>
-              <IconButton onClick={handleForceZip}>
+              <IconButton disabled={!isTooLong} onClick={handleForceZip}>
                 <DownloadingOutlined fontSize="small" />
               </IconButton>
             </span>
-          </Tooltip>
-          <Tooltip title={t_i18n('Download this support package')}>
+          </Tooltip>}
+          {isReady && <Tooltip title={t_i18n('Download this support package')}>
             <span>
               <IconButton
-                disabled={isProgress || !data.package_url}
+                disabled={!data.package_url}
                 href={`${APP_BASE_PATH}/storage/get/${encodeURIComponent(
                   data.package_url || '',
                 )}`}
               >
                 <GetAppOutlined fontSize="small" />
               </IconButton>
             </span>
-          </Tooltip>
+          </Tooltip>}
           <Tooltip title={t_i18n('Delete this support package')}>
             <span>
               <IconButton
-                disabled={false}
+                disabled={!isReady && !isTooLong}
                 color='inherit'
                 onClick={handleOpenDelete}
                 size="small"

diff --git a/opencti-platform/opencti-front/src/schema/relay.schema.graphql b/opencti-platform/opencti-front/src/schema/relay.schema.graphql
@@ -11514,7 +11514,6 @@ type SupportPackage implements InternalObject & BasicObject {
   package_status: PackageStatus!
   package_url: String
   package_upload_dir: String
-  nodes_status: [SupportNodeStatus]
   nodes_count: Int!
   createdBy: Individual
   creators: [Creator!]
@@ -11525,11 +11524,6 @@ type SupportPackageConnection {
   edges: [SupportPackageEdge!]!
 }
 
-type SupportNodeStatus {
-  node_id: String!
-  package_status: PackageStatus!
-}
-
 type SupportPackageEdge {
   cursor: String!
   node: SupportPackage!

diff --git a/opencti-platform/opencti-graphql/src/database/redis.ts b/opencti-platform/opencti-graphql/src/database/redis.ts
@@ -850,3 +850,38 @@ export const getLastPlaybookExecutions = async (playbookId: string) => {
   });
 };
 // endregion
+
+// region - support package handling
+export const SUPPORT_NODE_STATUS_IN_PROGRESS = 0;
+export const SUPPORT_NODE_STATUS_READY = 10;
+export const SUPPORT_NODE_STATUS_IN_ERROR = 100;
+
+/**
+ * Add or update for a given support package, one node status.
+ * @param supportPackageId
+ * @param nodeId
+ * @param nodeStatus one of SUPPORT_NODE_STATUS_IN_PROGRESS, SUPPORT_NODE_STATUS_READY, SUPPORT_NODE_STATUS_IN_ERROR
+ */
+export const redisStoreSupportPackageNodeStatus = (supportPackageId:string, nodeId: string, nodeStatus: number) => {
+  const setKeyId = `support:${supportPackageId}`;
+  // redis score =  nodeStatus
+  // redis member = nodeId
+  return getClientBase().zadd(setKeyId, nodeStatus, nodeId);
+};
+
+/**
+ * Count for a support package the number of node with a status.
+ * @param supportPackageId
+ * @param nodeStatus
+ */
+export const redisCountSupportPackageNodeWithStatus = (supportPackageId: string, nodeStatus: number) => {
+  const setKeyId = `support:${supportPackageId}`;
+  return getClientBase().zcount(setKeyId, nodeStatus, nodeStatus);
+};
+
+export const redisDeleteSupportPackageNodeStatus = (supportPackageId: string) => {
+  const setKeyId = `support:${supportPackageId}`;
+  return getClientBase().del(setKeyId);
+};
+
+// endregion - support package handling
diff --git a/opencti-platform/opencti-graphql/src/generated/graphql.ts b/opencti-platform/opencti-graphql/src/generated/graphql.ts
@@ -24233,12 +24233,6 @@ export type SubscriptionWorkspaceArgs = {
   id: Scalars['ID']['input'];
 };
 
-export type SupportNodeStatus = {
-  __typename?: 'SupportNodeStatus';
-  node_id: Scalars['String']['output'];
-  package_status: PackageStatus;
-};
-
 export type SupportPackage = BasicObject & InternalObject & {
   __typename?: 'SupportPackage';
   createdBy?: Maybe<Individual>;
@@ -24248,7 +24242,6 @@ export type SupportPackage = BasicObject & InternalObject & {
   id: Scalars['ID']['output'];
   name: Scalars['String']['output'];
   nodes_count: Scalars['Int']['output'];
-  nodes_status?: Maybe<Array<Maybe<SupportNodeStatus>>>;
   package_status: PackageStatus;
   package_upload_dir?: Maybe<Scalars['String']['output']>;
   package_url?: Maybe<Scalars['String']['output']>;
@@ -29213,7 +29206,6 @@ export type ResolversTypes = ResolversObject<{
   SubTypeEditMutations: ResolverTypeWrapper<Omit<SubTypeEditMutations, 'statusAdd' | 'statusDelete' | 'statusFieldPatch'> & { statusAdd?: Maybe<ResolversTypes['SubType']>, statusDelete?: Maybe<ResolversTypes['SubType']>, statusFieldPatch?: Maybe<ResolversTypes['SubType']> }>;
   SubTypesOrdering: SubTypesOrdering;
   Subscription: ResolverTypeWrapper<{}>;
-  SupportNodeStatus: ResolverTypeWrapper<SupportNodeStatus>;
   SupportPackage: ResolverTypeWrapper<BasicStoreEntitySupportPackage>;
   SupportPackageAddInput: SupportPackageAddInput;
   SupportPackageConnection: ResolverTypeWrapper<Omit<SupportPackageConnection, 'edges'> & { edges: Array<ResolversTypes['SupportPackageEdge']> }>;
@@ -29925,7 +29917,6 @@ export type ResolversParentTypes = ResolversObject<{
   SubTypeEdge: Omit<SubTypeEdge, 'node'> & { node: ResolversParentTypes['SubType'] };
   SubTypeEditMutations: Omit<SubTypeEditMutations, 'statusAdd' | 'statusDelete' | 'statusFieldPatch'> & { statusAdd?: Maybe<ResolversParentTypes['SubType']>, statusDelete?: Maybe<ResolversParentTypes['SubType']>, statusFieldPatch?: Maybe<ResolversParentTypes['SubType']> };
   Subscription: {};
-  SupportNodeStatus: SupportNodeStatus;
   SupportPackage: BasicStoreEntitySupportPackage;
   SupportPackageAddInput: SupportPackageAddInput;
   SupportPackageConnection: Omit<SupportPackageConnection, 'edges'> & { edges: Array<ResolversParentTypes['SupportPackageEdge']> };
@@ -37327,12 +37318,6 @@ export type SubscriptionResolvers<ContextType = any, ParentType extends Resolver
   workspace?: SubscriptionResolver<Maybe<ResolversTypes['Workspace']>, "workspace", ParentType, ContextType, RequireFields<SubscriptionWorkspaceArgs, 'id'>>;
 }>;
 
-export type SupportNodeStatusResolvers<ContextType = any, ParentType extends ResolversParentTypes['SupportNodeStatus'] = ResolversParentTypes['SupportNodeStatus']> = ResolversObject<{
-  node_id?: Resolver<ResolversTypes['String'], ParentType, ContextType>;
-  package_status?: Resolver<ResolversTypes['PackageStatus'], ParentType, ContextType>;
-  __isTypeOf?: IsTypeOfResolverFn<ParentType, ContextType>;
-}>;
-
 export type SupportPackageResolvers<ContextType = any, ParentType extends ResolversParentTypes['SupportPackage'] = ResolversParentTypes['SupportPackage']> = ResolversObject<{
   createdBy?: Resolver<Maybe<ResolversTypes['Individual']>, ParentType, ContextType>;
   created_at?: Resolver<ResolversTypes['DateTime'], ParentType, ContextType>;
@@ -37341,7 +37326,6 @@ export type SupportPackageResolvers<ContextType = any, ParentType extends Resolv
   id?: Resolver<ResolversTypes['ID'], ParentType, ContextType>;
   name?: Resolver<ResolversTypes['String'], ParentType, ContextType>;
   nodes_count?: Resolver<ResolversTypes['Int'], ParentType, ContextType>;
-  nodes_status?: Resolver<Maybe<Array<Maybe<ResolversTypes['SupportNodeStatus']>>>, ParentType, ContextType>;
   package_status?: Resolver<ResolversTypes['PackageStatus'], ParentType, ContextType>;
   package_upload_dir?: Resolver<Maybe<ResolversTypes['String']>, ParentType, ContextType>;
   package_url?: Resolver<Maybe<ResolversTypes['String']>, ParentType, ContextType>;
@@ -39139,7 +39123,6 @@ export type Resolvers<ContextType = any> = ResolversObject<{
   SubTypeEdge?: SubTypeEdgeResolvers<ContextType>;
   SubTypeEditMutations?: SubTypeEditMutationsResolvers<ContextType>;
   Subscription?: SubscriptionResolvers<ContextType>;
-  SupportNodeStatus?: SupportNodeStatusResolvers<ContextType>;
   SupportPackage?: SupportPackageResolvers<ContextType>;
   SupportPackageConnection?: SupportPackageConnectionResolvers<ContextType>;
   SupportPackageEdge?: SupportPackageEdgeResolvers<ContextType>;

diff --git a/opencti-platform/opencti-graphql/src/listener/supportPackageListener.ts b/opencti-platform/opencti-graphql/src/listener/supportPackageListener.ts
@@ -20,8 +20,8 @@ export const onSupportPackageMessage = async (event: { instance: BasicStoreEntit
   logApp.info(`[OPENCTI-MODULE] Support Package got event. ${event.instance.id} on node ${NODE_INSTANCE_ID}`);
   try {
     if (event.instance.entity_type === ENTITY_TYPE_SUPPORT_PACKAGE) {
-      await wait(Math.floor(Math.random() * 500));
       await registerNodeInSupportPackage(context, SYSTEM_USER, event.instance.id, PackageStatus.InProgress);
+      await wait(5000); // Wait for all nodes to register in Redis
       await sendCurrentNodeSupportLogToS3(context, SYSTEM_USER, event.instance as StoreEntitySupportPackage);
       await registerNodeInSupportPackage(context, SYSTEM_USER, event.instance.id, PackageStatus.Ready);
     } else {

diff --git a/opencti-platform/opencti-graphql/src/modules/support/support-converter.ts b/opencti-platform/opencti-graphql/src/modules/support/support-converter.ts
@@ -10,7 +10,6 @@ const convertSupportPackageToStix = (instance: StoreEntitySupportPackage): StixS
     package_status: instance.package_status,
     package_url: instance.package_url,
     package_upload_dir: instance.package_upload_dir,
-    nodes_status: instance.nodes_status,
     nodes_count: instance.nodes_count,
     extensions: {
       [STIX_EXT_OCTI]: cleanObject({