From 884c938136edeb98b38d0b5ed1f4ef1230315048 Mon Sep 17 00:00:00 2001
From: Yanju Chen <chyanju@gmail.com>
Date: Wed, 28 Feb 2024 11:46:23 -0800
Subject: [PATCH] update CI and fixed infoleak detector

---
 tests/scripts/test-infoleak.py      |  7 +++++-
 tests/scripts/test-unused.py        |  2 +-
 tests/test4.ipynb                   | 37 ++++++++++++++++-------------
 vanguard/aleo/detectors/infoleak.py |  3 +++
 4 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/tests/scripts/test-infoleak.py b/tests/scripts/test-infoleak.py
index be202e0..bfa5c92 100644
--- a/tests/scripts/test-infoleak.py
+++ b/tests/scripts/test-infoleak.py
@@ -6,6 +6,11 @@
     build_path = f"./tests/public/{project_name}/build/"
     env = AleoEnvironment(build_path)
     main = env.main
+
     fid = "ex0"
     res, info = detector_infoleak(env, main.id, fid)
-    assert res, f"Test failed for {fid}, expected: True, got: {res}"
\ No newline at end of file
+    assert res, f"Test failed for {fid}, expected: True, got: {res}"
+
+    fid = "ex1"
+    res, info = detector_infoleak(env, main.id, fid)
+    assert not res, f"Test failed for {fid}, expected: False, got: {res}"
\ No newline at end of file
diff --git a/tests/scripts/test-unused.py b/tests/scripts/test-unused.py
index c907b7c..a5f8960 100644
--- a/tests/scripts/test-unused.py
+++ b/tests/scripts/test-unused.py
@@ -9,7 +9,7 @@
     
     fid = "ex0"
     res, info = detector_unused(env, main.id, fid)
-    assert not res, f"Test failed for {fid}, expected: True, got: {res}"
+    assert not res, f"Test failed for {fid}, expected: False, got: {res}"
 
     fid = "ex1"
     res, info = detector_unused(env, main.id, fid)
diff --git a/tests/test4.ipynb b/tests/test4.ipynb
index 47815ce..11cbf13 100644
--- a/tests/test4.ipynb
+++ b/tests/test4.ipynb
@@ -38,14 +38,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "# [debug] deploy: main.aleo\n",
-      "# [debug] deploy: helpers.aleo\n"
+      "# [debug] deploy: main.aleo\n"
      ]
     }
    ],
    "source": [
-    "project_name = \"divz0\"\n",
-    "# project_name = \"infoleak0\"\n",
+    "# project_name = \"divz0\"\n",
+    "project_name = \"infoleak0\"\n",
     "# project_name = \"unused0\"\n",
     "# project_name = \"rtcnst0\"\n",
     "build_path = f\"./tests/public/{project_name}/build/\"\n",
@@ -63,26 +62,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 4,
    "id": "0dc6fcff-f8b2-44cb-8919-fda9b23c0b49",
    "metadata": {},
    "outputs": [
     {
-     "ename": "KeyError",
-     "evalue": "'ex26'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[9], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m fid \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mex26\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 2\u001b[0m \u001b[43mdetector_divz\u001b[49m\u001b[43m(\u001b[49m\u001b[43menv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mid\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfid\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreadable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/Desktop/UCSB/24winter/vanguard-aleo/vanguard/aleo/detectors/divz.py:158\u001b[0m, in \u001b[0;36mdetector_divz\u001b[0;34m(env, pid, fid, readable)\u001b[0m\n\u001b[1;32m    155\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdetector_divz\u001b[39m(env: AleoEnvironment, pid: \u001b[38;5;28mstr\u001b[39m, fid: \u001b[38;5;28mstr\u001b[39m, readable\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m):\n\u001b[1;32m    156\u001b[0m     \u001b[38;5;66;03m# initialize\u001b[39;00m\n\u001b[1;32m    157\u001b[0m     prog: AleoProgram \u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39mprograms[pid]\n\u001b[0;32m--> 158\u001b[0m     func: AleoFunction \u001b[38;5;241m=\u001b[39m \u001b[43mprog\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunctions\u001b[49m\u001b[43m[\u001b[49m\u001b[43mfid\u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m    160\u001b[0m     \u001b[38;5;66;03m# store problematic instructions\u001b[39;00m\n\u001b[1;32m    161\u001b[0m     lines \u001b[38;5;241m=\u001b[39m []\n",
-      "\u001b[0;31mKeyError\u001b[0m: 'ex26'"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "r0 -> r0\n",
+      "# [debug] sig_in: r0\n"
      ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "(True, [('r0', 'r0')])"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "fid = \"ex26\"\n",
-    "detector_divz(env, main.id, fid, readable=True)"
+    "fid = \"ex0\"\n",
+    "detector_infoleak(env, main.id, fid, readable=True)"
    ]
   },
   {
diff --git a/vanguard/aleo/detectors/infoleak.py b/vanguard/aleo/detectors/infoleak.py
index f3725bf..f198a1a 100644
--- a/vanguard/aleo/detectors/infoleak.py
+++ b/vanguard/aleo/detectors/infoleak.py
@@ -9,6 +9,9 @@ def detector_infoleak(env: AleoEnvironment, pid: str, fid: str, readable=False):
     func: AleoFunction = prog.functions[fid]
     
     edges = get_dfg_edges(env, pid, fid, hash=False)
+    # special edge a -> a so that direct returning of a variable is accounted
+    for k,v in func.inputs:
+        edges.append((k, k))
 
     # DEBUG
     for p in edges: