From f9abacb9ec3cf003872eaa112ebf32543ff9a8e3 Mon Sep 17 00:00:00 2001
From: Rnine <yangyizhuo@buaa.edu.cn>
Date: Thu, 12 Dec 2024 16:28:41 +0800
Subject: [PATCH] recognize AppStorage singleton method

Signed-off-by: Rnine <yangyizhuo@buaa.edu.cn>
---
 src/callgraph/pointerAnalysis/PagBuilder.ts   | 47 +++++++++++++++++--
 .../pta/AppStorage/AppStorageCase.ts          | 18 ++++---
 2 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/src/callgraph/pointerAnalysis/PagBuilder.ts b/src/callgraph/pointerAnalysis/PagBuilder.ts
index 44023d0b..195560f1 100644
--- a/src/callgraph/pointerAnalysis/PagBuilder.ts
+++ b/src/callgraph/pointerAnalysis/PagBuilder.ts
@@ -1154,7 +1154,11 @@ export class PagBuilder {
 
         let returnValues = arkMethod.getReturnValues();
 
-        let result = this.isValueConnected([...funcPag.getInternalEdges()!], heapObjects, returnValues);
+        let result = this.isValueConnected(
+            [...funcPag.getInternalEdges()!], 
+            [...funcPag.getNormalCallSites()], 
+            heapObjects, returnValues
+        );
         this.singletonFuncMap.set(funcID, result);
         if (result) {
             logger.info(`function ${funcID} is marked as singleton function`);
@@ -1162,10 +1166,33 @@ export class PagBuilder {
         return result;
     }
 
-    private isValueConnected(edges: IntraProceduralEdge[], leftNodes: Value[], targetNodes: Value[]): boolean {
+    private isValueConnected(edges: IntraProceduralEdge[], callSites: CallSite[], leftNodes: Value[], targetNodes: Value[]): boolean {
         // build funcPag graph
         const graph = new Map<Value, Value[]>();
         let hasStaticFieldOrGlobalVar: boolean = false;
+        let appStorageMap = new Map<string, Value>();
+
+        for (const callSite of callSites) {
+            let callee = callSite.callStmt.getInvokeExpr()?.getMethodSignature();
+            if (callee?.toString() === '@ohos/component/common_ts_ets_api.d.ts: AppStorage.[static]setOrCreate(string, T)') {
+                let keyValue = callSite.args![0];
+                if (keyValue instanceof Constant) {
+                    appStorageMap.set(keyValue.getValue(), callSite.args![1]);
+                }
+            } else if (callee?.toString() === '@ohos/component/common_ts_ets_api.d.ts: AppStorage.[static]get(string)') {
+                let keyValue = callSite.args![0];
+                if (keyValue instanceof Constant) {
+                    let retrievedValue: Value | undefined;
+                    if (callSite.callStmt instanceof ArkAssignStmt) {
+                        retrievedValue = callSite.callStmt.getLeftOp();
+                    }
+                    if (retrievedValue) {
+                        // save the return value with key: key + '_retrieved'
+                        appStorageMap.set(keyValue.getValue() + '_retrieved', retrievedValue);
+                    }
+                }
+            }
+        }
 
         for (const edge of edges) {
             let dst = this.getRealInstanceRef(edge.dst);
@@ -1184,7 +1211,7 @@ export class PagBuilder {
             graph.get(src)!.push(dst);
         }
 
-        if (!hasStaticFieldOrGlobalVar) {
+        if (!hasStaticFieldOrGlobalVar && appStorageMap.size === 0) {
             return false;
         }
 
@@ -1192,10 +1219,24 @@ export class PagBuilder {
             for (const leftNode of leftNodes) {
                 const visited = new Set<Value>();
                 let meetStaticField = false;
+
                 if (this.funcPagDfs(graph, visited, leftNode, targetNode, meetStaticField)) {
                     return true; // a value pair that satisfy condition
                 }
 
+                // Check for connections via appStorage
+                for (const [key, value] of appStorageMap) {
+                    /**
+                     * check `setOrCreate` API with the init key
+                     * check `get` API with the constructed key 
+                     * if init value is connected with init key and return value is connected with constructed key, return true
+                     **/ 
+                    if (this.funcPagDfs(graph, visited, leftNode, value, true) &&
+                        this.funcPagDfs(graph, visited, appStorageMap.get(key + '_retrieved')!, targetNode, true)) {
+                        return true;
+                    }
+                }
+
                 if (!meetStaticField) {
                     break; // heap obj will not deal any more
                 }
diff --git a/tests/resources/pta/AppStorage/AppStorageCase.ts b/tests/resources/pta/AppStorage/AppStorageCase.ts
index 864e30bb..9009d130 100644
--- a/tests/resources/pta/AppStorage/AppStorageCase.ts
+++ b/tests/resources/pta/AppStorage/AppStorageCase.ts
@@ -18,15 +18,13 @@ namespace AppStorageCase {
 
     }
 
-    function main() {
-        let a: A = new A();
-
-        AppStorage.setOrCreate('propertyA', a);
-        let b = AppStorage.get('propertyA');
-
-        let link: SubscribedAbstractProperty<A> = AppStorage.link('propertyB');
-        link.set(b);
-        let prop: SubscribedAbstractProperty<A> = AppStorage.prop('propertyB');
-        prop.set(new A());
+    class Main {
+        static main() {
+            if (AppStorage.get('PropA') == null) {
+                AppStorage.setOrCreate('PropA', new A());
+            }
+            let manager: A = AppStorage.get('PropA');
+            return manager;
+        }
     }
 }
\ No newline at end of file
-- 
Gitee