HYCEXITADMIN/mirror_server
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

P1: 架构加固

Browse Source 
- 线程池服务器: 连接级读写超时 + 有界队列(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
Download Patch File Download Diff File Expand all files Collapse all files
core/server.py
+38 -5
View File
@@ -7,6 +7,7 @@ import os
import sys
import ssl
import signal
import socket
import mimetypes
import threading
import socketserver
@@ -18,29 +19,59 @@ from .mirror_sync import MirrorSyncManager
class ThreadPoolHTTPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
"""使用线程池的 HTTP 服务器"""
"""使用线程池的 HTTP 服务器（有界队列 + 连接超时，防慢速 DoS）"""
allow_reuse_address = 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.queue_size = queue_size
self.socket_timeout = socket_timeout
self._executor = ThreadPoolExecutor(
max_workers=max_workers,
thread_name_prefix="http_handler"
)
# 有界并发槽位（活跃 + 排队），超过上限直接拒绝新连接，
# 防止连接洪水导致任务无限堆积在内存
self._slots = threading.BoundedSemaphore(max_workers + queue_size)
super().__init__(server_address, RequestHandlerClass)
def process_request(self, 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):
"""实际处理请求"""
try:
# 连接级读写超时，防慢客户端无限占用工作线程
try:
request.settimeout(self.socket_timeout)
except Exception:
pass
self.finish_request(request, client_address)
except Exception:
self.handle_error(request, client_address)
finally:
self._slots.release()
self.shutdown_request(request)
def server_close(self):
@@ -118,7 +149,9 @@ class MirrorServer:
self.server = ThreadPoolHTTPServer(
server_address,
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)
# 设置超时
# 设置 accept 轮询超时
self.server.timeout = self.config.get('timeout', 30)
# 启用HTTPS
mirrors/docker.py
+2 -2
View File
@@ -141,7 +141,7 @@ class DockerMirror:
auth = base64.b64encode(f"{self.username}:{self.password}".encode()).decode()
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()
# 缓存
@@ -189,7 +189,7 @@ class DockerMirror:
auth = base64.b64encode(f"{self.username}:{self.password}".encode()).decode()
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()
# 缓存
mirrors/http.py
+51 -10
View File
@@ -45,6 +45,10 @@ class HttpMirror:
self.base_dir = config.get('base_dir', './downloads')
self.cache_enabled = config.get('cache_enabled', True)
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)
@@ -267,7 +271,11 @@ class HttpMirror:
handler.end_headers()
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
except Exception as e:
@@ -323,13 +331,27 @@ class HttpMirror:
handler.send_header('X-Upstream', self.upstream_url)
handler.end_headers()
# 读取并转发内容，同时缓存
data = response.read()
# 流式转发（避免整文件读入内存导致 OOM），小响应可选缓存
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:
self._set_cache(cache_key, data, content_type)
handler.wfile.write(data)
if self.cache_enabled and response.status == 200 and total <= self.max_cache_bytes:
self._set_cache(cache_key, bytes(cache_buf), content_type)
return True
except Exception as e:
@@ -337,13 +359,32 @@ class HttpMirror:
return False
def _serve_cached(self, handler, cached: dict, url: str):
"""服务缓存的文件"""
"""服务缓存的文件（支持单段 Range 请求）"""
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-Length', len(cached.get('data', b'')))
handler.send_header('Content-Length', str(end - start + 1))
handler.send_header('X-Cached', 'true')
handler.end_headers()
handler.wfile.write(cached.get('data', b''))
handler.wfile.write(data[start:end + 1])
return True
def _get_cache_key(self, url: str) -> str:
mirrors/pypi.py
+20 -7
View File
@@ -524,16 +524,29 @@ class PyPIMirror:
req.add_header('User-Agent', 'PyPI-Mirror/1.0')
with urllib.request.urlopen(req, timeout=120) as response:
data = response.read()
if True:
self._set_cache(cache_key, data)
# 流式转发，避免大文件整读入内存；小响应可缓存
handler.send_response(200)
handler.send_header('Content-Type', 'application/octet-stream')
handler.send_header('Content-Length', str(len(data)))
handler.send_header('Content-Length', response.headers.get('Content-Length', ''))
handler.end_headers()
handler.wfile.write(data)
total = 0
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)
if total <= max_cache:
self._set_cache(cache_key, bytes(cache_buf))
return True
except Exception as e: