Python Socket编程如何限制单IP每分钟TCP请求次数以防护DDoS攻击

可行性说明 #

该需求完全可以实现，针对你的小型TCP服务场景，使用固定窗口速率限制算法即可满足需求，实现复杂度低、性能开销小，足以抵御小型DDoS攻击和恶意高频请求。

核心实现逻辑 #

• 用字典存储每个IP的历史请求时间戳列表
• 每次新请求到来时，先清理该IP列表中超出1分钟时间窗口的旧记录
• 若剩余记录数超过你设定的阈值，则直接拒绝该请求
• 多线程场景下加线程锁保证数据读写安全，避免并发冲突

具体实现步骤 #

1. 新增依赖导入和全局配置 #

在代码头部新增时间模块和线程锁导入，添加速率限制相关配置项：

import time
from threading import Thread, Lock

新增全局变量：

# 速率限制配置：单IP每分钟最多允许10次请求，可按需调整
RATE_LIMIT = 10
TIME_WINDOW = 60
# 存储每个IP的请求时间列表
ip_request_records = {}
# 线程锁，避免多线程同时修改请求记录字典出现数据错误
rate_lock = Lock()

2. 新增速率校验函数 #

定义独立的速率校验函数，返回True表示允许访问，False表示超出限制：

def check_rate_limit(ip):
    current_time = time.time()
    with rate_lock:
        # 初始化新IP的记录
        if ip not in ip_request_records:
            ip_request_records[ip] = []
        # 清理超出时间窗口的旧记录
        ip_request_records[ip] = [t for t in ip_request_records[ip] if current_time - t < TIME_WINDOW]
        # 判断是否超出限制
        if len(ip_request_records[ip]) < RATE_LIMIT:
            ip_request_records[ip].append(current_time)
            return True
        return False

3. 修改主连接处理逻辑 #

修复原有代码先开线程再校验连接的逻辑错误，新增速率校验步骤：

while True:
    conn, addr = s.accept()
    client_ip = addr[0]
    # 第一步：先校验速率是否超限
    if not check_rate_limit(client_ip):
        print(f"速率限制：拒绝来自{client_ip}:{addr[1]}的连接，请求超出每分钟上限")
        conn.close()
        continue
    # 第二步：校验是否已经存在活跃连接
    if client_ip in con:
        print(f"连接拒绝：来自{client_ip}:{addr[1]}的重复连接")
        conn.close()
        continue
    # 校验通过才添加连接记录、启动处理线程
    con.append(client_ip)
    print(f"连接接受：来自{client_ip}:{addr[1]}的连接")
    start_new_math_thread(conn, addr)

4. 可选优化：定期清理过期IP记录 #

避免长时间运行后请求记录字典占用过多内存，可新增一个定时清理线程：

def clean_expired_records():
    while True:
        time.sleep(300) # 每5分钟清理一次
        current_time = time.time()
        with rate_lock:
            # 删除超过10分钟没有新请求的IP记录
            expired_ips = [ip for ip, records in ip_request_records.items() if current_time - records[-1] > 600]
            for ip in expired_ips:
                del ip_request_records[ip]

# 服务启动时启动清理线程
clean_thread = Thread(target=clean_expired_records, daemon=True)
clean_thread.start()

修改后完整可运行代码 #

import socket
import re
import time
from subprocess import Popen, STDOUT, PIPE
from threading import Thread, Lock

# 速率限制配置
RATE_LIMIT = 10
TIME_WINDOW = 60
ip_request_records = {}
rate_lock = Lock()

def check_rate_limit(ip):
    current_time = time.time()
    with rate_lock:
        if ip not in ip_request_records:
            ip_request_records[ip] = []
        ip_request_records[ip] = [t for t in ip_request_records[ip] if current_time - t < TIME_WINDOW]
        if len(ip_request_records[ip]) < RATE_LIMIT:
            ip_request_records[ip].append(current_time)
            return True
        return False

def clean_expired_records():
    while True:
        time.sleep(300)
        current_time = time.time()
        with rate_lock:
            expired_ips = [ip for ip, records in ip_request_records.items() if current_time - records[-1] > 600]
            for ip in expired_ips:
                del ip_request_records[ip]

def start_new_math_thread(conn, addr):
    t = MathServerCommunicationThread(conn, addr)
    t.start()

class ProcessOutputThread(Thread):
    def __init__(self, proc, conn):
        Thread.__init__(self)
        self.proc = proc
        self.conn = conn
        
    def run(self):
        while not self.proc.stdout.closed and not self.conn._closed:
            try:
                self.conn.sendall(self.proc.stdout.readline())
            except:
                pass
            
class MathServerCommunicationThread(Thread):
    def __init__(self, conn, addr):
        Thread.__init__(self)
        self.conn = conn
        self.addr = addr
    
    def run(self):
        self.conn.sendall("Simple Math Server developed by LAHTP Version-2.o ('integrated with Firewall' ) \n\nGive any math expressions, and I will answer you...\U0001F609 \n\n$ ".encode())
        p = Popen(['bc'], stdout=PIPE, stdin=PIPE, stderr=STDOUT)
        output = ProcessOutputThread(p, self.conn)
        output.start()

        while not p.stdout.closed or not self.conn._closed:
            try:
                data = self.conn.recv(1024)
                if not data:
                    break
                else:
                    try:
                        data = data.decode()
                        query = data.strip()
                        if query == 'quit' or query == 'exit':
                            print(f"{self.addr[0]} disconnected...!")
                            con.remove(self.addr[0])
                            p.communicate(query.encode(), timeout=1)
                            if p.poll() is not None:
                                break
                            
                        query = query + '\n'
                        if not re.match("^[0-9|+|-|*|/|^|=]{1,15}$",query):
                            self.conn.send(b"Only Integers are allowed & it should range between 0to15 digits\n")
                            self.conn.close()
                            print (f"{self.addr[0]} is trying to exploit the resources so the connection is closed.")
                            if self.addr[0] in con:
                                con.remove(self.addr[0])
                        else:
                            p.stdin.write(query.encode())
                            p.stdin.flush()
                    except Exception as e:
                        pass
            except:
                pass
        self.conn.close()

HOST = ''
PORT = 8877
con=[]

# 启动过期记录清理线程
clean_thread = Thread(target=clean_expired_records, daemon=True)
clean_thread.start()

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
s.bind((HOST, PORT))
print("Server initialized\n Waiting for the incoming connections...")
s.listen()
while True:
    conn, addr = s.accept()
    client_ip = addr[0]
    # 速率校验
    if not check_rate_limit(client_ip):
        print(f"Rate limit: reject connection from {client_ip}:{addr[1]}, exceed max requests per minute")
        conn.close()
        continue
    # 重复连接校验
    if client_ip in con:
        print(f"connection rejected from {client_ip}:{addr[1]}")
        conn.close()
        continue
    con.append(client_ip)
    print(f"connection accepted from {client_ip}:{addr[1]}")
    start_new_math_thread(conn, addr)

注意事项 #

• 你可以根据实际业务需求调整RATE_LIMIT的数值，比如设置为20就是单IP每分钟允许20次请求
• 该实现是应用层限流，若遇到大流量DDoS攻击，建议搭配操作系统层面的防火墙规则或者云厂商的DDoS防护服务使用
• 代码中已经修复了原有逻辑的两处小问题：修复了先启动线程再校验连接的逻辑错误、修复了线程内部调用全局变量addr和conn的异常问题

内容的提问来源于stack exchange，提问作者Rogith