302 lines
11 KiB
Python
302 lines
11 KiB
Python
"""调度器安全相关测试：SSRF 防护 + 单实例防护 + 统一线程池。"""
|
||
from modules.job import scheduler as sch
|
||
|
||
|
||
def test_build_full_url():
|
||
assert sch.build_full_url('/api/x') == f"http://{sch.FASTAPI_HOST}:{sch.FASTAPI_PORT}/api/x"
|
||
assert sch.build_full_url('http://a.com/x') == 'http://a.com/x'
|
||
assert sch.build_full_url('ftp://a.com/x') == 'ftp://a.com/x'
|
||
|
||
|
||
def test_validate_http_url_rejects_internal_ips():
|
||
for url in (
|
||
'http://127.0.0.1:6379/', # 回环
|
||
'http://10.0.0.1/x', # 私有A类
|
||
'http://169.254.169.254/latest', # 云元数据
|
||
'http://192.168.1.1/x', # 私有C类
|
||
'http://172.16.5.5/x', # 私有B类
|
||
):
|
||
ok, reason = sch._validate_http_url(url)
|
||
assert ok is False, url
|
||
assert '禁止' in reason or '内网' in reason
|
||
|
||
|
||
def test_validate_http_url_rejects_non_http_scheme():
|
||
for url in ('ftp://a.com/x', 'file:///etc/passwd', 'javascript:alert(1)'):
|
||
ok, reason = sch._validate_http_url(url)
|
||
assert ok is False, url
|
||
assert '协议' in reason
|
||
|
||
|
||
def test_validate_http_url_accepts_public_ip():
|
||
# 公网 IP 字面量，无需真实 DNS
|
||
ok, reason = sch._validate_http_url('http://93.184.216.34/')
|
||
assert ok is True, reason
|
||
|
||
|
||
def test_execute_by_url_blocks_internal_url(monkeypatch):
|
||
sent = []
|
||
monkeypatch.setattr(
|
||
sch, '_send_sync_request',
|
||
lambda url, data, jid, pinned_ip=None: sent.append(url) or {'success': True})
|
||
result = sch._execute_by_url(1, 'http://127.0.0.1:6379', 1)
|
||
assert result['success'] is False
|
||
assert '禁止' in result['error'] or '内网' in result['error']
|
||
assert sent == [] # 未实际发起请求
|
||
|
||
|
||
def test_execute_by_url_relative_self_call_allowed(monkeypatch):
|
||
sent = []
|
||
monkeypatch.setattr(
|
||
sch, '_send_sync_request',
|
||
lambda url, data, jid, pinned_ip=None: sent.append(url) or {'success': True})
|
||
result = sch._execute_by_url(1, '/api/ping', 1)
|
||
assert result['success'] is True
|
||
assert sent == [f"http://{sch.FASTAPI_HOST}:{sch.FASTAPI_PORT}/api/ping"]
|
||
|
||
|
||
# ============================================================
|
||
# SSRF DNS 重绑定防护：预解析后固定 IP 直连，杜绝二次解析域名
|
||
# ============================================================
|
||
|
||
def test_resolve_and_pin_returns_pinned_ip():
|
||
# 公网 IP 字面量，无需真实 DNS
|
||
ok, reason, ip = sch._resolve_and_pin('http://93.184.216.34/x')
|
||
assert ok is True, reason
|
||
assert ip == '93.184.216.34'
|
||
|
||
|
||
def test_resolve_and_pin_rejects_custom_port():
|
||
ok, reason, ip = sch._resolve_and_pin('http://93.184.216.34:8080/x')
|
||
assert ok is False
|
||
assert '端口' in reason
|
||
assert ip is None
|
||
|
||
|
||
def test_resolve_and_pin_allows_80_443():
|
||
for url in ('http://93.184.216.34:80/x', 'https://93.184.216.34:443/x'):
|
||
ok, reason, ip = sch._resolve_and_pin(url)
|
||
assert ok is True, reason
|
||
assert ip == '93.184.216.34'
|
||
|
||
|
||
def test_pinned_connection_pins_dns_host_keeps_real_host():
|
||
# HTTP 与 HTTPS 连接：_dns_host 固定为预校验 IP（实际建连目标），
|
||
# host 保持真实主机名（Host 头 / TLS SNI）——同域名只解析一次，杜绝重绑定。
|
||
for cls, port in ((sch._PinnedHTTPConnection, 80), (sch._PinnedHTTPSConnection, 443)):
|
||
conn = cls(host='example.com', port=port, pinned_ip='93.184.216.34')
|
||
assert conn._dns_host == '93.184.216.34'
|
||
assert conn.host == 'example.com'
|
||
|
||
|
||
def test_send_sync_request_uses_pinned_ip_when_given(monkeypatch):
|
||
calls = []
|
||
|
||
class _FakeResp:
|
||
status_code = 200
|
||
|
||
def raise_for_status(self):
|
||
pass
|
||
|
||
def json(self):
|
||
return {'ok': 1}
|
||
|
||
monkeypatch.setattr(
|
||
sch, '_post_pinned',
|
||
lambda url, data, headers, timeout, pinned_ip: calls.append((url, pinned_ip)) or _FakeResp())
|
||
result = sch._send_sync_request('http://example.com/x', {'a': 1}, 1, '93.184.216.34')
|
||
assert result['success'] is True
|
||
assert calls == [('http://example.com/x', '93.184.216.34')]
|
||
|
||
|
||
def test_send_sync_request_without_pinned_ip_uses_plain_post(monkeypatch):
|
||
calls = []
|
||
|
||
class _FakeResp:
|
||
status_code = 200
|
||
|
||
def raise_for_status(self):
|
||
pass
|
||
|
||
def json(self):
|
||
return {'ok': 1}
|
||
|
||
monkeypatch.setattr(
|
||
sch, '_post_pinned',
|
||
lambda *a, **k: calls.append('pinned') or _FakeResp())
|
||
monkeypatch.setattr(
|
||
sch.requests, 'post',
|
||
lambda url, **kw: calls.append('plain') or _FakeResp())
|
||
|
||
# 相对路径自调用等可信场景无固定 IP，应走常规 requests.post 而非 pinned 路径
|
||
result = sch._send_sync_request('http://127.0.0.1:8031/api/ping', {'a': 1}, 1)
|
||
assert result['success'] is True
|
||
assert calls == ['plain']
|
||
|
||
|
||
def test_job_single_instance_claim(monkeypatch):
|
||
# 禁用 Redis 锁（退化进程内互斥），使进程内 set 行为可独立验证
|
||
monkeypatch.setattr(sch, '_get_lock_redis', lambda: None)
|
||
assert sch._try_claim_job(42) is True
|
||
assert sch._try_claim_job(42) is False # 已占用，拒绝重入
|
||
sch._release_job(42)
|
||
assert sch._try_claim_job(42) is True
|
||
sch._release_job(42)
|
||
|
||
|
||
# ============================================================
|
||
# Redis 分布式锁：跨 worker 进程单实例互斥
|
||
# ============================================================
|
||
|
||
def test_try_claim_job_respects_cross_process_redis_lock(monkeypatch):
|
||
import fakeredis
|
||
|
||
fake = fakeredis.FakeRedis()
|
||
monkeypatch.setattr(sch, '_get_lock_redis', lambda: fake)
|
||
|
||
# 进程内重复触发仍被本地 set 拦截
|
||
assert sch._try_claim_job(7) is True
|
||
assert sch._try_claim_job(7) is False
|
||
sch._release_job(7)
|
||
assert sch._try_claim_job(7) is True
|
||
sch._release_job(7)
|
||
|
||
# 模拟他进程已持锁：Redis 锁被占用 → 本进程拒绝，且本地占用已回滚
|
||
fake.set(f'{sch._JOB_LOCK_PREFIX}9', '1', nx=True, ex=300)
|
||
assert sch._try_claim_job(9) is False
|
||
assert 9 not in sch._running_jobs
|
||
# 释放后（他进程删除锁）可再次占用
|
||
fake.delete(f'{sch._JOB_LOCK_PREFIX}9')
|
||
assert sch._try_claim_job(9) is True
|
||
sch._release_job(9)
|
||
|
||
|
||
def test_try_claim_job_degrades_when_redis_unavailable(monkeypatch):
|
||
# Redis 不可用 → 退化进程内互斥，不中断任务执行
|
||
monkeypatch.setattr(sch, '_get_lock_redis', lambda: None)
|
||
assert sch._try_claim_job(10) is True
|
||
assert sch._try_claim_job(10) is False # 本地 set 仍生效
|
||
sch._release_job(10)
|
||
|
||
|
||
# ============================================================
|
||
# _handle_scheduler_update：运行中任务编辑后必须重加调度器刷新执行参数
|
||
# ============================================================
|
||
|
||
def _patch_scheduler_calls(monkeypatch):
|
||
"""打桩 job.service 中的调度器函数，记录 add/remove 调用。"""
|
||
from modules.job import service as job_svc
|
||
calls = {'add': [], 'remove': []}
|
||
monkeypatch.setattr(
|
||
job_svc, 'add_job_to_scheduler',
|
||
lambda data: calls['add'].append(data) or True)
|
||
monkeypatch.setattr(
|
||
job_svc, 'remove_job_from_scheduler',
|
||
lambda jid: calls['remove'].append(jid) or True)
|
||
return job_svc, calls
|
||
|
||
|
||
def test_handle_scheduler_update_running_edit_readds(monkeypatch):
|
||
"""运行中任务编辑 url/params：cron 未变也需重加，否则调度器仍按旧参数执行。"""
|
||
job_svc, calls = _patch_scheduler_calls(monkeypatch)
|
||
ok = job_svc._handle_scheduler_update(
|
||
{'id': 5, 'status': 1, 'url': '/api/new'}, old_status=1, old_cron='0 * * * *')
|
||
assert ok is True
|
||
assert calls['remove'] == []
|
||
assert len(calls['add']) == 1
|
||
# 重加时携带完整任务状态与有效 cron（部分更新缺 cron 时沿用旧值）
|
||
added = calls['add'][0]
|
||
assert added['id'] == 5
|
||
assert added['url'] == '/api/new'
|
||
assert added['cron_expression'] == '0 * * * *'
|
||
|
||
|
||
def test_handle_scheduler_update_running_explicit_new_cron(monkeypatch):
|
||
"""运行中任务显式改 cron：以新 cron 重加。"""
|
||
job_svc, calls = _patch_scheduler_calls(monkeypatch)
|
||
ok = job_svc._handle_scheduler_update(
|
||
{'id': 5, 'status': 1, 'cron_expression': '*/5 * * * *'}, old_status=1, old_cron='0 * * * *')
|
||
assert ok is True
|
||
assert calls['add'][0]['cron_expression'] == '*/5 * * * *'
|
||
|
||
|
||
def test_handle_scheduler_update_running_to_disabled_removes(monkeypatch):
|
||
"""运行中 → 禁用：从调度器移除。"""
|
||
job_svc, calls = _patch_scheduler_calls(monkeypatch)
|
||
ok = job_svc._handle_scheduler_update(
|
||
{'id': 5, 'status': 2}, old_status=1, old_cron='0 * * * *')
|
||
assert ok is True
|
||
assert calls['add'] == []
|
||
assert calls['remove'] == [5]
|
||
|
||
|
||
def test_handle_scheduler_update_disabled_to_running_readds(monkeypatch):
|
||
"""暂停 → 启用且带 cron：重加。"""
|
||
job_svc, calls = _patch_scheduler_calls(monkeypatch)
|
||
ok = job_svc._handle_scheduler_update(
|
||
{'id': 5, 'status': 1, 'cron_expression': '0 * * * *'}, old_status=2, old_cron='0 * * * *')
|
||
assert ok is True
|
||
assert len(calls['add']) == 1
|
||
assert calls['remove'] == []
|
||
|
||
|
||
def test_handle_scheduler_update_disabled_to_running_without_cron(monkeypatch):
|
||
"""暂停 → 启用但无 cron：无法调度，移除（保持调度器一致性）。"""
|
||
job_svc, calls = _patch_scheduler_calls(monkeypatch)
|
||
ok = job_svc._handle_scheduler_update(
|
||
{'id': 5, 'status': 1}, old_status=2, old_cron=None)
|
||
assert ok is True
|
||
assert calls['add'] == []
|
||
assert calls['remove'] == [5]
|
||
|
||
|
||
def test_handle_scheduler_update_no_status_change_noop(monkeypatch):
|
||
"""保持暂停：不触发任何调度器操作。"""
|
||
job_svc, calls = _patch_scheduler_calls(monkeypatch)
|
||
ok = job_svc._handle_scheduler_update(
|
||
{'id': 5, 'status': 2}, old_status=2, old_cron=None)
|
||
assert ok is True
|
||
assert calls['add'] == []
|
||
assert calls['remove'] == []
|
||
|
||
|
||
# ============================================================
|
||
# _build_cron_trigger：5/6 字段 cron，6 字段首字段为秒（秒级触发）
|
||
# ============================================================
|
||
|
||
def _fire_seconds_away(cron_expr):
|
||
"""返回 (next_fire, 距 now 的秒数)。now 取本地 aware 时间，与触发器默认时区对齐。"""
|
||
from datetime import datetime
|
||
now = datetime.now().astimezone()
|
||
nf = sch._build_cron_trigger(cron_expr).get_next_fire_time(None, now)
|
||
return nf, (nf - now).total_seconds() if nf else None
|
||
|
||
|
||
def test_build_cron_trigger_5_fields_minute_granularity():
|
||
nf, delta = _fire_seconds_away('0 * * * *') # 每小时整点
|
||
assert nf is not None
|
||
assert 0 < delta <= 3600
|
||
|
||
|
||
def test_build_cron_trigger_6_fields_second_granularity():
|
||
nf, delta = _fire_seconds_away('*/15 * * * * *') # 每 15 秒
|
||
assert nf is not None
|
||
assert 0 < delta <= 15, f"秒级 cron 应 15s 内触发，实际 {delta}s"
|
||
|
||
|
||
def test_build_cron_trigger_5_field_matches_from_crontab():
|
||
"""5 字段行为与 CronTrigger.from_crontab 完全一致（向后兼容）。"""
|
||
from apscheduler.triggers.cron import CronTrigger
|
||
from datetime import datetime
|
||
now = datetime.now().astimezone()
|
||
t1 = sch._build_cron_trigger('0 * * * *').get_next_fire_time(None, now)
|
||
t2 = CronTrigger.from_crontab('0 * * * *').get_next_fire_time(None, now)
|
||
assert t1 == t2
|
||
|
||
|
||
def test_build_cron_trigger_invalid_field_count_raises():
|
||
import pytest
|
||
for bad in ('', '* * * *', '* * * * * * *'): # 0/4/7 字段均非法
|
||
with pytest.raises(ValueError):
|
||
sch._build_cron_trigger(bad)
|