[Misc] Refactor: Improve argument handling for conda command
#20481
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
There was a problem hiding this comment.
Summary of Changes
Hello @reidliu41, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request focuses on a critical security improvement by addressing potential command injection vulnerabilities in the codebase. It refactors how external commands are executed, moving away from insecure shell-dependent execution and introducing proactive measures to block dangerous commands, thereby enhancing the overall robustness and safety of the system's command execution utility.
Highlights
- Security Enhancement: Eliminated the use of
shell=Trueinsubprocess.Popencalls within therunfunction. This change significantly mitigates command injection vulnerabilities by preventing the shell from interpreting untrusted input as commands. - Safe Command Parsing: Implemented
shlex.splitto safely parse string commands into a list format. This ensures that commands are executed securely as a list of arguments, rather than relying on the shell to interpret a single command string, which is a common source of injection risks. - Command Blacklisting: Introduced a
COMMAND_BLACKLISTand anis_dangerous_cmdhelper function. Therunfunction now checks if the command being executed is on this blacklist (e.g.,rm,wget,bash) and raises aRuntimeErrorto prevent the execution of potentially dangerous system commands. - Robust Error Handling: Added specific error handling for
FileNotFoundErrorwithin therunfunction. If a command executable is not found, it now returns a distinct error code (127) and a clear error message, improving diagnostic capabilities.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request removes shell=True and adds a command blacklist, which is a great step forward in securing the environment collection script. My review highlights Python 2 compatibility issues that were introduced, which would break the script in a Python 2 environment. I've also pointed out a logic bug and a potential bypass in the security blacklist. Addressing these points will ensure the changes are both robust and correct.
vllm/collect_env.py
Outdated
| command = shlex.split(command) | ||
|
|
||
| if is_dangerous_cmd(command[0]): | ||
| raise RuntimeError(f"Blocked dangerous command: {command[0]}") |
There was a problem hiding this comment.
This file has a comment stating it must be Python 2 compliant. The f-string used here was introduced in Python 3.6 and will cause a SyntaxError when run with a Python 2 interpreter. Please use a different string formatting method, such as str.format() or the % operator, to maintain compatibility.
| raise RuntimeError(f"Blocked dangerous command: {command[0]}") | |
| raise RuntimeError("Blocked dangerous command: {}".format(command[0])) |
vllm/collect_env.py
Outdated
| except FileNotFoundError: | ||
| return 127, '', f"Command not found: {command[0]}" |
There was a problem hiding this comment.
This code introduces two Python 2 compatibility issues:
FileNotFoundErrorwas introduced in Python 3.3 and does not exist in Python 2. This will raise aNameError.- The f-string on line 148 is not supported in Python 2.
To fix this, you should catch OSError and check its errno for errno.ENOENT to handle a missing command, and use a compatible string formatting method.
| except FileNotFoundError: | |
| return 127, '', f"Command not found: {command[0]}" | |
| except OSError as e: | |
| import errno | |
| if e.errno == errno.ENOENT: | |
| return 127, '', "Command not found: {}".format(command[0]) | |
| raise |
vllm/collect_env.py
Outdated
| def is_dangerous_cmd(cmd): | ||
| cmd_base = os.path.basename(cmd) | ||
| return cmd_base in COMMAND_BLACKLIST |
There was a problem hiding this comment.
The current implementation of is_dangerous_cmd can be bypassed using symbolic links. It checks os.path.basename(cmd), but an attacker could create a symlink (e.g., ln -s /bin/rm my_command) which would not be caught by the blacklist but would execute the dangerous command.
| else: | ||
| enc = locale.getpreferredencoding() | ||
| output = raw_output.decode(enc) | ||
| if command == 'nvidia-smi topo -m': |
There was a problem hiding this comment.
After splitting the command string with shlex.split, the command variable becomes a list of strings. This comparison, command == 'nvidia-smi topo -m', will always evaluate to False, breaking the special handling for this command's output.
You should compare against the list representation of the command instead.
| if command == 'nvidia-smi topo -m': | |
| if command == ['nvidia-smi', 'topo', '-m']: |
vllm/collect_env.py
Outdated
| COMMAND_BLACKLIST = [ | ||
| 'rm', | ||
| 'wget', | ||
| 'curl', | ||
| 'nc', | ||
| 'bash', | ||
| 'sh', | ||
| 'mkfs', | ||
| 'dd', | ||
| 'scp', | ||
| 'ftp', | ||
| 'poweroff', | ||
| 'reboot', | ||
| 'init', | ||
| ] |
There was a problem hiding this comment.
For performance, it's better to define COMMAND_BLACKLIST as a set instead of a list. Membership testing (in) is on average O(1) for sets, while it's O(n) for lists.
| COMMAND_BLACKLIST = [ | |
| 'rm', | |
| 'wget', | |
| 'curl', | |
| 'nc', | |
| 'bash', | |
| 'sh', | |
| 'mkfs', | |
| 'dd', | |
| 'scp', | |
| 'ftp', | |
| 'poweroff', | |
| 'reboot', | |
| 'init', | |
| ] | |
| COMMAND_BLACKLIST = { | |
| 'rm', | |
| 'wget', | |
| 'curl', | |
| 'nc', | |
| 'bash', | |
| 'sh', | |
| 'mkfs', | |
| 'dd', | |
| 'scp', | |
| 'ftp', | |
| 'poweroff', | |
| 'reboot', | |
| 'init', | |
| } |
There was a problem hiding this comment.
LGTM but I'll have @youkaichao @mgoin take another look just to be sure
|
Hi team, could someone else help take a look? |
There was a problem hiding this comment.
why do we need to take care of security here? who can inject arbitrary command in this file? this file follows https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py .
|
According to the PR description:
|
Improving the code is fine, but I don't consider this a security vulnerability. "if combined with untrusted input" -- I'm not sure why or when that would ever be the case? |
…plugins] [-V] COMMAND ...
conda is a tool for managing and deploying applications, environments and packages.
options:
-h, --help Show this help message and exit.
-v, --verbose Can be used multiple times. Once for detailed output,
twice for INFO logging, thrice for DEBUG logging, four
times for TRACE logging.
--no-plugins Disable all plugins that are not built into conda.
-V, --version Show the conda version number and exit.
commands:
The following built-in and plugins subcommands are available.
COMMAND
activate Activate a conda environment.
clean Remove unused packages and caches.
commands List all available conda subcommands (including those
from plugins). Generally only used by tab-completion.
compare Compare packages between conda environments.
config Modify configuration values in .condarc.
content-trust Signing and verification tools for Conda
create Create a new conda environment from a list of
specified packages.
deactivate Deactivate the current active conda environment.
doctor Display a health report for your environment.
export Export a given environment
info Display information about current conda install.
init Initialize conda for shell interaction.
install Install a list of packages into a specified conda
environment.
list List installed packages in a conda environment.
notices Retrieve latest channel notifications.
package Create low-level conda packages. (EXPERIMENTAL)
remove (uninstall) Remove a list of packages from a specified conda
environment.
rename Rename an existing environment.
repoquery Advanced search for repodata.
run Run an executable in a conda environment.
search Search for packages and display associated information
using the MatchSpec format.
tos A subcommand for viewing, accepting, rejecting, and
otherwise interacting with a channel's Terms of
Service (ToS). This plugin periodically checks for
updated Terms of Service for the active/selected
channels. Channels with a Terms of Service will need
to be accepted or rejected prior to use. Conda will
only allow package installation from channels without
a Terms of Service or with an accepted Terms of
Service. Attempting to use a channel with a rejected
Terms of Service will result in an error.
update (upgrade) Update conda packages to the latest compatible
version. command
Signed-off-by: reidliu41 <reid201711@gmail.com>
0a00a64 to
2274fd6
Compare
conda command
|
Thank you for your insightful feedback. It helped me refine my approach, and You've raised very valid points. After a deeper look, I agree this isn't a practical The new change is now surgically focused on the one case where an environment variable |
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Himanshu Jaju <hj@mistral.ai>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: avigny <47987522+avigny@users.noreply.github.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: x22x22 <wadeking@qq.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Paul Pak <paulpak58@gmail.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Diego-Castan <diego.castan@ibm.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com>
…project#20481) Signed-off-by: reidliu41 <reid201711@gmail.com>
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
Remove shell=True and add command blacklist to prevent command injection
The original code uses shell=True to support string commands, which may lead to command injection risks if combined with untrusted inputs or environment variables, e.g.
CONDA_EXE="conda; rm -rf /".Changes
Test Plan
Test Result
(Optional) Documentation Update