addressing all Required items by claude code review

- Path Traversal Protection
- Replace Bare Exception Handlers
- Test Discovery Logging

Author: mohamedashrraf222

codeflash/benchmarking/trace_benchmarks.py

@@ -10,6 +10,9 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
 
+# Constant for subprocess drain timeout after killing
+SUBPROCESS_DRAIN_TIMEOUT_SECONDS = 5
+
 
 def trace_benchmarks_pytest(
     benchmarks_root: Path, tests_root: Path, project_root: Path, trace_file: Path, timeout: int = 300
@@ -48,7 +51,14 @@ def trace_benchmarks_pytest(
         except subprocess.TimeoutExpired:
             with contextlib.suppress(OSError):
                 process.kill()
-            stdout_content, stderr_content = process.communicate(timeout=5)
+            try:
+                # Try to drain remaining output after killing
+                stdout_content, stderr_content = process.communicate(timeout=SUBPROCESS_DRAIN_TIMEOUT_SECONDS)
+            except subprocess.TimeoutExpired:
+                # Last resort: terminate and get partial output
+                with contextlib.suppress(OSError):
+                    process.terminate()
+                stdout_content, stderr_content = "", "Process killed after timeout"
             raise subprocess.TimeoutExpired(cmd_list, timeout, output=stdout_content, stderr=stderr_content) from None
         result = subprocess.CompletedProcess(cmd_list, returncode, stdout_content, stderr_content)
     else:

codeflash/code_utils/code_replacer.py

@@ -447,7 +447,7 @@ def replace_function_definitions_in_module(
 
     new_code: str = replace_functions_and_add_imports(
         # adding the global assignments before replacing the code, not after
-        # becuase of an "edge case" where the optimized code intoduced a new import and a global assignment using that import
+        # because of an "edge case" where the optimized code intoduced a new import and a global assignment using that import
         # and that import wasn't used before, so it was ignored when calling AddImportsVisitor.add_needed_import inside replace_functions_and_add_imports (because the global assignment wasn't added yet)
         # this was added at https://github.com/codeflash-ai/codeflash/pull/448
         add_global_assignments(code_to_apply, source_code) if should_add_global_assignments else source_code,

codeflash/code_utils/git_worktree_utils.py

@@ -21,6 +21,10 @@
 worktree_dirs = codeflash_cache_dir / "worktrees"
 patches_dir = codeflash_cache_dir / "patches"
 
+# Constants for Windows retry logic
+MAX_WINDOWS_RETRIES = 3
+INITIAL_RETRY_DELAY_SECONDS = 0.5
+
 
 def create_worktree_snapshot_commit(worktree_dir: Path, commit_message: str) -> None:
     repository = git.Repo(worktree_dir, search_parent_directories=True)
@@ -113,14 +117,17 @@ def remove_worktree(worktree_dir: Path) -> None:
         return
 
     is_windows = sys.platform == "win32"
-    max_retries = 3 if is_windows else 1
-    retry_delay = 0.5  # Start with 500ms delay
+    max_retries = MAX_WINDOWS_RETRIES if is_windows else 1
+    retry_delay = INITIAL_RETRY_DELAY_SECONDS
 
     # Try to get the repository and git root for worktree removal
     try:
         repository = git.Repo(worktree_dir, search_parent_directories=True)
-    except Exception:
+    except (git.exc.InvalidGitRepositoryError, OSError, PermissionError) as e:
         # If we can't access the repository, try manual cleanup
+        # Log at debug level since this is expected in some edge cases
+        from codeflash.cli_cmds.console import logger
+        logger.debug(f"Could not access git repository at {worktree_dir}: {e}. Attempting manual cleanup.")
         _manual_cleanup_worktree_directory(worktree_dir)
         return
 
@@ -140,17 +147,20 @@ def remove_worktree(worktree_dir: Path) -> None:
             else:
                 # Last attempt failed or non-permission error
                 break
-        except Exception:
+        except (OSError, PermissionError) as e:
+            # Log unexpected errors for debugging
+            from codeflash.cli_cmds.console import logger
+            logger.debug(f"Worktree removal attempt {attempt + 1} failed with unexpected error: {e}")
             break
 
     # Fallback: Try to remove worktree entry from git, then manually delete directory
-    with contextlib.suppress(Exception):
+    with contextlib.suppress(git.exc.GitCommandError, OSError, PermissionError):
         # Try to prune the worktree entry from git (this doesn't delete the directory)
         # Use git worktree prune to remove stale entries
         repository.git.worktree("prune")
 
     # Manually remove the directory (always attempt, even if prune failed)
-    with contextlib.suppress(Exception):
+    with contextlib.suppress(OSError, PermissionError):
         _manual_cleanup_worktree_directory(worktree_dir)
 
 
@@ -178,8 +188,8 @@ def _manual_cleanup_worktree_directory(worktree_dir: Path) -> None:
 
     # Attempt removal with retries on Windows
     is_windows = sys.platform == "win32"
-    max_retries = 3 if is_windows else 1
-    retry_delay = 0.5
+    max_retries = MAX_WINDOWS_RETRIES if is_windows else 1
+    retry_delay = INITIAL_RETRY_DELAY_SECONDS
 
     for attempt in range(max_retries):
         attempt_num = attempt + 1

codeflash/discovery/discover_unit_tests.py

@@ -331,6 +331,8 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
             # Be conservative except when an alias is used (which requires exact method matching)
             for target_func in fnames:
                 if "." in target_func:
+                    # Split to extract class name; method name is intentionally discarded (leading underscore)
+                    # as we only need to check if the imported class matches the target function's class
                     class_name, _method_name = target_func.split(".", 1)
                     if aname == class_name and not alias.asname:
                         self.found_any_target_function = True
@@ -604,23 +606,38 @@ def discover_tests_pytest(
                 check=False,
                 **run_kwargs,
             )
-        except subprocess.TimeoutExpired:
+        except subprocess.TimeoutExpired as e:
+            logger.error(
+                f"Test discovery subprocess timed out after {run_kwargs.get('timeout', 600)} seconds. "
+                f"Command: {discovery_script}"
+            )
             result = subprocess.CompletedProcess(args=[], returncode=-1, stdout="", stderr="Timeout")
-        except Exception as e:
+        except (OSError, subprocess.SubprocessError, ValueError) as e:
+            logger.error(
+                f"Test discovery subprocess failed with error: {e}. "
+                f"Command: {discovery_script}, "
+                f"Project root: {project_root}, Tests root: {tests_root}"
+            )
             result = subprocess.CompletedProcess(args=[], returncode=-1, stdout="", stderr=str(e))
 
     try:
         # Check if pickle file exists before trying to read it
         if not tmp_pickle_path.exists():
             tests, pytest_rootdir = [], None
-            logger.warning(
+            logger.error(
                 f"Test discovery pickle file not found. "
-                f"Subprocess return code: {result.returncode}, stdout: {result.stdout}, stderr: {result.stderr}"
+                f"Subprocess return code: {result.returncode}, stdout: {result.stdout[:500]}, stderr: {result.stderr[:500]}"
             )
             exitcode = result.returncode if result.returncode != 0 else -1
         else:
             with tmp_pickle_path.open(mode="rb") as f:
                 exitcode, tests, pytest_rootdir = pickle.load(f)
+            # Log error if subprocess failed even though pickle file exists
+            if exitcode != 0:
+                logger.error(
+                    f"Test discovery subprocess returned non-zero exit code: {exitcode}. "
+                    f"Subprocess return code: {result.returncode}, stdout: {result.stdout[:500]}, stderr: {result.stderr[:500]}"
+                )
     except Exception as e:
         tests, pytest_rootdir = [], None
         logger.exception(

codeflash/discovery/functions_to_optimize.py

@@ -665,12 +665,54 @@ def filter_functions(
     blocklist_funcs_removed_count: int = 0
     previous_checkpoint_functions_removed_count: int = 0
 
+    def _validate_path_no_traversal(path: Path | str) -> bool:
+        """Validate that a path does not contain path traversal components.
+        
+        This prevents path traversal attacks by rejecting paths with '..' components.
+        Paths passed to this function should be from trusted sources (git operations,
+        file system discovery), but we validate defensively.
+        
+        Args:
+            path: Path to validate
+            
+        Returns:
+            True if path is safe (no traversal components), False otherwise
+        """
+        path_str = str(path)
+        # Check for path traversal attempts
+        if ".." in path_str:
+            return False
+        # Check for absolute paths that might escape (additional safety check)
+        # Note: We allow absolute paths as they're needed for worktree paths
+        return True
+
     def _resolve_path(path: Path | str) -> Path:
         # Use strict=False so we don't fail on paths that don't exist yet (e.g. worktree paths)
+        # SECURITY: Validate path before resolution to prevent traversal attacks
+        if not _validate_path_no_traversal(path):
+            raise ValueError(f"Path contains traversal components: {path}")
         return Path(path).resolve(strict=False)
 
     def _resolve_path_consistent(path: Path | str) -> Path:
-        """Resolve path consistently: use strict resolution if path exists, otherwise non-strict."""
+        """Resolve path consistently: use strict resolution if path exists, otherwise non-strict.
+        
+        SECURITY: This function validates paths to prevent traversal attacks before resolution.
+        Paths should come from trusted sources (git operations, file system discovery),
+        but we validate defensively.
+        
+        Args:
+            path: Path to resolve (from trusted sources like git diff or file discovery)
+            
+        Returns:
+            Resolved absolute Path
+            
+        Raises:
+            ValueError: If path contains traversal components
+        """
+        # SECURITY: Validate path before any resolution to prevent traversal attacks
+        if not _validate_path_no_traversal(path):
+            raise ValueError(f"Path contains traversal components: {path}")
+            
         path_obj = Path(path)
         if path_obj.exists():
             try:
@@ -691,6 +733,10 @@ def _resolve_path_consistent(path: Path | str) -> Path:
     # We desperately need Python 3.10+ only support to make this code readable with structural pattern matching
     for file_path_path, functions in modified_functions.items():
         _functions = functions
+        # SECURITY: Validate file path before processing to prevent traversal attacks
+        if not _validate_path_no_traversal(file_path_path):
+            logger.warning(f"Skipping file with traversal components: {file_path_path}")
+            continue
         # Resolve file path to absolute path
         # Convert to Path if it's a string (e.g., from get_functions_within_git_diff)
         file_path_obj = Path(file_path_path)

codeflash/lsp/lsp_message.py

@@ -11,7 +11,7 @@
 json_primitive_types = (str, float, int, bool)
 max_code_lines_before_collapse = 45
 
-# \\u241F is the message delimiter becuase it can be more than one message sent over the same message, so we need something to separate each message
+# \\u241F is the message delimiter because it can be more than one message sent over the same message, so we need something to separate each message
 message_delimiter = "\\u241F"
 
 

codeflash/verification/test_runner.py

@@ -19,6 +19,9 @@
 BEHAVIORAL_BLOCKLISTED_PLUGINS = ["benchmark", "codspeed", "xdist", "sugar"]
 BENCHMARKING_BLOCKLISTED_PLUGINS = ["codspeed", "cov", "benchmark", "profiling", "xdist", "sugar"]
 
+# Constant for subprocess drain timeout after killing
+SUBPROCESS_DRAIN_TIMEOUT_SECONDS = 5
+
 
 def execute_test_subprocess(
     cmd_list: list[str], cwd: Path, env: dict[str, str] | None, timeout: int = 600
@@ -74,7 +77,13 @@ def execute_test_subprocess(
                     process.kill()
 
                 # Drain remaining output after killing
-                stdout_content, stderr_content = process.communicate(timeout=5)
+                try:
+                    stdout_content, stderr_content = process.communicate(timeout=SUBPROCESS_DRAIN_TIMEOUT_SECONDS)
+                except subprocess.TimeoutExpired:
+                    # Last resort: terminate and get partial output
+                    with contextlib.suppress(OSError):
+                        process.terminate()
+                    stdout_content, stderr_content = "", "Process killed after timeout"
                 raise subprocess.TimeoutExpired(
                     cmd_list, timeout, output=stdout_content, stderr=stderr_content
                 ) from None