feat: add network filter

competitions/runner.py

@@ -15,7 +15,7 @@ from loguru import logger
 
 from competitions.enums import SubmissionStatus, ErrorMessage
 from competitions.info import CompetitionInfo
-from competitions.utils import user_token_api, space_cleaner
+from competitions.utils import user_token_api, space_cleaner, dockerfile_modifier
 
 
 @dataclass
@@ -201,6 +201,14 @@ class JobRunner:
         )
         with open(f"{client_code_local_dir}/README.md", "w", encoding="utf-8") as f:
             f.write(self._create_readme(space_id))
+        shutil.copyfile("./other_files/network_filter.so", os.path.join(client_code_local_dir, "network_filter.so"))
+        for filename in os.listdir(client_code_local_dir):
+            if filename.lower() == "dockerfile":
+                filepath = os.path.join(client_code_local_dir, filename)
+                with open(filepath, "r", encoding="utf-8") as f:
+                    dockerfile_content = f.read()
+                with open(filepath, "w", encoding="utf-8") as f:
+                    f.write(dockerfile_modifier.modify_dockerfile_content(dockerfile_content)[0])
         try:
             api.upload_folder(
                 repo_id=space_id,

competitions/utils.py

@@ -3,12 +3,12 @@ import json
 import os
 import shlex
 import subprocess
-import traceback
+import re
 import threading
 import uuid
 import base64
 import glob
-from typing import Optional, Dict, Any, List, Literal
+from typing import Optional, Dict, Any, List, Literal, Tuple
 from collections import defaultdict
 from datetime import datetime, timezone, timedelta
 
@@ -755,3 +755,163 @@ leaderboard_api = LeaderboardApi(
     hf_token=os.environ.get("HF_TOKEN", None),
     competition_id=os.environ.get("COMPETITION_ID")
 )
+
+
+
+class DockerfileModifier:
+    def __init__(self, allowed_hosts: str, source_so_path: str = "./network_filter.so"):
+        self.allowed_hosts = allowed_hosts
+        self.source_so_path = source_so_path
+        self.tatget_so_dir = "/_app_extensions"
+        self.tatget_so_path = os.path.join(self.tatget_so_dir, "network_filter.so")
+        self.preload_prefix = f'LD_PRELOAD={self.tatget_so_path} ALLOWED_HOSTS="{allowed_hosts}"'
+        
+    def parse_dockerfile_line(self, line: str) -> Tuple[str, str, str]:
+        """
+        解析 Dockerfile 行，返回 (指令名, 原始命令, 格式类型)
+        格式类型: 'exec' (JSON数组) 或 'shell' (shell命令)
+        """
+        line = line.strip()
+        
+        # 匹配 CMD 或 ENTRYPOINT
+        cmd_match = re.match(r'^(CMD|ENTRYPOINT)\s+(.+)$', line, re.IGNORECASE)
+        if not cmd_match:
+            return "", "", ""
+            
+        instruction = cmd_match.group(1).upper()
+        command_part = cmd_match.group(2).strip()
+        
+        # 判断是 exec 格式 (JSON数组) 还是 shell 格式
+        if command_part.startswith('[') and command_part.endswith(']'):
+            return instruction, command_part, "exec"
+        else:
+            return instruction, command_part, "shell"
+    
+    def modify_shell_format(self, command: str) -> str:
+        """修改 shell 格式的命令"""
+        # 在原命令前添加环境变量
+        return f'{self.preload_prefix} {command}'
+    
+    def modify_exec_format(self, command: str) -> str:
+        """修改 exec 格式 (JSON数组) 的命令"""
+        try:
+            # 解析 JSON 数组格式
+            # 移除外层的方括号
+            inner = command[1:-1].strip()
+            
+            # 简单的 JSON 数组解析
+            parts = []
+            current = ""
+            in_quotes = False
+            escape_next = False
+            
+            for char in inner:
+                if escape_next:
+                    current += char
+                    escape_next = False
+                elif char == '\\':
+                    current += char
+                    escape_next = True
+                elif char == '"' and not escape_next:
+                    in_quotes = not in_quotes
+                    current += char
+                elif char == ',' and not in_quotes:
+                    parts.append(current.strip())
+                    current = ""
+                else:
+                    current += char
+            
+            if current.strip():
+                parts.append(current.strip())
+            
+            # 移除引号并处理转义
+            cleaned_parts = []
+            for part in parts:
+                part = part.strip()
+                if part.startswith('"') and part.endswith('"'):
+                    part = part[1:-1]
+                # 处理基本的转义字符
+                part = part.replace('\\"', '"').replace('\\\\', '\\')
+                cleaned_parts.append(part)
+            
+            if not cleaned_parts:
+                return command
+            
+            # 构建新的命令
+            # 第一个元素通常是 shell (/bin/sh, /bin/bash 等)
+            # 如果第一个元素是 shell，修改执行的命令
+            if len(cleaned_parts) >= 3 and cleaned_parts[0] in ['/bin/sh', '/bin/bash', 'sh', 'bash']:
+                if cleaned_parts[1] == '-c':
+                    # 格式: ["/bin/sh", "-c", "command"]
+                    original_cmd = cleaned_parts[2]
+                    new_cmd = f'{self.preload_prefix} {original_cmd}'
+                    new_parts = [cleaned_parts[0], cleaned_parts[1], new_cmd] + cleaned_parts[3:]
+                else:
+                    # 直接在现有命令前添加环境变量，通过 shell 执行
+                    original_cmd = ' '.join(cleaned_parts[1:])
+                    new_cmd = f'{self.preload_prefix} {original_cmd}'
+                    new_parts = [cleaned_parts[0], '-c', new_cmd]
+            else:
+                # 直接执行的命令，需要通过 shell 包装
+                original_cmd = ' '.join(cleaned_parts)
+                new_parts = ['/bin/sh', '-c', f'{self.preload_prefix} {original_cmd}']
+            
+            # 重新构建 JSON 数组
+            escaped_parts = []
+            for part in new_parts:
+                # 转义引号和反斜杠
+                escaped = part.replace('\\', '\\\\').replace('"', '\\"')
+                escaped_parts.append(f'"{escaped}"')
+            
+            return '[' + ', '.join(escaped_parts) + ']'
+            
+        except Exception as e:
+            print(f"警告: 解析 exec 格式失败: {e}")
+            print(f"原始命令: {command}")
+            # 如果解析失败，转换为 shell 格式
+            return f'{self.preload_prefix} {command}'
+    
+    def modify_dockerfile_content(self, content: str) -> Tuple[str, List[str]]:
+        """
+        修改 Dockerfile 内容
+        返回: (修改后的内容, 修改日志)
+        """
+        lines = content.splitlines()
+        modified_lines = []
+        changes = []
+        
+        for i, line in enumerate(lines, 1):
+            instruction, command, format_type = self.parse_dockerfile_line(line)
+            
+            if instruction in ['CMD', 'ENTRYPOINT'] and command:
+                if format_type == "shell":
+                    new_command = self.modify_shell_format(command)
+                    new_line = f'{instruction} {new_command}'
+                elif format_type == "exec":
+                    new_command = self.modify_exec_format(command)
+                    new_line = f'{instruction} {new_command}'
+                else:
+                    new_line = line
+                
+                changes.append(f"第 {i} 行: {instruction} 指令已修改")
+                changes.append(f"  原始: {line}")
+                changes.append(f"  修改: {new_line}")
+                modified_lines.append(new_line)
+            else:
+                modified_lines.append(line)
+
+        last_user = None
+        for line in modified_lines[::-1]:
+            if line.startswith("USER"):
+                last_user = line.split()[1].strip()
+
+        if last_user is None:
+            modified_lines.insert(-1, f"COPY {self.source_so_path}" + f" {self.tatget_so_path}")
+        else:
+            modified_lines.insert(-1, f"COPY --chown={last_user}:{last_user} {self.source_so_path} {self.tatget_so_path}")
+            modified_lines.insert(-1, f"RUN chown -R {last_user}:{last_user} {self.tatget_so_dir}")
+
+        return '\n'.join(modified_lines), changes
+
+
+dockerfile_modifier = DockerfileModifier("127.0.0.1,xdimlab-hugsim-web-server-0.hf.space")