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)

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

结论

通过添加 IP 白名单功能，ComfyUI 可以更好地控制谁可以访问其服务，从而增强了整体的安全性。这种方法对于希望限制其服务只对特定用户或网络开放的应用程序尤其有用。当然，安全是一个多层次的话题，IP白名单只是其中的一部分。在实际应用中，我们还需要考虑更多的安全措施来保护我们的应用和用户。

感谢阅读，希望这篇博客对您有所帮助！