From ee4ec8ae9a22684a975d476fe323c1908022e33e Mon Sep 17 00:00:00 2001
From: Risson <aime.risson.1@gmail.com>
Date: Mon, 7 Aug 2023 16:40:15 +0200
Subject: [PATCH] blacklist methods using tagging

---
 .../interpreter/function_interpreter.ex       | 52 ++++++++++++--
 .../interpreter/library/common/chain_impl.ex  |  7 +-
 .../interpreter/library/common/code.ex        |  3 +
 .../interpreter/library/common/crypto.ex      |  3 +
 .../interpreter/library/common/http_impl.ex   |  5 ++
 .../interpreter/library/common/json.ex        |  3 +
 .../interpreter/library/common/list.ex        |  3 +
 .../interpreter/library/common/map.ex         |  3 +
 .../interpreter/library/common/regex.ex       |  3 +
 .../interpreter/library/common/string.ex      |  3 +
 .../interpreter/library/common/time.ex        |  3 +
 .../interpreter/library/common/token_impl.ex  |  4 ++
 .../contracts/interpreter/library/contract.ex | 14 ++++
 .../interpreter/function_interpreter_test.exs | 68 ++++++++++++++++++-
 .../methods/call_contract_function_test.exs   | 11 +--
 15 files changed, 172 insertions(+), 13 deletions(-)

diff --git a/lib/archethic/contracts/interpreter/function_interpreter.ex b/lib/archethic/contracts/interpreter/function_interpreter.ex
index 29d47201f4..57ef9fad34 100644
--- a/lib/archethic/contracts/interpreter/function_interpreter.ex
+++ b/lib/archethic/contracts/interpreter/function_interpreter.ex
@@ -4,6 +4,7 @@ defmodule Archethic.Contracts.Interpreter.FunctionInterpreter do
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
   alias Archethic.Contracts.Interpreter.Scope
   alias Archethic.Contracts.Interpreter.CommonInterpreter
+  alias Knigge
   require Logger
 
   @doc """
@@ -18,7 +19,7 @@ defmodule Archethic.Contracts.Interpreter.FunctionInterpreter do
     {:ok, function_name, args, ast}
   catch
     {:error, node} ->
-      {:error, node, "unexpectted term"}
+      {:error, node, "unexpected term"}
 
     {:error, node, reason} ->
       {:error, node, reason}
@@ -35,14 +36,14 @@ defmodule Archethic.Contracts.Interpreter.FunctionInterpreter do
     {:ok, function_name, args, ast}
   catch
     {:error, node} ->
-      {:error, node, "unexpecccccted term"}
+      {:error, node, "unexpected term"}
 
     {:error, node, reason} ->
       {:error, node, reason}
   end
 
   def parse(node, _) do
-    {:error, node, "unexpecteeeeed term"}
+    {:error, node, "unexpected term"}
   end
 
   @doc """
@@ -97,11 +98,50 @@ defmodule Archethic.Contracts.Interpreter.FunctionInterpreter do
   # ----------------------------------------------------------------------
   # Ban access to Contract module
   defp prewalk(
-         node = {:__aliases__, _, [atom: "Contract"]},
-         _,
+         node =
+           {{:., _meta, [{:__aliases__, _, [atom: "Contract"]}, {:atom, function_name}]}, _, _},
+         acc,
          _visibility
        ) do
-    throw({:error, node, "Contract is not allowed in function"})
+    absolute_module_atom =
+      Code.ensure_loaded!(
+        String.to_existing_atom("Elixir.Archethic.Contracts.Interpreter.Library.Contract")
+      )
+
+    if absolute_module_atom.tagged_with?(String.to_atom(function_name), :write_contract) do
+      throw({:error, node, "Write contract functions are not allowed in custom functions"})
+    end
+
+    CommonInterpreter.prewalk(node, acc)
+  end
+
+  defp prewalk(
+         node =
+           {{:., _meta, [{:__aliases__, _, [atom: module_name]}, {:atom, function_name}]}, _, _},
+         acc,
+         true
+       ) do
+    absolute_module_atom =
+      Code.ensure_loaded!(
+        String.to_existing_atom(
+          "Elixir.Archethic.Contracts.Interpreter.Library.Common.#{module_name}"
+        )
+      )
+
+    absolute_module_atom =
+      try do
+        %Knigge.Options{default: module} = Knigge.options!(absolute_module_atom)
+        module
+      rescue
+        _ ->
+          absolute_module_atom
+      end
+
+    if absolute_module_atom.tagged_with?(String.to_atom(function_name), :io) do
+      throw({:error, node, "IO function calls not allowed in public functions"})
+    end
+
+    CommonInterpreter.prewalk(node, acc)
   end
 
   defp prewalk(node = {{:atom, function_name}, _, args}, _acc, true)
diff --git a/lib/archethic/contracts/interpreter/library/common/chain_impl.ex b/lib/archethic/contracts/interpreter/library/common/chain_impl.ex
index 26388e1317..4e07b66ebe 100644
--- a/lib/archethic/contracts/interpreter/library/common/chain_impl.ex
+++ b/lib/archethic/contracts/interpreter/library/common/chain_impl.ex
@@ -5,16 +5,19 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.ChainImpl do
   alias Archethic.Contracts.Interpreter.Library.Common.Chain
   alias Archethic.Contracts.Interpreter.Legacy.UtilsInterpreter
   alias Archethic.Contracts.ContractConstants, as: Constants
-
+  alias Archethic.Tag
   alias Archethic.TransactionChain.Transaction.ValidationStamp.LedgerOperations
 
   @behaviour Chain
+  use Tag
 
+  @tag [:io]
   @impl Chain
   defdelegate get_genesis_address(address),
     to: Legacy.Library,
     as: :get_genesis_address
 
+  @tag [:io]
   @impl Chain
   def get_first_transaction_address(address) do
     try do
@@ -24,6 +27,7 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.ChainImpl do
     end
   end
 
+  @tag [:io]
   @impl Chain
   def get_genesis_public_key(public_key) do
     try do
@@ -33,6 +37,7 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.ChainImpl do
     end
   end
 
+  @tag [:io]
   @impl Chain
   def get_transaction(address) do
     address
diff --git a/lib/archethic/contracts/interpreter/library/common/code.ex b/lib/archethic/contracts/interpreter/library/common/code.ex
index a7a16b28da..4401fdd280 100644
--- a/lib/archethic/contracts/interpreter/library/common/code.ex
+++ b/lib/archethic/contracts/interpreter/library/common/code.ex
@@ -2,9 +2,12 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.Code do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
 
+  use Tag
+
   @spec is_same?(binary(), binary()) :: boolean()
   def is_same?(first_code, second_code) do
     first_ast = first_code |> Interpreter.sanitize_code(ignore_meta?: true)
diff --git a/lib/archethic/contracts/interpreter/library/common/crypto.ex b/lib/archethic/contracts/interpreter/library/common/crypto.ex
index 6e4f5cb69b..5888df10c3 100644
--- a/lib/archethic/contracts/interpreter/library/common/crypto.ex
+++ b/lib/archethic/contracts/interpreter/library/common/crypto.ex
@@ -2,10 +2,13 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.Crypto do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
   alias Archethic.Contracts.Interpreter.Legacy
   alias Archethic.Contracts.Interpreter.Legacy.UtilsInterpreter
 
+  use Tag
+
   @spec hash(binary(), binary()) :: binary()
   def hash(content, algo \\ "sha256")
 
diff --git a/lib/archethic/contracts/interpreter/library/common/http_impl.ex b/lib/archethic/contracts/interpreter/library/common/http_impl.ex
index 53a8aa3f40..493dc26309 100644
--- a/lib/archethic/contracts/interpreter/library/common/http_impl.ex
+++ b/lib/archethic/contracts/interpreter/library/common/http_impl.ex
@@ -7,14 +7,18 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.HttpImpl do
   other processes, we use it from inside a Task.
   """
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.Library
   alias Archethic.Contracts.Interpreter.Library.Common.Http
   alias Archethic.TaskSupervisor
 
+  use Tag
+
   @behaviour Http
   @threshold 256 * 1024
   @timeout Application.compile_env(:archethic, [__MODULE__, :timeout], 2_000)
 
+  @tag [:io]
   @impl Http
   def fetch(uri) do
     task =
@@ -45,6 +49,7 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.HttpImpl do
     end
   end
 
+  @tag [:io]
   @impl Http
   def fetch_many(uris) do
     uris_count = length(uris)
diff --git a/lib/archethic/contracts/interpreter/library/common/json.ex b/lib/archethic/contracts/interpreter/library/common/json.ex
index 3078b9f39b..0d2e47c46d 100644
--- a/lib/archethic/contracts/interpreter/library/common/json.ex
+++ b/lib/archethic/contracts/interpreter/library/common/json.ex
@@ -2,9 +2,12 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.Json do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
   alias Archethic.Contracts.Interpreter.Legacy
 
+  use Tag
+
   @spec path_extract(String.t(), String.t()) :: String.t()
   defdelegate path_extract(text, path),
     to: Legacy.Library,
diff --git a/lib/archethic/contracts/interpreter/library/common/list.ex b/lib/archethic/contracts/interpreter/library/common/list.ex
index 23b1c7b9b1..0d4a62f064 100644
--- a/lib/archethic/contracts/interpreter/library/common/list.ex
+++ b/lib/archethic/contracts/interpreter/library/common/list.ex
@@ -2,8 +2,11 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.List do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
 
+  use Tag
+
   @spec at(list(), integer() | float()) :: any()
   def at(list, idx) do
     cond do
diff --git a/lib/archethic/contracts/interpreter/library/common/map.ex b/lib/archethic/contracts/interpreter/library/common/map.ex
index 797600f9d5..0388e9abe6 100644
--- a/lib/archethic/contracts/interpreter/library/common/map.ex
+++ b/lib/archethic/contracts/interpreter/library/common/map.ex
@@ -2,8 +2,11 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.Map do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
 
+  use Tag
+
   @spec new() :: map()
   defdelegate new(),
     to: Map
diff --git a/lib/archethic/contracts/interpreter/library/common/regex.ex b/lib/archethic/contracts/interpreter/library/common/regex.ex
index 4625cbd0d6..57e7564783 100644
--- a/lib/archethic/contracts/interpreter/library/common/regex.ex
+++ b/lib/archethic/contracts/interpreter/library/common/regex.ex
@@ -2,9 +2,12 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.Regex do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
   alias Archethic.Contracts.Interpreter.Legacy
 
+  use Tag
+
   @spec match?(binary(), binary()) :: boolean()
   defdelegate match?(text, pattern),
     to: Legacy.Library,
diff --git a/lib/archethic/contracts/interpreter/library/common/string.ex b/lib/archethic/contracts/interpreter/library/common/string.ex
index 432063b1a1..fda1d88d29 100644
--- a/lib/archethic/contracts/interpreter/library/common/string.ex
+++ b/lib/archethic/contracts/interpreter/library/common/string.ex
@@ -2,8 +2,11 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.String do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.ASTHelper, as: AST
 
+  use Tag
+
   @spec size(String.t()) :: integer()
   defdelegate size(str),
     to: String,
diff --git a/lib/archethic/contracts/interpreter/library/common/time.ex b/lib/archethic/contracts/interpreter/library/common/time.ex
index 778d91c9b3..e126604c8e 100644
--- a/lib/archethic/contracts/interpreter/library/common/time.ex
+++ b/lib/archethic/contracts/interpreter/library/common/time.ex
@@ -2,8 +2,11 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.Time do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.Scope
 
+  use Tag
+
   @doc """
   Returns the Unix timestamp of the trigger (it is approximately the same as current time).
   We cannot use "now" because it is not determinist.
diff --git a/lib/archethic/contracts/interpreter/library/common/token_impl.ex b/lib/archethic/contracts/interpreter/library/common/token_impl.ex
index d08784833a..d05e76130f 100644
--- a/lib/archethic/contracts/interpreter/library/common/token_impl.ex
+++ b/lib/archethic/contracts/interpreter/library/common/token_impl.ex
@@ -2,8 +2,12 @@ defmodule Archethic.Contracts.Interpreter.Library.Common.TokenImpl do
   @moduledoc false
   @behaviour Archethic.Contracts.Interpreter.Library.Common.Token
 
+  alias Archethic.Tag
   alias Archethic.Contracts.Interpreter.Legacy
 
+  use Tag
+
+  @tag [:io]
   @impl Archethic.Contracts.Interpreter.Library.Common.Token
   defdelegate fetch_id_from_address(address),
     to: Legacy.Library,
diff --git a/lib/archethic/contracts/interpreter/library/contract.ex b/lib/archethic/contracts/interpreter/library/contract.ex
index 14fc1d5097..21f3e235c8 100644
--- a/lib/archethic/contracts/interpreter/library/contract.ex
+++ b/lib/archethic/contracts/interpreter/library/contract.ex
@@ -11,11 +11,16 @@ defmodule Archethic.Contracts.Interpreter.Library.Contract do
   alias Archethic.TransactionChain.Transaction
   alias Archethic.Contracts.Interpreter.Legacy.TransactionStatements
   alias Archethic.Utils
+  alias Archethic.Tag
 
+  use Tag
+
+  @tag [:write_contract]
   @spec set_type(Transaction.t(), binary()) :: Transaction.t()
   defdelegate set_type(next_tx, type),
     to: TransactionStatements
 
+  @tag [:write_contract]
   @spec set_content(Transaction.t(), binary() | integer() | float()) :: Transaction.t()
   def set_content(next_tx, content) when is_binary(content) do
     put_in(next_tx, [Access.key(:data), Access.key(:content)], content)
@@ -29,45 +34,54 @@ defmodule Archethic.Contracts.Interpreter.Library.Contract do
     )
   end
 
+  @tag [:write_contract]
   @spec set_code(Transaction.t(), binary()) :: Transaction.t()
   defdelegate set_code(next_tx, args),
     to: TransactionStatements
 
+  @tag [:write_contract]
   @spec add_recipient(Transaction.t(), binary()) :: Transaction.t()
   defdelegate add_recipient(next_tx, args),
     to: TransactionStatements
 
+  @tag [:write_contract]
   @spec add_recipients(Transaction.t(), list(binary())) :: Transaction.t()
   defdelegate add_recipients(next_tx, args),
     to: TransactionStatements
 
+  @tag [:write_contract]
   @spec add_uco_transfer(Transaction.t(), map()) :: Transaction.t()
   def add_uco_transfer(next_tx, args) do
     args = Map.update!(args, "amount", &Utils.to_bigint/1)
     TransactionStatements.add_uco_transfer(next_tx, Map.to_list(args))
   end
 
+  @tag [:write_contract]
   @spec add_uco_transfers(Transaction.t(), list(map())) :: Transaction.t()
   def add_uco_transfers(next_tx, args) do
     Enum.reduce(args, next_tx, &add_uco_transfer(&2, &1))
   end
 
+  @tag [:write_contract]
   @spec add_token_transfer(Transaction.t(), map()) :: Transaction.t()
   def add_token_transfer(next_tx, args) do
     args = Map.update!(args, "amount", &Utils.to_bigint/1)
     TransactionStatements.add_token_transfer(next_tx, Map.to_list(args))
   end
 
+  @tag [:write_contract]
   @spec add_token_transfers(Transaction.t(), list(map())) :: Transaction.t()
   def add_token_transfers(next_tx, args) do
     Enum.reduce(args, next_tx, &add_token_transfer(&2, &1))
   end
 
+  @tag [:write_contract]
   @spec add_ownership(Transaction.t(), map()) :: Transaction.t()
   def add_ownership(next_tx, args) do
     TransactionStatements.add_ownership(next_tx, Map.to_list(args))
   end
 
+  @tag [:write_contract]
   @spec add_ownerships(Transaction.t(), list(map())) :: Transaction.t()
   def add_ownerships(next_tx, args) do
     casted_args = Enum.map(args, &Map.to_list/1)
diff --git a/test/archethic/contracts/interpreter/function_interpreter_test.exs b/test/archethic/contracts/interpreter/function_interpreter_test.exs
index 970af5020a..d66116d269 100644
--- a/test/archethic/contracts/interpreter/function_interpreter_test.exs
+++ b/test/archethic/contracts/interpreter/function_interpreter_test.exs
@@ -74,14 +74,66 @@ defmodule Archethic.Contracts.Interpreter.FunctionInterpreterTest do
       end
       """
 
-      assert {:error, _, "Contract is not allowed in function"} =
+      assert {:error, _, "Write contract functions are not allowed in custom functions"} =
+               code
+               |> Interpreter.sanitize_code()
+               |> elem(1)
+               |> FunctionInterpreter.parse([])
+
+      code = ~S"""
+      export fun test_public do
+        Contract.set_content "hello"
+      end
+      """
+
+      assert {:error, _, "Write contract functions are not allowed in custom functions"} =
                code
                |> Interpreter.sanitize_code()
                |> elem(1)
                |> FunctionInterpreter.parse([])
     end
 
-    test "should be able to parse when there is whitelisted module" do
+    test "should not be able to use IO functions in public function" do
+      code = ~S"""
+      export fun test_public do
+        Chain.get_genesis_address("hello")
+      end
+      """
+
+      assert {:error, _, "IO function calls not allowed in public functions"} =
+               code
+               |> Interpreter.sanitize_code()
+               |> elem(1)
+               |> FunctionInterpreter.parse([])
+    end
+
+    test "should be able to use IO functions in private function" do
+      code = ~S"""
+      fun test_private do
+        Chain.get_genesis_address("hello")
+      end
+      """
+
+      assert {:ok, _, _, _} =
+               code
+               |> Interpreter.sanitize_code()
+               |> elem(1)
+               |> FunctionInterpreter.parse([])
+
+      code = ~S"""
+      export fun test_public do
+        Chain.get_burn_address()
+      end
+      """
+
+      assert {:ok, _, _, _} =
+               code
+               |> Interpreter.sanitize_code()
+               |> elem(1)
+               |> FunctionInterpreter.parse([])
+    end
+
+    test "should be able to parse public function when a module's function is not IO" do
       code = ~S"""
       fun test do
        Json.to_string "[1,2,3]"
@@ -93,6 +145,18 @@ defmodule Archethic.Contracts.Interpreter.FunctionInterpreterTest do
                |> Interpreter.sanitize_code()
                |> elem(1)
                |> FunctionInterpreter.parse([])
+
+      code = ~S"""
+      export fun test_public do
+        Chain.get_burn_address()
+      end
+      """
+
+      assert {:ok, _, _, _} =
+               code
+               |> Interpreter.sanitize_code()
+               |> elem(1)
+               |> FunctionInterpreter.parse([])
     end
 
     test "should not be able to call non declared function" do
diff --git a/test/archethic_web/api/jsonrpc/methods/call_contract_function_test.exs b/test/archethic_web/api/jsonrpc/methods/call_contract_function_test.exs
index 890a93ddb2..f420d6f1b1 100644
--- a/test/archethic_web/api/jsonrpc/methods/call_contract_function_test.exs
+++ b/test/archethic_web/api/jsonrpc/methods/call_contract_function_test.exs
@@ -201,6 +201,7 @@ defmodule ArchethicWeb.API.JsonRPC.Methods.CallContractFunctionTest do
       export fun bob() do 
         "hello bob"
       end
+
       export fun hello() do 
         bob()
       end
@@ -240,8 +241,9 @@ defmodule ArchethicWeb.API.JsonRPC.Methods.CallContractFunctionTest do
         args: []
       }
 
-      assert {:error, :function_failure, "There was an error while executing the function",
-              "hello/0"} = CallContractFunction.execute(params)
+      assert {:error, :parsing_contract, _,
+              "not allowed to call function from public function - bob - L8"} =
+               CallContractFunction.execute(params)
     end
 
     test "should not be able to call a private function from a public function" do
@@ -290,8 +292,9 @@ defmodule ArchethicWeb.API.JsonRPC.Methods.CallContractFunctionTest do
         args: []
       }
 
-      assert {:error, :function_failure, "There was an error while executing the function",
-              "hello/0"} = CallContractFunction.execute(params)
+      assert {:error, :parsing_contract, _,
+              "not allowed to call function from public function - bob - L7"} =
+               CallContractFunction.execute(params)
     end
 
     test "should return error when called function does not exist" do