refactor: upgrade Gemini SDK, improve logging, and update Tree-sitter bindings

benchmark/Cpp/toy/RACE/race.c

@@ -0,0 +1,106 @@
+//gcc chall.c -m32 -pie -fstack-protector-all -o chall
+
+#include<stdio.h>
+#include<stdlib.h>
+#include<string.h>
+#include <pthread.h>
+
+unsigned int a = 0;
+unsigned int b = 0;
+unsigned int a_sleep = 0;
+int flag = 1;
+int pstr1 = 1;
+int ret1;
+pthread_t th1;
+void * th_ret = NULL;
+
+void menu_go(){
+    if(a_sleep == 0){
+        a = a + 5;
+    }else{
+        a_sleep = 0;
+    }
+
+    b = b + 2;
+}
+
+int *menu_chance(){
+    if(a<=b){
+        puts("No");
+        return 0;
+    }
+
+    if(flag == 1){
+        a_sleep = 1;
+        sleep(1);
+        flag = 0;
+    }
+    else{
+        puts("Only have one chance");
+    }
+    return 0;
+}
+
+
+void menu_test(){
+    if( b>a ){
+        puts("Win!");
+        system("/bin/sh");
+        exit(0);
+    }else{
+        puts("Lose!");
+        exit(0);
+    }
+}
+
+void menu_exit(){
+    puts("Bye");
+    exit(0);
+}
+
+void menu(){
+    printf("***** race *****\n");
+    printf("*** 1:Go\n*** 2:Chance\n*** 3:Test\n*** 4:Exit \n");
+    printf("*************************************\n");
+    printf("Choice> ");
+    int choose;
+    scanf("%d",&choose);
+    switch(choose)
+    {
+    case 1:
+        menu_go();
+        break;
+    case 2:
+        ret1 = pthread_create(&th1, NULL, menu_chance, &pstr1);
+        break;
+    case 3:
+        menu_test();
+        break;
+    case 4:
+        menu_exit();
+        break;
+    default:
+        return;
+    }
+    return;
+
+}
+
+
+void init(){
+    setbuf(stdin, 0LL);
+    setbuf(stdout, 0LL);
+    setbuf(stderr, 0LL);
+
+    while (1)
+    {
+        menu();
+    }
+
+}
+
+int main(){
+    init();
+    return 0;
+}
+

lib/build.py

@@ -1,8 +1,8 @@
 import os
-
-from tree_sitter import Language, Parser
 from pathlib import Path
 
+from tree_sitter import Language
+
 cwd = Path(__file__).resolve().parent.absolute()
 
 # clone tree-sitter if necessary

requirements.txt

@@ -5,7 +5,7 @@ torch
 tiktoken
 replicate
 openai
-google-generativeai
+google-genai
 tqdm
 networkx
 streamlit

src/agent/dfbscan.py

@@ -17,6 +17,7 @@
 from tstool.dfbscan_extractor.Cpp.Cpp_MLK_extractor import *
 from tstool.dfbscan_extractor.Cpp.Cpp_NPD_extractor import *
 from tstool.dfbscan_extractor.Cpp.Cpp_UAF_extractor import *
+from tstool.dfbscan_extractor.Cpp.Cpp_Race_extractor import *
 from tstool.dfbscan_extractor.Java.Java_NPD_extractor import *
 from tstool.dfbscan_extractor.Python.Python_NPD_extractor import *
 from tstool.dfbscan_extractor.Go.Go_NPD_extractor import *
@@ -103,6 +104,8 @@ def __obtain_extractor(self) -> DFBScanExtractor:
                 return Cpp_NPD_Extractor(self.ts_analyzer)
             elif self.bug_type == "UAF":
                 return Cpp_UAF_Extractor(self.ts_analyzer)
+            elif self.bug_type == "RACE":
+                return Cpp_Race_Extractor(self.ts_analyzer)
         elif self.language == "Java":
             if self.bug_type == "NPD":
                 return Java_NPD_Extractor(self.ts_analyzer)

src/llmtool/LLM_utils.py

@@ -2,7 +2,8 @@
 from openai import *
 from pathlib import Path
 from typing import Tuple
-import google.generativeai as genai
+from google import genai
+from google.genai import types
 import anthropic
 import signal
 import sys
@@ -87,27 +88,53 @@ def run_with_timeout(self, func, timeout):
                 ("Operation timed out")
                 return ""
             except Exception as e:
+                self.logger.print_console(f"Operation failed: {e}")
                 self.logger.print_log(f"Operation failed: {e}")
                 return ""
 
     def infer_with_gemini(self, message: str) -> str:
-        """Infer using the Gemini model from Google Generative AI"""
-        gemini_model = genai.GenerativeModel("gemini-pro")
+        """Infer using the latest Gemini SDK (google-genai)"""
+        api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
+
+        if not api_key:
+            raise EnvironmentError(
+                "Please set the GOOGLE_API_KEY or GEMINI_API_KEY environment variable."
+            )
+
+        client = genai.Client(api_key=api_key)
+
+        model_name = self.online_model_name
+        if model_name == "gemini-pro":
+            model_name = "gemini-2.0-flash"
 
         def call_api():
-            message_with_role = self.systemRole + "\n" + message
             safety_settings = [
-                {
-                    "category": "HARM_CATEGORY_DANGEROUS",
-                    "threshold": "BLOCK_NONE",
-                },
-                # ...existing safety settings...
+                types.SafetySetting(
+                    category=types.HarmCategory.HARM_CATEGORY_HARASSMENT,
+                    threshold=types.HarmBlockThreshold.BLOCK_NONE,
+                ),
+                types.SafetySetting(
+                    category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+                    threshold=types.HarmBlockThreshold.BLOCK_NONE,
+                ),
+                types.SafetySetting(
+                    category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+                    threshold=types.HarmBlockThreshold.BLOCK_NONE,
+                ),
+                types.SafetySetting(
+                    category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                    threshold=types.HarmBlockThreshold.BLOCK_NONE,
+                ),
             ]
-            response = gemini_model.generate_content(
-                message_with_role,
-                safety_settings=safety_settings,
-                generation_config=genai.types.GenerationConfig(
-                    temperature=self.temperature
+
+            response = client.models.generate_content(
+                model=model_name,
+                contents=message,
+                config=types.GenerateContentConfig(
+                    system_instruction=self.systemRole,
+                    temperature=self.temperature,
+                    max_output_tokens=self.max_output_length,
+                    safety_settings=safety_settings,
                 ),
             )
             return response.text
@@ -118,14 +145,13 @@ def call_api():
             try:
                 output = self.run_with_timeout(call_api, timeout=50)
                 if output:
-                    self.logger.print_log("Inference succeeded...")
+                    self.logger.print_log(f"Gemini ({model_name}) inference succeeded...")
                     return output
             except Exception as e:
-                self.logger.print_log(f"API error: {e}")
+                self.logger.print_log(f"Gemini API error: {e}")
             time.sleep(2)
 
         return ""
-
     def infer_with_openai_model(self, message):
         """Infer using the OpenAI model"""
         api_key = os.environ.get("OPENAI_API_KEY").split(":")[0]
@@ -188,7 +214,8 @@ def infer_with_deepseek_model(self, message):
         """
         Infer using the DeepSeek model
         """
-        api_key = os.environ.get("DEEPSEEK_API_KEY2")
+        self.logger.print_console(f"Calling DeepSeek API ({self.online_model_name})...")
+        api_key = os.environ.get("DEEPSEEK_API_KEY") or os.environ.get("OPENAI_API_KEY2")
         model_input = [
             {
                 "role": "system",
@@ -298,6 +325,7 @@ def call_api():
 
     def infer_with_claude_key(self, message):
         """Infer using the Claude model via API key, with thinking mode for 3.7"""
+        self.logger.print_console(f"Calling Claude API ({self.online_model_name})...")
         api_key = os.environ.get("ANTHROPIC_API_KEY")
         if not api_key:
             raise EnvironmentError(

src/prompt/Cpp/dfbscan/path_validator.json

@@ -10,6 +10,7 @@
     "- If the function exits or returns before reaching the sink or relevant propagation sites (such as call sites), then the path is unreachable, so answer No.",
     "- Analyze the conditions on each sub-path within a function. You should infer the outcome of these conditions from branch details and then check whether the conditions across sub-paths conflict. If they do, then the overall path is unreachable.",
     "- Examine the values of relevant variables. If those values contradict the related branch conditions necessary to trigger the bug, the path is unreachable and you should answer No.",
+    "- In the RACE detection, if the shared resource is accessed within a critical section (protected by locks like mutex) or if the access is atomic, then consider the path safe (unreachable for bug) and answer No.",
     "In summary, evaluate the condition of each sub-path, verify possible conflicts, and then decide whether the entire propagation path is reachable."
   ],
   "question_template": [
@@ -132,6 +133,31 @@
     "2. In the 'flag' branch, the condition at line 5 checks if p is not NULL.",
     "3. Since p remains NULL, the condition fails and the else branch at line 7 is executed, preventing any dereference at line 6.",
     "Therefore, this guarded path is unreachable and does not cause the NPD bug.",
+    "Answer: No.",
+    "",
+  "Example 5:",
+    "User:",
+    "Consider the following program which updates a shared resource based on a condition:",
+    "```cpp",
+    "1. int g_resource = 0;",
+    "2. std::mutex resource_mutex;",
+    "3. void process_resource() {",
+    "4.      // Start of critical section",
+    "5.      resource_mutex.lock();",
+    "6.      if (g_resource < 100) {  // Read access",
+    "7.          g_resource += 5;     // Write access",
+    "8.      }",
+    "9.      resource_mutex.unlock();",
+    "10.     // End of critical section",
+    "11. }",
+    "```",
+    "Does the following propagation path cause the RACE bug?",
+    "`g_resource` at line 1 --> `g_resource += 5` at line 7",
+    "Explanation:",
+    "1. The code performs a 'Check-Then-Act' operation (Read at line 6, Write at line 7) on the shared variable `g_resource`.",
+    "2. Both the Read and Write operations are strictly enclosed between `resource_mutex.lock()` (line 5) and `resource_mutex.unlock()` (line 9).",
+    "3. This constitutes a valid critical section. The mutex ensures atomicity: no other thread can modify `g_resource` between the check (line 6) and the update (line 7).",
+    "4. Since the access is fully protected by a lock, it is thread-safe.",
     "Answer: No."
   ],
   "additional_fact": [

src/repoaudit.py

@@ -14,7 +14,7 @@
 from typing import List
 
 default_dfbscan_checkers = {
-    "Cpp": ["MLK", "NPD", "UAF"],
+    "Cpp": ["MLK", "NPD", "UAF", "RACE"],
     "Java": ["NPD"],
     "Python": ["NPD"],
     "Go": ["NPD"],

src/run_repoaudit.sh

@@ -3,10 +3,11 @@ set -euo pipefail
 IFS=$'\n\t'
 
 # --- Defaults ---
-LANGUAGE="Python"
-MODEL="claude-3.7"
+LANGUAGE="Cpp"
+MODEL="deepseek-chat"
+# MODEL="claude-3.7"
 DEFAULT_PROJECT_NAME="toy"
-DEFAULT_BUG_TYPE="NPD"     # allowed: MLK, NPD, UAF
+DEFAULT_BUG_TYPE="RACE"     # allowed: MLK, NPD, UAF, RACE
 SCAN_TYPE="dfbscan"
 
 # Construct the default project *path* from LANGUAGE + DEFAULT_PROJECT_NAME
@@ -19,12 +20,13 @@ Usage: run_scan.sh [PROJECT_PATH] [BUG_TYPE]
 Arguments:
   PROJECT_PATH   Optional absolute/relative path to the subject project.
                  Defaults to: ../benchmark/Python/toy
-  BUG_TYPE       Optional bug type. One of: MLK, NPD, UAF. Defaults to: NPD
+  BUG_TYPE       Optional bug type. One of: MLK, NPD, UAF, RACE. Defaults to: NPD
 
 Bug type meanings:
   MLK  - Memory Leak
   NPD  - Null Pointer Dereference
   UAF  - Use After Free
+  RACE - Race Condition
 
 Examples:
   ./run_scan.sh
@@ -48,10 +50,10 @@ BUG_TYPE="$(echo "$BUG_TYPE_RAW" | tr '[:lower:]' '[:upper:]')"
 
 # --- Validate BUG_TYPE ---
 case "$BUG_TYPE" in
-  MLK|NPD|UAF) : ;;
+  MLK|NPD|UAF|RACE) : ;;
   *)
-    echo "Error: BUG_TYPE must be one of: MLK, NPD, UAF (got '$BUG_TYPE_RAW')." >&2
-    echo "       MLK = Memory Leak; NPD = Null Pointer Dereference; UAF = Use After Free." >&2
+    echo "Error: BUG_TYPE must be one of: MLK, NPD, UAF, RACE (got '$BUG_TYPE_RAW')." >&2
+    echo "       MLK = Memory Leak; NPD = Null Pointer Dereference; UAF = Use After Free; RACE = Race Condition." >&2
     exit 1
     ;;
 esac