🕳️

Pythonで実装する、IPアドレスベースの同時接続数制限

に公開
Python
tech

はじめに

前回の記事では、アイドル状態のTCP接続が大量に存在すると、サーバーのメモリが枯渇しうること を確認しました。悪意のあるユーザーや、バグのあるクライアントからの過剰な接続は、サービス全体を停止させる原因となります。
では、特定のユーザー（IPアドレス）からの過剰な接続を防ぐには、どうすれば良いのでしょうか？
今回は、IPアドレスごとに同時接続数を制限するシンプルなレートリミッタを、Pythonのasyncioを使って実装します。

マシンスペック

MacBook Air M2 arm64

準備

Dockerファイルの作成

# Dockerfile
FROM python:3.11-slim

# procps: freeコマンドなど、メモリやプロセス情報を確認するツール群
RUN apt-get update && apt-get install -y procps iproute2 netcat-traditional

WORKDIR /app

サーバ

# ratelimit_server.py
import asyncio
from collections import defaultdict

# IPアドレスごとの接続数を保持する辞書
# defaultdictを使うと、キーが存在しない場合に自動で0を返してくれるので便利
CONNECTION_COUNTS = defaultdict(int)
# 1IPアドレスあたりの最大同時接続数
CONNECTION_LIMIT = 5

async def handle_connection(reader, writer):
    # 接続元の情報を取得
    peername = writer.get_extra_info('peername')
    if peername is None:
        writer.close()
        await writer.wait_closed()
        return
        
    ip = peername[0]

    # --- 接続数制限のロジック ---
    # MutexやLockなしで安全にアクセスするため、ここで一度に処理
    current_count = CONNECTION_COUNTS.get(ip, 0)
    if current_count >= CONNECTION_LIMIT:
        print(f"[拒否] {ip}からの接続は上限({CONNECTION_LIMIT})を超えました。")
        writer.close()
        await writer.wait_closed()
        return

    # 上限に達していなければ、カウントを増やして接続を許可
    CONNECTION_COUNTS[ip] = current_count + 1
    print(f"[許可] {ip}からの新しい接続。 (現在: {CONNECTION_COUNTS[ip]}接続)")
    # -----------------------------

    try:
        # クライアントが接続を切るまで待機
        while True:
            data = await reader.read(1024)
            if not data:
                break
    except Exception as e:
        print(f"エラー発生: {e}")
    finally:
        # --- 接続終了時の後処理 ---
        # 接続が切れたら、必ずカウントを減らす
        CONNECTION_COUNTS[ip] -= 1
        print(f"[切断] {ip}が接続を終了しました。 (残り: {CONNECTION_COUNTS[ip]}接続)")
        # カウントが0になったら、メモリ節約のために辞書からキーを削除
        if CONNECTION_COUNTS[ip] == 0:
            del CONNECTION_COUNTS[ip]
        
        writer.close()
        await writer.wait_closed()

async def main():
    server = await asyncio.start_server(handle_connection, '0.0.0.0', 8888)
    addr = server.sockets[0].getsockname()
    print(f'レートリミッタ付きサーバーが {addr} で起動しました (接続上限: {CONNECTION_LIMIT}/IP)')

    async with server:
        await server.serve_forever()

if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nサーバーを停止します。")

terminalのメインプロセスで実施（サーバ）

docker build -t memory-test .
# RAMは256に制限
docker run --rm -it --name mem-test --memory=256m -v "$PWD":/app memory-test bash
python3 ratelimit_server.py

terminalの別プロセスで実施(クライアント）

docker exec -it mem-test bash
for i in $(seq 1 6); do nc localhost 8888 & done

結果

クライアントからの接続時に下記のメッセージが出ました。

レートリミッタ付きサーバーが ('0.0.0.0', 8888) で起動しました (接続上限: 5/IP)
[許可] 127.0.0.1からの新しい接続。 (現在: 1接続)
[許可] 127.0.0.1からの新しい接続。 (現在: 2接続)
[許可] 127.0.0.1からの新しい接続。 (現在: 3接続)
[許可] 127.0.0.1からの新しい接続。 (現在: 4接続)
[許可] 127.0.0.1からの新しい接続。 (現在: 5接続)
[拒否] 127.0.0.1からの接続は上限(5)を超えました。

まとめ

今回は、クライアントの同時接続数を制限する実験をしました。皆様の学習の参考になれば幸いです。

Discussion