P1: 架构加固

- 线程池服务器: 连接级读写超时 + 有界队列(BoundedSemaphore),拒绝慢速/洪水连接
- 镜像代理流式转发(64KB 分块 + 4GB 上限),小响应才内存缓存(16MB 上限)
- _serve_cached 支持单段 Range(206);pypi 下载流式化;docker urlopen 补超时
This commit is contained in:
HYC Fixer
2026-08-30 12:23:29 +08:00
parent 084164296d
commit 237b99b480
4 changed files with 588 additions and 501 deletions
+39 -6
View File
@@ -7,6 +7,7 @@ import os
import sys import sys
import ssl import ssl
import signal import signal
import socket
import mimetypes import mimetypes
import threading import threading
import socketserver import socketserver
@@ -18,29 +19,59 @@ from .mirror_sync import MirrorSyncManager
class ThreadPoolHTTPServer(socketserver.ThreadingMixIn, socketserver.TCPServer): class ThreadPoolHTTPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
"""使用线程池的 HTTP 服务器""" """使用线程池的 HTTP 服务器(有界队列 + 连接超时,防慢速 DoS)"""
allow_reuse_address = True allow_reuse_address = True
daemon_threads = True # 使用守护线程 daemon_threads = True # 使用守护线程
def __init__(self, server_address, RequestHandlerClass, max_workers=50): def __init__(self, server_address, RequestHandlerClass, max_workers=50,
queue_size=50, socket_timeout=30):
self.max_workers = max_workers self.max_workers = max_workers
self.queue_size = queue_size
self.socket_timeout = socket_timeout
self._executor = ThreadPoolExecutor( self._executor = ThreadPoolExecutor(
max_workers=max_workers, max_workers=max_workers,
thread_name_prefix="http_handler" thread_name_prefix="http_handler"
) )
# 有界并发槽位(活跃 + 排队),超过上限直接拒绝新连接,
# 防止连接洪水导致任务无限堆积在内存
self._slots = threading.BoundedSemaphore(max_workers + queue_size)
super().__init__(server_address, RequestHandlerClass) super().__init__(server_address, RequestHandlerClass)
def process_request(self, request, client_address): def process_request(self, request, client_address):
"""使用线程池处理请求""" """使用线程池处理请求(有界队列)"""
self._executor.submit(self._handle_request, request, client_address) if not self._slots.acquire(blocking=False):
# 队列已满:拒绝并立即关闭连接
try:
request.shutdown(socket.SHUT_RDWR)
except Exception:
pass
try:
request.close()
except Exception:
pass
return
try:
self._executor.submit(self._handle_request, request, client_address)
except Exception:
self._slots.release()
try:
request.close()
except Exception:
pass
def _handle_request(self, request, client_address): def _handle_request(self, request, client_address):
"""实际处理请求""" """实际处理请求"""
try: try:
# 连接级读写超时,防慢客户端无限占用工作线程
try:
request.settimeout(self.socket_timeout)
except Exception:
pass
self.finish_request(request, client_address) self.finish_request(request, client_address)
except Exception: except Exception:
self.handle_error(request, client_address) self.handle_error(request, client_address)
finally: finally:
self._slots.release()
self.shutdown_request(request) self.shutdown_request(request)
def server_close(self): def server_close(self):
@@ -118,7 +149,9 @@ class MirrorServer:
self.server = ThreadPoolHTTPServer( self.server = ThreadPoolHTTPServer(
server_address, server_address,
MirrorServerHandler, MirrorServerHandler,
max_workers=max_workers max_workers=max_workers,
queue_size=max_workers * 5,
socket_timeout=self.config.get('timeout', 30)
) )
# 传递配置到处理器 # 传递配置到处理器
@@ -128,7 +161,7 @@ class MirrorServer:
# 设置调试模式 # 设置调试模式
MirrorServerHandler._setup_debug(self.config) MirrorServerHandler._setup_debug(self.config)
# 设置超时 # 设置 accept 轮询超时
self.server.timeout = self.config.get('timeout', 30) self.server.timeout = self.config.get('timeout', 30)
# 启用HTTPS # 启用HTTPS
+2 -2
View File
@@ -141,7 +141,7 @@ class DockerMirror:
auth = base64.b64encode(f"{self.username}:{self.password}".encode()).decode() auth = base64.b64encode(f"{self.username}:{self.password}".encode()).decode()
req.add_header('Authorization', f"Basic {auth}") req.add_header('Authorization', f"Basic {auth}")
with urllib.request.urlopen(req) as response: with urllib.request.urlopen(req, timeout=60) as response:
data = response.read() data = response.read()
# 缓存 # 缓存
@@ -189,7 +189,7 @@ class DockerMirror:
auth = base64.b64encode(f"{self.username}:{self.password}".encode()).decode() auth = base64.b64encode(f"{self.username}:{self.password}".encode()).decode()
req.add_header('Authorization', f"Basic {auth}") req.add_header('Authorization', f"Basic {auth}")
with urllib.request.urlopen(req) as response: with urllib.request.urlopen(req, timeout=60) as response:
data = response.read() data = response.read()
# 缓存 # 缓存
+52 -11
View File
@@ -45,6 +45,10 @@ class HttpMirror:
self.base_dir = config.get('base_dir', './downloads') self.base_dir = config.get('base_dir', './downloads')
self.cache_enabled = config.get('cache_enabled', True) self.cache_enabled = config.get('cache_enabled', True)
self.cache_ttl = config.get('cache_ttl', 3600) # 默认1小时 self.cache_ttl = config.get('cache_ttl', 3600) # 默认1小时
# 仅缓存小响应(内存缓存上限),大文件只做流式转发
self.max_cache_bytes = config.get('max_cache_bytes', 16 * 1024 * 1024)
# 单次转发上限(防止无界拉取)
self.max_stream_bytes = config.get('max_stream_bytes', 4 * 1024 * 1024 * 1024)
# 确保存储目录存在 # 确保存储目录存在
os.makedirs(self.storage_dir, exist_ok=True) os.makedirs(self.storage_dir, exist_ok=True)
@@ -267,7 +271,11 @@ class HttpMirror:
handler.end_headers() handler.end_headers()
with open(local_path, 'rb') as f: with open(local_path, 'rb') as f:
handler.wfile.write(f.read()) while True:
chunk = f.read(64 * 1024)
if not chunk:
break
handler.wfile.write(chunk)
return True return True
except Exception as e: except Exception as e:
@@ -323,13 +331,27 @@ class HttpMirror:
handler.send_header('X-Upstream', self.upstream_url) handler.send_header('X-Upstream', self.upstream_url)
handler.end_headers() handler.end_headers()
# 读取并转发内容,同时缓存 # 流式转发(避免整文件读入内存导致 OOM),小响应可选缓存
data = response.read() total = 0
cache_buf = bytearray()
try:
while True:
chunk = response.read(64 * 1024)
if not chunk:
break
total += len(chunk)
if total > self.max_stream_bytes:
break
handler.wfile.write(chunk)
if self.cache_enabled and response.status == 200 and len(cache_buf) < self.max_cache_bytes:
cache_buf.extend(chunk)
except (BrokenPipeError, ConnectionResetError):
return False
finally:
response.close()
if self.cache_enabled and response.status == 200: if self.cache_enabled and response.status == 200 and total <= self.max_cache_bytes:
self._set_cache(cache_key, data, content_type) self._set_cache(cache_key, bytes(cache_buf), content_type)
handler.wfile.write(data)
return True return True
except Exception as e: except Exception as e:
@@ -337,13 +359,32 @@ class HttpMirror:
return False return False
def _serve_cached(self, handler, cached: dict, url: str): def _serve_cached(self, handler, cached: dict, url: str):
"""服务缓存的文件""" """服务缓存的文件(支持单段 Range 请求)"""
handler.send_response(200) data = cached.get('data', b'')
total = len(data)
start, end = 0, total - 1
range_header = handler.headers.get('Range')
if range_header:
m = re.match(r'bytes=(\d+)-(\d*)', range_header)
if m:
start = int(m.group(1))
if start >= total:
handler.send_error(416, 'Requested Range Not Satisfiable')
return False
end = int(m.group(2)) if m.group(2) else total - 1
end = min(end, total - 1)
if start > 0 or end < total - 1:
handler.send_response(206)
handler.send_header('Content-Range', f'bytes {start}-{end}/{total}')
else:
handler.send_response(200)
handler.send_header('Content-Type', cached.get('content_type', 'application/octet-stream')) handler.send_header('Content-Type', cached.get('content_type', 'application/octet-stream'))
handler.send_header('Content-Length', len(cached.get('data', b''))) handler.send_header('Content-Length', str(end - start + 1))
handler.send_header('X-Cached', 'true') handler.send_header('X-Cached', 'true')
handler.end_headers() handler.end_headers()
handler.wfile.write(cached.get('data', b'')) handler.wfile.write(data[start:end + 1])
return True return True
def _get_cache_key(self, url: str) -> str: def _get_cache_key(self, url: str) -> str:
+21 -8
View File
@@ -524,16 +524,29 @@ class PyPIMirror:
req.add_header('User-Agent', 'PyPI-Mirror/1.0') req.add_header('User-Agent', 'PyPI-Mirror/1.0')
with urllib.request.urlopen(req, timeout=120) as response: with urllib.request.urlopen(req, timeout=120) as response:
data = response.read() # 流式转发,避免大文件整读入内存;小响应可缓存
handler.send_response(200)
handler.send_header('Content-Type', 'application/octet-stream')
handler.send_header('Content-Length', response.headers.get('Content-Length', ''))
handler.end_headers()
if True: total = 0
self._set_cache(cache_key, data) cache_buf = bytearray()
max_cache = 16 * 1024 * 1024 # 仅缓存 <=16MB 的文件
max_stream = 4 * 1024 * 1024 * 1024
while True:
chunk = response.read(64 * 1024)
if not chunk:
break
total += len(chunk)
if total > max_stream:
break
handler.wfile.write(chunk)
if len(cache_buf) < max_cache:
cache_buf.extend(chunk)
handler.send_response(200) if total <= max_cache:
handler.send_header('Content-Type', 'application/octet-stream') self._set_cache(cache_key, bytes(cache_buf))
handler.send_header('Content-Length', str(len(data)))
handler.end_headers()
handler.wfile.write(data)
return True return True
except Exception as e: except Exception as e: