Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 17 additions & 13 deletions providers/ssh/src/airflow/providers/ssh/utils/remote_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import base64
import re
import secrets
import shlex
import string
from dataclasses import dataclass
from typing import Literal
Expand Down Expand Up @@ -201,12 +202,12 @@ def build_posix_wrapper_command(
# ``setsid`` (some macOS/BSD hosts) ``$$`` is just the job's own PID and
# cancellation degrades to the previous single-process behaviour.
wrapper = f"""set -euo pipefail
job_dir='{paths.job_dir}'
log_file='{paths.log_file}'
exit_code_file='{paths.exit_code_file}'
exit_code_tmp='{paths.exit_code_tmp_file}'
pid_file='{paths.pid_file}'
status_file='{paths.status_file}'
job_dir={shlex.quote(paths.job_dir)}
log_file={shlex.quote(paths.log_file)}
exit_code_file={shlex.quote(paths.exit_code_file)}
exit_code_tmp={shlex.quote(paths.exit_code_tmp_file)}
pid_file={shlex.quote(paths.pid_file)}
status_file={shlex.quote(paths.status_file)}

mkdir -p "$job_dir"
: > "$log_file"
Expand All @@ -228,7 +229,7 @@ def build_posix_wrapper_command(
else
nohup bash -c "$job_script" >/dev/null 2>&1 &
fi
echo "{paths.job_id}"
echo {shlex.quote(paths.job_id)}
"""
return wrapper

Expand Down Expand Up @@ -312,7 +313,7 @@ def build_posix_log_tail_command(log_file: str, offset: int, max_bytes: int) ->
"""
# tail -c +N is 1-indexed, so offset 0 means start at byte 1
tail_offset = offset + 1
return f"tail -c +{tail_offset} '{log_file}' 2>/dev/null | head -c {max_bytes} || true"
return f"tail -c +{tail_offset} {shlex.quote(log_file)} 2>/dev/null | head -c {max_bytes} || true"


def build_windows_log_tail_command(log_file: str, offset: int, max_bytes: int) -> str:
Expand Down Expand Up @@ -350,7 +351,8 @@ def build_posix_file_size_command(file_path: str) -> str:
:param file_path: Path to the file
:return: Shell command that outputs the file size
"""
return f"stat -c%s '{file_path}' 2>/dev/null || stat -f%z '{file_path}' 2>/dev/null || echo 0"
quoted = shlex.quote(file_path)
return f"stat -c%s {quoted} 2>/dev/null || stat -f%z {quoted} 2>/dev/null || echo 0"


def build_windows_file_size_command(file_path: str) -> str:
Expand Down Expand Up @@ -379,7 +381,8 @@ def build_posix_completion_check_command(exit_code_file: str) -> str:
:param exit_code_file: Path to the exit code file
:return: Shell command that outputs exit code if done, empty otherwise
"""
return f"test -s '{exit_code_file}' && cat '{exit_code_file}' || true"
quoted = shlex.quote(exit_code_file)
return f"test -s {quoted} && cat {quoted} || true"


def build_windows_completion_check_command(exit_code_file: str) -> str:
Expand Down Expand Up @@ -419,9 +422,10 @@ def build_posix_kill_command(pid_file: str) -> str:
:param pid_file: Path to the PID file
:return: Shell command to kill the process
"""
quoted = shlex.quote(pid_file)
return (
f"if test -f '{pid_file}'; then "
f"p=\"$(cat '{pid_file}')\"; "
f"if test -f {quoted}; then "
f'p="$(cat {quoted})"; '
f'if [ "$p" -gt 1 ] 2>/dev/null; then '
f'kill -TERM -"$p" 2>/dev/null || kill -TERM "$p" 2>/dev/null || true; '
"fi; fi"
Expand Down Expand Up @@ -461,7 +465,7 @@ def build_posix_cleanup_command(job_dir: str) -> str:
:raises ValueError: If job_dir is not under the expected base directory
"""
_validate_job_dir(job_dir, "posix")
return f"rm -rf '{job_dir}'"
return f"rm -rf {shlex.quote(job_dir)}"


def build_windows_cleanup_command(job_dir: str) -> str:
Expand Down
32 changes: 31 additions & 1 deletion providers/ssh/tests/unit/ssh/utils/test_remote_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ class TestKillCommands:
def test_posix_kill_signals_process_group_then_falls_back(self):
"""POSIX kill targets the process group first, then a single PID as fallback."""
cmd = build_posix_kill_command("/tmp/pid")
assert "cat '/tmp/pid'" in cmd
assert "cat /tmp/pid" in cmd
# Negative PID => signal the whole process group (kills the job's children too)
assert 'kill -TERM -"$p"' in cmd
# Fallback for jobs that are not group leaders (host without setsid)
Expand Down Expand Up @@ -413,3 +413,33 @@ def test_windows_cleanup_rejects_invalid_path(self):
"""Test Windows cleanup rejects paths outside expected base directory."""
with pytest.raises(ValueError, match="Invalid job directory"):
build_windows_cleanup_command("C:\\temp\\other_dir")


class TestPosixPathQuoting:
"""A shell metacharacter in remote_base_dir must stay data, never become a command."""

@staticmethod
def _builders(paths):
return {
"wrapper": lambda: build_posix_wrapper_command("true", paths),
"cleanup": lambda: build_posix_cleanup_command(paths.job_dir),
"kill": lambda: build_posix_kill_command(paths.pid_file),
"log_tail": lambda: build_posix_log_tail_command(paths.log_file, 0, 64),
"file_size": lambda: build_posix_file_size_command(paths.log_file),
"completion": lambda: build_posix_completion_check_command(paths.exit_code_file),
}

@pytest.mark.parametrize(
"builder",
["wrapper", "cleanup", "kill", "log_tail", "file_size", "completion"],
)
def test_single_quote_in_base_dir_does_not_execute(self, builder, tmp_path):
marker = tmp_path / "injected"
# The prefix keeps job_dir under POSIX_DEFAULT_BASE_DIR so _validate_job_dir passes.
base_dir = f"/tmp/airflow-ssh-jobs/x'; touch {marker}; :'"
paths = RemoteJobPaths(job_id="job_123", remote_os="posix", base_dir=base_dir)

cmd = self._builders(paths)[builder]()
subprocess.run(["sh", "-c", cmd], capture_output=True, check=False)

assert not marker.exists()