Merge pull request #461 from NomicFoundation/fix/limit-library-in-calls

restrict calls and static calls to non library contracts
NomicFoundation · Sep 13, 2023 · b89c9ed · b89c9ed
2 parents 25bbada + f6b48b2
commit b89c9ed
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 9 deletions.
diff --git a/packages/core/src/internal/module-builder.ts b/packages/core/src/internal/module-builder.ts
@@ -6,6 +6,7 @@ import {
   isAccountRuntimeValue,
   isAddressResolvableFuture,
   isArtifactType,
+  isCallableContractFuture,
   isContractFuture,
   isFuture,
   isModuleParameterRuntimeValue,
@@ -18,6 +19,7 @@ import {
   ArtifactContractAtFuture,
   ArtifactContractDeploymentFuture,
   ArtifactLibraryDeploymentFuture,
+  CallableContractFuture,
   ContractFuture,
   IgnitionModule,
   IgnitionModuleResult,
@@ -331,7 +333,7 @@ class IgnitionModuleBuilderImplementation<
   }
 
   public call<ContractNameT extends string, FunctionNameT extends string>(
-    contractFuture: ContractFuture<ContractNameT>,
+    contractFuture: CallableContractFuture<ContractNameT>,
     functionName: FunctionNameT,
     args: ArgumentType[] = [],
     options: CallOptions = {}
@@ -344,7 +346,7 @@ class IgnitionModuleBuilderImplementation<
     this._assertUniqueCallId(futureId);
     this._assertValidValue(options.value, this.call);
     this._assertValidFrom(options.from, this.call);
-    this._assertValidContract(contractFuture, this.call);
+    this._assertValidCallableContract(contractFuture, this.call);
     /* validation end */
 
     const future = new NamedContractCallFutureImplementation(
@@ -373,7 +375,7 @@ class IgnitionModuleBuilderImplementation<
   }
 
   public staticCall<ContractNameT extends string, FunctionNameT extends string>(
-    contractFuture: ContractFuture<ContractNameT>,
+    contractFuture: CallableContractFuture<ContractNameT>,
     functionName: FunctionNameT,
     args: ArgumentType[] = [],
     options: StaticCallOptions = {}
@@ -384,7 +386,7 @@ class IgnitionModuleBuilderImplementation<
     /* validation start */
     this._assertUniqueStaticCallId(futureId);
     this._assertValidFrom(options.from, this.staticCall);
-    this._assertValidContract(contractFuture, this.staticCall);
+    this._assertValidCallableContract(contractFuture, this.staticCall);
     /* validation end */
 
     const future = new NamedStaticCallFutureImplementation(
@@ -754,11 +756,11 @@ class IgnitionModuleBuilderImplementation<
     }
   }
 
-  private _assertValidContract(
-    contract: ContractFuture<string>,
+  private _assertValidCallableContract(
+    contract: CallableContractFuture<string>,
     func: (...[]: any[]) => any
   ) {
-    if (!isContractFuture(contract)) {
+    if (!isCallableContractFuture(contract)) {
       this._throwErrorWithStackTrace(`Invalid contract given`, func);
     }
   }

diff --git a/packages/core/src/type-guards.ts b/packages/core/src/type-guards.ts
@@ -5,6 +5,7 @@ import {
   ArtifactContractAtFuture,
   ArtifactContractDeploymentFuture,
   ArtifactLibraryDeploymentFuture,
+  CallableContractFuture,
   ContractFuture,
   DeploymentFuture,
   FunctionCallFuture,
@@ -102,6 +103,26 @@ export function isContractFuture(
   }
 }
 
+/**
+ * Returns true if future is of type CallableContractFuture<string>.
+ *
+ * @beta
+ */
+export function isCallableContractFuture(
+  future: Future
+): future is CallableContractFuture<string> {
+  switch (future.type) {
+    case FutureType.NAMED_CONTRACT_DEPLOYMENT:
+    case FutureType.ARTIFACT_CONTRACT_DEPLOYMENT:
+    case FutureType.NAMED_CONTRACT_AT:
+    case FutureType.ARTIFACT_CONTRACT_AT:
+      return true;
+
+    default:
+      return false;
+  }
+}
+
 /**
  * Returns true if future is of type AddressResolvable.
  *

diff --git a/packages/core/src/types/module-builder.ts b/packages/core/src/types/module-builder.ts
@@ -6,6 +6,7 @@ import {
   ArtifactContractAtFuture,
   ArtifactContractDeploymentFuture,
   ArtifactLibraryDeploymentFuture,
+  CallableContractFuture,
   ContractFuture,
   Future,
   IgnitionModule,
@@ -177,14 +178,14 @@ export interface IgnitionModuleBuilder {
   ): ArtifactLibraryDeploymentFuture;
 
   call<ContractNameT extends string, FunctionNameT extends string>(
-    contractFuture: ContractFuture<ContractNameT>,
+    contractFuture: CallableContractFuture<ContractNameT>,
     functionName: FunctionNameT,
     args?: ArgumentType[],
     options?: CallOptions
   ): NamedContractCallFuture<ContractNameT, FunctionNameT>;
 
   staticCall<ContractNameT extends string, FunctionNameT extends string>(
-    contractFuture: ContractFuture<ContractNameT>,
+    contractFuture: CallableContractFuture<ContractNameT>,
     functionName: FunctionNameT,
     args?: ArgumentType[],
     options?: StaticCallOptions

diff --git a/packages/core/src/types/module.ts b/packages/core/src/types/module.ts
@@ -75,6 +75,18 @@ export type ContractFuture<ContractNameT extends string> =
   | NamedContractAtFuture<ContractNameT>
   | ArtifactContractAtFuture;
 
+/**
+ * A future representing only contracts that can be called off-chain (i.e. not libraries).
+ * Either an existing one or one that will be deployed.
+ *
+ * @beta
+ */
+export type CallableContractFuture<ContractNameT extends string> =
+  | NamedContractDeploymentFuture<ContractNameT>
+  | ArtifactContractDeploymentFuture
+  | NamedContractAtFuture<ContractNameT>
+  | ArtifactContractAtFuture;
+
 /**
  * A future representing a deployment.
  *

diff --git a/packages/core/test/call.ts b/packages/core/test/call.ts
@@ -464,6 +464,20 @@ describe("call", () => {
           /Invalid contract given/
         );
       });
+
+      it("should not validate a library", () => {
+        assert.throws(
+          () =>
+            buildModule("Module1", (m) => {
+              const another = m.library("Another");
+
+              m.call(another as any, "test");
+
+              return { another };
+            }),
+          /Invalid contract given/
+        );
+      });
     });
 
     describe("stage one", () => {

diff --git a/packages/core/test/staticCall.ts b/packages/core/test/staticCall.ts
@@ -434,6 +434,20 @@ describe("static call", () => {
           /Invalid contract given/
         );
       });
+
+      it("should not validate a library", () => {
+        assert.throws(
+          () =>
+            buildModule("Module1", (m) => {
+              const another = m.library("Another");
+
+              m.staticCall(another as any, "test");
+
+              return { another };
+            }),
+          /Invalid contract given/
+        );
+      });
     });
 
     describe("stage one", () => {