Refactor path validation and error handling in discovery functions

- Consolidated path validation logic to prevent traversal attacks.
- Improved error handling in subprocess calls for test discovery, preserving command arguments for better context.
- Updated Windows-specific wait times for file release during cleanup.
- Enhanced readability and maintainability of path resolution functions.

Author: mohamedashrraf222

codeflash/code_utils/git_worktree_utils.py

@@ -8,10 +8,7 @@
 import tempfile
 import time
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
-
-if TYPE_CHECKING:
-    from collections.abc import Callable
+from typing import Any, Callable, Optional
 
 import git
 
@@ -24,6 +21,8 @@
 # Constants for Windows retry logic
 MAX_WINDOWS_RETRIES = 3
 INITIAL_RETRY_DELAY_SECONDS = 0.5
+WINDOWS_FILE_RELEASE_WAIT_LONG = 0.3
+WINDOWS_FILE_RELEASE_WAIT_SHORT = 0.1
 
 
 def create_worktree_snapshot_commit(worktree_dir: Path, commit_message: str) -> None:
@@ -212,14 +211,18 @@ def _manual_cleanup_worktree_directory(worktree_dir: Path) -> None:
 
             # Brief wait on Windows to allow file handles to be released
             if is_windows:
-                wait_time = 0.3 if attempt_num < max_retries else 0.1
+                wait_time = (
+                    WINDOWS_FILE_RELEASE_WAIT_LONG if attempt_num < max_retries else WINDOWS_FILE_RELEASE_WAIT_SHORT
+                )
                 time.sleep(wait_time)
 
             # Check if removal was successful
             if not worktree_dir.exists():
                 return
 
-        except Exception:  # noqa: S110
+        except (OSError, PermissionError):
+            # Silently ignore permission errors during cleanup attempts
+            # These are expected on Windows when files are locked
             pass
 
 
@@ -251,9 +254,11 @@ def _validate_worktree_path_safety(worktree_dir: Path) -> bool:
     return worktree_dir_resolved != worktree_dirs_resolved
 
 
-def _create_windows_rmtree_error_handler() -> Callable[
-    [Callable[[str], None], str, tuple[type[BaseException], BaseException, Any]], None
-]:
+# Type alias for shutil.rmtree error handler signature
+ErrorHandler = Callable[[Callable[[str], None], str, tuple[type[BaseException], BaseException, Any]], None]
+
+
+def _create_windows_rmtree_error_handler() -> ErrorHandler:
     """Create an error handler for shutil.rmtree that handles Windows-specific issues.
 
     This handler attempts to remove read-only attributes when encountering permission errors.

codeflash/discovery/discover_unit_tests.py

@@ -599,26 +599,25 @@ def discover_tests_pytest(
         # Prevent console window spawning on Windows which can cause hangs in LSP
         run_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
 
+    cmd_list = [SAFE_SYS_EXECUTABLE, discovery_script, str(project_root), str(tests_root), str(tmp_pickle_path)]
     with custom_addopts():
         try:
-            result = subprocess.run(
-                [SAFE_SYS_EXECUTABLE, discovery_script, str(project_root), str(tests_root), str(tmp_pickle_path)],
-                check=False,
-                **run_kwargs,
-            )
+            result = subprocess.run(cmd_list, check=False, **run_kwargs)
         except subprocess.TimeoutExpired:
             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")
+            # Preserve command args for better error context
+            result = subprocess.CompletedProcess(args=cmd_list, returncode=-1, stdout="", stderr="Timeout")
         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))
+            # Preserve command args for better error context
+            result = subprocess.CompletedProcess(args=cmd_list, returncode=-1, stdout="", stderr=str(e))
 
     try:
         # Check if pickle file exists before trying to read it

codeflash/discovery/functions_to_optimize.py

@@ -679,22 +679,17 @@ def _validate_path_no_traversal(path: Path | str) -> bool:
             True if path is safe (no traversal components), False otherwise
 
         """
-        path_str = str(path)
-        # Check for path traversal attempts
-        # Check for absolute paths that might escape (additional safety check)
-        # Note: We allow absolute paths as they're needed for worktree paths
-        return ".." not in path_str
-
-    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):
-            error_msg = f"Path contains traversal components: {path}"
-            raise ValueError(error_msg)
-        return Path(path).resolve(strict=False)
+        try:
+            path_obj = Path(path)
+        except (ValueError, OSError):
+            return False
+        else:
+            # Check if any part is exactly '..' (not just containing '..' as a substring)
+            # This avoids false positives like files named "config..bak"
+            return ".." not in path_obj.parts
 
-    def _resolve_path_consistent(path: Path | str) -> Path:
-        """Resolve path consistently: use strict resolution if path exists, otherwise non-strict.
+    def _resolve_path_safely(path: Path | str) -> Path:
+        """Resolve path, preferring strict resolution but falling back to non-strict.
 
         SECURITY: This function validates paths to prevent traversal attacks before resolution.
         Paths should come from trusted sources (git operations, file system discovery),
@@ -716,42 +711,34 @@ def _resolve_path_consistent(path: Path | str) -> Path:
             raise ValueError(error_msg)
 
         path_obj = Path(path)
-        if path_obj.exists():
-            try:
-                return path_obj.resolve()
-            except (OSError, RuntimeError):
-                return _resolve_path(path)
-        return _resolve_path(path)
+        try:
+            # Prefer strict resolution if path exists, otherwise use non-strict
+            return path_obj.resolve(strict=True) if path_obj.exists() else path_obj.resolve(strict=False)
+        except (OSError, RuntimeError):
+            # Fallback to non-strict resolution on errors
+            return path_obj.resolve(strict=False)
 
     # Resolve all root paths to absolute paths for consistent comparison
     # Use consistent resolution: strict for existing paths, non-strict for non-existent
-    tests_root_resolved = _resolve_path_consistent(tests_root)
-    module_root_resolved = _resolve_path_consistent(module_root)
+    tests_root_resolved = _resolve_path_safely(tests_root)
+    module_root_resolved = _resolve_path_safely(module_root)
 
     # Resolve ignore paths and submodule paths
-    ignore_paths_resolved = [_resolve_path_consistent(p) for p in ignore_paths]
-    submodule_paths_resolved = [_resolve_path_consistent(p) for p in submodule_paths]
+    ignore_paths_resolved = [_resolve_path_safely(p) for p in ignore_paths]
+    submodule_paths_resolved = [_resolve_path_safely(p) for p in submodule_paths]
 
     # 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
+        # Note: Paths come from trusted sources (git operations), but we validate defensively
         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)
-        # Use consistent resolution: try strict first (for existing files), fall back to non-strict
-        if file_path_obj.exists():
-            try:
-                file_path_resolved = file_path_obj.resolve()
-            except (OSError, RuntimeError):
-                file_path_resolved = _resolve_path(file_path_path)
-        else:
-            file_path_resolved = _resolve_path(file_path_path)
-
-        file_path = str(file_path_obj)
+            error_msg = f"Path contains traversal components: {file_path_path}"
+            logger.error(error_msg)
+            raise ValueError(error_msg)
+        # Resolve file path to absolute path using consolidated resolution logic
+        file_path_resolved = _resolve_path_safely(file_path_path)
+        file_path = str(file_path_path)  # Keep original path string for compatibility
 
         # Check if file is in tests root using resolved paths
         try: