ComfyUI添加IP白名单功能-Toy模板网

这篇具有很好参考价值的文章主要介绍了ComfyUI添加IP白名单功能。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

AI生图很火，相信你对ComfyUI不陌生，查看ComfyUI的源码可以发现它是使用aiohttp来作为服务端的。那么我们在使用ComfyUI的时候可能需要做一些安全的限制，接下来我们将探讨如何在 ComfyUI 中添加 IP 白名单功能，以确保只有特定的用户可以访问我们的服务。

什么是IP白名单？

IP白名单是一种安全措施，用于限制只有特定的IP地址可以访问网络服务。这种方法非常适合那些只希望向特定用户或内部网络开放的应用。

实现步骤

为了在 ComfyUI 中添加 IP 白名单功能，我们需要创建一个中间件来检查每个进入的请求。如果请求来自不在白名单中的 IP 地址，我们将阻止该请求并返回一个 403 状态码。

第一步：定义IP白名单

将以下代码粘贴在ComfyUI根目录中的server.py文件当中

# 定义IP白名单
ALLOWED_IPS = ['192.168.1.1', '192.168.1.2']

在这里，我们定义了一个简单的 IP 白名单，只允许来自 192.168.1.1 和 192.168.1.2 的请求。

第二步：创建IP白名单中间件

将以下代码粘贴在ComfyUI根目录中的server.py文件当中

@web.middleware
async def ip_whitelist_middleware(request: web.Request, handler):
    peername = request.transport.get_extra_info('peername')
    if peername is not None:
        host, _ = peername
        if host not in ALLOWED_IPS:
            return web.Response(text="Access denied.", status=403)

    return await handler(request)

这个中间件检查每个请求的来源IP是否在我们定义的白名单中。如果不在，它会立即返回一个 403 状态的响应。

第三步：将中间件添加到应用

在 PromptServer 类中，我们确保无论其他配置如何，都添加了 IP 白名单中间件。

在middlewares字典中追加上ip_whitelist_middleware，代码如下：

class PromptServer():
    def __init__(self, loop):
        # ... 省略其他初始化代码 ...

        middlewares = [cache_control, ip_whitelist_middleware]
        if args.enable_cors_header:
            middlewares.append(create_cors_middleware(args.enable_cors_header))

        max_upload_size = round(args.max_upload_size * 1024 * 1024)
        self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)

        # ... 省略其他初始化代码 ...