[8.14] [Security Solution] [Attack discovery] Fixes attack discovery …

…cache issues (#183005) (#183009) # Backport This will backport the following commits from `main` to `8.14`: - [[Security Solution] [Attack discovery] Fixes attack discovery cache issues (#183005)](#183005)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Andrew Macri <[email protected]>
elastic · May 9, 2024 · b032740 · b032740
1 parent a14ee25
commit b032740
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 25 deletions.
diff --git a/x-pack/plugins/security_solution/public/attack_discovery/pages/session_storage/index.ts b/x-pack/plugins/security_solution/public/attack_discovery/pages/session_storage/index.ts
@@ -6,7 +6,9 @@
  */
 
 import type { Replacements } from '@kbn/elastic-assistant-common';
-import type { AttackDiscovery } from '../../types';
+import { isEmpty } from 'lodash/fp';
+
+import type { AttackDiscovery, GenerationInterval } from '../../types';
 
 export interface CachedAttackDiscoveries {
   connectorId: string;
@@ -19,7 +21,7 @@ export const encodeCachedAttackDiscoveries = (
   cachedAttackDiscoveries: Record<string, CachedAttackDiscoveries>
 ): string | null => {
   try {
-    return JSON.stringify(cachedAttackDiscoveries, null, 2);
+    return JSON.stringify(cachedAttackDiscoveries);
   } catch {
     return null;
   }
@@ -34,3 +36,75 @@ export const decodeCachedAttackDiscoveries = (
     return null;
   }
 };
+
+export const getSessionStorageCachedAttackDiscoveries = (
+  key: string
+): Record<string, CachedAttackDiscoveries> | null => {
+  if (!isEmpty(key)) {
+    return decodeCachedAttackDiscoveries(sessionStorage.getItem(key) ?? '');
+  }
+
+  return null;
+};
+
+export const setSessionStorageCachedAttackDiscoveries = ({
+  key,
+  cachedAttackDiscoveries,
+}: {
+  key: string;
+  cachedAttackDiscoveries: Record<string, CachedAttackDiscoveries>;
+}) => {
+  if (!isEmpty(key)) {
+    const encoded = encodeCachedAttackDiscoveries(cachedAttackDiscoveries);
+
+    if (encoded != null) {
+      sessionStorage.setItem(key, encoded);
+    }
+  }
+};
+
+export const encodeGenerationIntervals = (
+  generationIntervals: Record<string, GenerationInterval[]>
+): string | null => {
+  try {
+    return JSON.stringify(generationIntervals);
+  } catch {
+    return null;
+  }
+};
+
+export const decodeGenerationIntervals = (
+  generationIntervals: string
+): Record<string, GenerationInterval[]> | null => {
+  try {
+    return JSON.parse(generationIntervals);
+  } catch {
+    return null;
+  }
+};
+
+export const getLocalStorageGenerationIntervals = (
+  key: string
+): Record<string, GenerationInterval[]> | null => {
+  if (!isEmpty(key)) {
+    return decodeGenerationIntervals(sessionStorage.getItem(key) ?? '');
+  }
+
+  return null;
+};
+
+export const setLocalStorageGenerationIntervals = ({
+  key,
+  generationIntervals,
+}: {
+  key: string;
+  generationIntervals: Record<string, GenerationInterval[]>;
+}) => {
+  if (!isEmpty(key)) {
+    const encoded = encodeGenerationIntervals(generationIntervals);
+
+    if (encoded != null) {
+      localStorage.setItem(key, encoded);
+    }
+  }
+};
diff --git a/x-pack/plugins/security_solution/public/attack_discovery/use_attack_discovery/index.tsx b/x-pack/plugins/security_solution/public/attack_discovery/use_attack_discovery/index.tsx
@@ -18,8 +18,7 @@ import {
 } from '@kbn/elastic-assistant-common';
 import { uniq } from 'lodash/fp';
 import moment from 'moment';
-import React, { useCallback, useMemo, useState } from 'react';
-import { useLocalStorage, useSessionStorage } from 'react-use';
+import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import * as uuid from 'uuid';
 import { useFetchAnonymizationFields } from '@kbn/elastic-assistant/impl/assistant/api/anonymization_fields/use_fetch_anonymization_fields';
 
@@ -36,8 +35,10 @@ import {
 import { getAverageIntervalSeconds } from '../pages/loading_callout/countdown/last_times_popover/helpers';
 import type { CachedAttackDiscoveries } from '../pages/session_storage';
 import {
-  encodeCachedAttackDiscoveries,
-  decodeCachedAttackDiscoveries,
+  getLocalStorageGenerationIntervals,
+  getSessionStorageCachedAttackDiscoveries,
+  setLocalStorageGenerationIntervals,
+  setSessionStorageCachedAttackDiscoveries,
 } from '../pages/session_storage';
 import { ERROR_GENERATING_ATTACK_DISCOVERIES } from '../pages/translations';
 import type { AttackDiscovery, GenerationInterval } from '../types';
@@ -67,7 +68,7 @@ export const useAttackDiscovery = ({
   setLoadingConnectorId?: (loadingConnectorId: string | null) => void;
 }): UseAttackDiscovery => {
   const { reportAttackDiscoveriesGenerated } = useAttackDiscoveryTelemetry();
-  const spaceId = useSpaceId() ?? 'default';
+  const spaceId: string | undefined = useSpaceId();
 
   // get Kibana services and connectors
   const {
@@ -86,24 +87,60 @@ export const useAttackDiscovery = ({
 
   const { data: anonymizationFields } = useFetchAnonymizationFields();
 
-  // get cached attack discoveries from session storage:
-  const [sessionStorageCachedAttackDiscoveries, setSessionStorageCachedAttackDiscoveries] =
-    useSessionStorage<string>(
-      `${DEFAULT_ASSISTANT_NAMESPACE}.${ATTACK_DISCOVERY_STORAGE_KEY}.${spaceId}.${CACHED_ATTACK_DISCOVERIES_SESSION_STORAGE_KEY}`
-    );
+  const sessionStorageKey = useMemo(
+    () =>
+      spaceId != null // spaceId is undefined while the useSpaceId hook is loading
+        ? `${DEFAULT_ASSISTANT_NAMESPACE}.${ATTACK_DISCOVERY_STORAGE_KEY}.${spaceId}.${CACHED_ATTACK_DISCOVERIES_SESSION_STORAGE_KEY}`
+        : '',
+    [spaceId]
+  );
+
   const [cachedAttackDiscoveries, setCachedAttackDiscoveries] = useState<
     Record<string, CachedAttackDiscoveries>
-  >(decodeCachedAttackDiscoveries(sessionStorageCachedAttackDiscoveries) ?? {});
+  >({});
+
+  useEffect(() => {
+    const decoded = getSessionStorageCachedAttackDiscoveries(sessionStorageKey);
+
+    if (decoded != null) {
+      setCachedAttackDiscoveries(decoded);
+
+      const decodedAttackDiscoveries = decoded[connectorId ?? '']?.attackDiscoveries;
+      if (decodedAttackDiscoveries != null) {
+        setAttackDiscoveries(decodedAttackDiscoveries);
+      }
+
+      const decodedReplacements = decoded[connectorId ?? '']?.replacements;
+      if (decodedReplacements != null) {
+        setReplacements(decodedReplacements);
+      }
 
-  // get generation intervals from local storage:
-  const [localStorageGenerationIntervals, setLocalStorageGenerationIntervals] = useLocalStorage<
-    Record<string, GenerationInterval[]>
-  >(
-    `${DEFAULT_ASSISTANT_NAMESPACE}.${ATTACK_DISCOVERY_STORAGE_KEY}.${spaceId}.${GENERATION_INTERVALS_LOCAL_STORAGE_KEY}`
+      const decodedLastUpdated = decoded[connectorId ?? '']?.updated;
+      if (decodedLastUpdated != null) {
+        setLastUpdated(decodedLastUpdated);
+      }
+    }
+  }, [connectorId, sessionStorageKey]);
+
+  const localStorageKey = useMemo(
+    () =>
+      spaceId != null // spaceId is undefined while the useSpaceId hook is loading
+        ? `${DEFAULT_ASSISTANT_NAMESPACE}.${ATTACK_DISCOVERY_STORAGE_KEY}.${spaceId}.${GENERATION_INTERVALS_LOCAL_STORAGE_KEY}`
+        : '',
+    [spaceId]
   );
+
   const [generationIntervals, setGenerationIntervals] = React.useState<
     Record<string, GenerationInterval[]> | undefined
-  >(localStorageGenerationIntervals);
+  >(undefined);
+
+  useEffect(() => {
+    const decoded = getLocalStorageGenerationIntervals(localStorageKey);
+
+    if (decoded != null) {
+      setGenerationIntervals(decoded);
+    }
+  }, [localStorageKey]);
 
   // get connector intervals from generation intervals:
   const connectorIntervals = useMemo(
@@ -199,9 +236,10 @@ export const useAttackDiscovery = ({
       };
 
       setCachedAttackDiscoveries(newCachedAttackDiscoveries);
-      setSessionStorageCachedAttackDiscoveries(
-        encodeCachedAttackDiscoveries(newCachedAttackDiscoveries) ?? ''
-      );
+      setSessionStorageCachedAttackDiscoveries({
+        key: sessionStorageKey,
+        cachedAttackDiscoveries: newCachedAttackDiscoveries,
+      });
 
       // update the generation intervals with the latest timing:
       const previousConnectorIntervals: GenerationInterval[] =
@@ -227,7 +265,10 @@ export const useAttackDiscovery = ({
       // only update the generation intervals if alerts were sent as context to the LLM:
       if (newAlertsContextCount != null && newAlertsContextCount > 0) {
         setGenerationIntervals(newGenerationIntervals);
-        setLocalStorageGenerationIntervals(newGenerationIntervals);
+        setLocalStorageGenerationIntervals({
+          key: localStorageKey,
+          generationIntervals: newGenerationIntervals,
+        });
       }
 
       setReplacements(newReplacements);
@@ -266,12 +307,12 @@ export const useAttackDiscovery = ({
     generationIntervals,
     http,
     knowledgeBase,
+    localStorageKey,
     replacements,
     reportAttackDiscoveriesGenerated,
+    sessionStorageKey,
     setConnectorId,
     setLoadingConnectorId,
-    setLocalStorageGenerationIntervals,
-    setSessionStorageCachedAttackDiscoveries,
     toasts,
     traceOptions,
   ]);