From 2b87fc8f16b7f408bc3b702c039f0d64f89694e0 Mon Sep 17 00:00:00 2001
From: Samuel Pitko <samuel.pitko@gmail.com>
Date: Sun, 17 Jul 2022 21:49:24 +0300
Subject: [PATCH] Add validation to blockHeight arg

---
 packages/query/src/graphql/graphql.module.ts  |  2 +
 .../plugins/BlockHeightArgValidationPlugin.ts | 66 +++++++++++++++++++
 2 files changed, 68 insertions(+)
 create mode 100644 packages/query/src/graphql/plugins/BlockHeightArgValidationPlugin.ts

diff --git a/packages/query/src/graphql/graphql.module.ts b/packages/query/src/graphql/graphql.module.ts
index b1939a6aed..d0e2c55360 100644
--- a/packages/query/src/graphql/graphql.module.ts
+++ b/packages/query/src/graphql/graphql.module.ts
@@ -21,6 +21,7 @@ import {Config} from '../configure';
 import {PinoConfig} from '../utils/logger';
 import {getYargsOption} from '../yargs';
 import {plugins} from './plugins';
+import {BlockHeightArgValidationPlugin} from './plugins/BlockHeightArgValidationPlugin';
 import {PgSubscriptionPlugin} from './plugins/PgSubscriptionPlugin';
 import {queryComplexityPlugin} from './plugins/QueryComplexityPlugin';
 import {ProjectService} from './project.service';
@@ -108,6 +109,7 @@ export class GraphqlModule implements OnModuleInit, OnModuleDestroy {
         ? ApolloServerPluginLandingPageGraphQLPlayground()
         : ApolloServerPluginLandingPageDisabled(),
       queryComplexityPlugin({schema, maxComplexity: argv['query-complexity']}),
+      BlockHeightArgValidationPlugin({dbSchema}),
     ];
 
     const server = new ApolloServer({
diff --git a/packages/query/src/graphql/plugins/BlockHeightArgValidationPlugin.ts b/packages/query/src/graphql/plugins/BlockHeightArgValidationPlugin.ts
new file mode 100644
index 0000000000..e48d753f06
--- /dev/null
+++ b/packages/query/src/graphql/plugins/BlockHeightArgValidationPlugin.ts
@@ -0,0 +1,66 @@
+// Copyright 2022 OnFinality Limited authors & contributors
+// SPDX-License-Identifier: Apache-2.0
+
+import {UserInputError} from 'apollo-server-express';
+import type {ApolloServerPlugin} from 'apollo-server-plugin-base';
+import {DocumentNode, GraphQLError, visit} from 'graphql';
+import {Pool} from 'pg';
+
+function parseBlockHeightArgs(document: DocumentNode): string[] {
+  const values = [];
+  visit(document, {
+    Argument(node) {
+      if (node.name.value === 'blockHeight' && node.value.kind === 'StringValue') {
+        values.push(node.value.value);
+      }
+    },
+  });
+  return values;
+}
+
+async function fetchLastProcessedHeight(pgClient: Pool, schemaName: string): Promise<bigint | null> {
+  const result = await pgClient.query(`select value from "${schemaName}"._metadata WHERE key = 'lastProcessedHeight'`);
+  if (!result.rowCount) {
+    return null;
+  }
+  return BigInt(result.rows[0].value);
+}
+
+function validateBlockHeightArgs(args: string[], lastProcessedHeight: bigint): GraphQLError | null {
+  for (const arg of args) {
+    if (!arg.match(/^\d+$/)) {
+      return new UserInputError(`Invalid block height: ${arg}`);
+    }
+    const value = BigInt(arg);
+    if (value > lastProcessedHeight) {
+      return new UserInputError(`Block height ${arg} is larger than last processed height: ${lastProcessedHeight}`);
+    }
+  }
+  return null;
+}
+
+// Plugin to check that any provided blockHeight arg is <= lastProcessedHeight
+export function BlockHeightArgValidationPlugin({dbSchema}: {dbSchema: string}): ApolloServerPlugin {
+  return {
+    // eslint-disable-next-line
+    requestDidStart: async () => {
+      return {
+        async responseForOperation({context, document}) {
+          const blockHeightArgs = parseBlockHeightArgs(document);
+          if (!blockHeightArgs) {
+            return null;
+          }
+          const lastProcessedHeight = await fetchLastProcessedHeight(context.pgClient, dbSchema);
+          if (!lastProcessedHeight) {
+            return null;
+          }
+          const error = validateBlockHeightArgs(blockHeightArgs, lastProcessedHeight);
+          if (!error) {
+            return null;
+          }
+          return {errors: [error]};
+        },
+      };
+    },
+  } as ApolloServerPlugin;
+}