From 91bcbc3d6f73003780fc3e786521ce907694fb84 Mon Sep 17 00:00:00 2001
From: Pujit Mehrotra <pujitm2009@gmail.com>
Date: Mon, 21 Oct 2024 12:16:45 -0400
Subject: [PATCH] fix(api): strip server id prefixes from graphql request
 variables

---
 api/src/unraid-api/graph/id-prefix-plugin.ts | 74 ++++++++++++--------
 api/src/utils.ts                             | 31 ++++++++
 2 files changed, 76 insertions(+), 29 deletions(-)

diff --git a/api/src/unraid-api/graph/id-prefix-plugin.ts b/api/src/unraid-api/graph/id-prefix-plugin.ts
index 59cd3fb57..382f4d02c 100644
--- a/api/src/unraid-api/graph/id-prefix-plugin.ts
+++ b/api/src/unraid-api/graph/id-prefix-plugin.ts
@@ -1,31 +1,44 @@
-import { type ApolloServerPlugin } from "@apollo/server";
-import { getServerIdentifier } from "@app/core/utils/server-identifier";
+import { type ApolloServerPlugin } from '@apollo/server';
+import { getServerIdentifier } from '@app/core/utils/server-identifier';
+import { updateObject } from '@app/utils';
+
+type ObjectModifier = (obj: object) => void;
 
 /**
- * Modify all ID fields in the GQL response object to include a prefix
- * @param obj GQL response object, to be modified in place
+ * Returns a function that takes an object and updates any 'id' properties to
+ * include the given serverId as a prefix.
+ *
+ * e.g. If the object is { id: '1234' }, the returned function will update it to
+ * { id: '<serverId>:1234' }.
+ *
+ * @param serverId - The server identifier to use as the prefix.
+ * @returns A function that takes an object and updates any 'id' properties with the given serverId.
  */
-const updateId = (obj: Record<string, unknown>) => {
-    const serverId = getServerIdentifier();
-    const stack = [obj];
-    let iterations = 0;
-    // Prevent infinite loops
-    while (stack.length > 0 && iterations < 100) {
-        const current = stack.pop();
-
-        if (current && typeof current === 'object') {
-            if ('id' in current && typeof current.id === 'string') {
-                current.id = `${serverId}:${current.id}`;
-            }
-
-            for (const value of Object.values(current)) {
-                if (value && typeof value === 'object') {
-                    stack.push(value as Record<string, unknown>);
-                }
-            }
+function prefixWithServerId(serverId: string): ObjectModifier {
+    return (currentObj) => {
+        if ('id' in currentObj && typeof currentObj.id === 'string') {
+            currentObj.id = `${serverId}:${currentObj.id}`;
         }
+    };
+}
 
-        iterations++;
+/**
+ * Takes an object and removes any server prefix from the 'id' property.
+ *
+ * e.g. If the object is { id: '<serverId>:1234' }, the returned function will update it to
+ * { id: '1234' }.
+ *
+ * @param current - The object to update. If it has an 'id' property that is a string and
+ *                  has a server prefix, the prefix is removed.
+ */
+const stripServerPrefixFromIds: ObjectModifier = (current) => {
+    if ('id' in current && typeof current.id === 'string') {
+        const parts = current.id.split(':');
+        // if there are more or less than 2 parts to the split,
+        // assume there is no server prefix and don't touch it.
+        if (parts.length === 2) {
+            current.id = parts[1];
+        }
     }
 };
 
@@ -37,13 +50,16 @@ export const idPrefixPlugin: ApolloServerPlugin = {
         }
         // If ID is requested, return an ID field with an extra prefix
         return {
+            async didResolveOperation({ request }) {
+                if (request.variables) {
+                    updateObject(request.variables, stripServerPrefixFromIds)
+                    stripServerPrefixFromIds(request.variables);
+                }
+            },
             async willSendResponse({ response }) {
-                if (
-                    response.body.kind === 'single' &&
-                    response.body.singleResult.data
-                ) {
-                    // Iteratively update all ID fields with a prefix
-                    updateId(response.body.singleResult.data);
+                if (response.body.kind === 'single' && response.body.singleResult.data) {
+                    const serverId = getServerIdentifier();
+                    updateObject(response.body.singleResult.data, prefixWithServerId(serverId));
                 }
             },
         };
diff --git a/api/src/utils.ts b/api/src/utils.ts
index 85b256018..7208c2a1d 100644
--- a/api/src/utils.ts
+++ b/api/src/utils.ts
@@ -61,3 +61,34 @@ export async function batchProcess<Input, T>(items: Input[], action: (id: Input)
         errorOccured: errors.length > 0,
     };
 }
+
+/**
+ * Traverses an object and its nested objects, passing each one to a callback function.
+ *
+ * This function iterates over the input object, using a stack to keep track of nested objects,
+ * and applies the given modifier function to each object it encounters.
+ * It prevents infinite loops by limiting the number of iterations.
+ *
+ * @param obj - The object to be traversed and modified.
+ * @param modifier - A callback function, taking an object. Modifications should happen in place.
+ */
+export function updateObject(obj: object, modifier: (currentObj: object) => void) {
+    const stack = [obj];
+    let iterations = 0;
+    // Prevent infinite loops
+    while (stack.length > 0 && iterations < 100) {
+        const current = stack.pop();
+
+        if (current && typeof current === 'object') {
+            modifier(current);
+
+            for (const value of Object.values(current)) {
+                if (value && typeof value === 'object') {
+                    stack.push(value);
+                }
+            }
+        }
+
+        iterations++;
+    }
+}