Flask/Quart WSGI中间件无法添加X-Time-Taken响应头问题

问题现象 #

• 适配Quart/Flask的WSGI中间件开发中，计划通过响应处理逻辑记录请求日志，同时将请求处理时长注入自定义响应头X-Time-Taken
• 代码运行后自定义头始终未出现在最终响应中，实际返回的响应头仅包含Content-Type: text/plain; charset=utf-8
• 参考公开实现编写的逻辑未按预期生效，需要定位具体故障点

故障复现代码 #

class SampleWSGI():    
    def __init__(self, app, config=None):
        self.app = app
        self.app_id = config["app_id"]
        self.secret_key = config["secret_key"]

    def __call__(self, environ, start_response):
        request = Request(environ)
        resp = Response(start_response)
        self._process_response(self.app_id, resp)
        apikey = request.args.get("apikey")
        if "favicon.ico" in request.path:
            return

        x = self.authorize_user(apikey=apikey, client_path=request.path)
        if x.status_code < 400:

            ##Injects duration into response time
            start_time = datetime.utcnow().timestamp()
            def injecting_start_response(status, headers, exc_info=None):
                end_time = datetime.utcnow().timestamp()
                time_taken = str(end_time - start_time)
                headers.append(('X-Time-Taken', time_taken))
                return start_response(status, headers, exc_info)
        
            return self.app(environ, injecting_start_response)
        
        res = Response(u'Authorization failed', mimetype= 'text/plain', status=x.status_code)
        return res(environ, start_response)

    @staticmethod
    def _process_response(app_id, response, request):
        #req = Request(environ, shallow=True)
        print(f'request header: {response.headers}')

故障根因 #

1. 提前触发start_response导致响应头提前发送
   方法入口处resp = Response(start_response)会直接触发WSGI协议规定的响应头发送流程，此时后续鉴权、计时、头注入逻辑都还没执行，服务器已经把仅包含默认Content-Type的头发送给客户端，后续所有头修改操作都不会生效。
2. 方法调用参数不匹配+日志打点时机错误
   定义的_process_response静态方法需要传入3个位置参数，实际调用时仅传入2个，代码运行到这一行会直接抛出参数缺失错误；同时在请求刚进入中间件、业务逻辑还没执行、响应完全没生成的时候就调用日志方法，根本不可能拿到最终的响应头数据。
3. favicon路径分支违反WSGI协议规范
   匹配到favicon请求时直接执行空return，既没有调用start_response设置响应状态和头，也没有返回可迭代的响应体，WSGI服务器处理该路径时会直接抛出协议错误，中断整个请求流程。
4. 响应头引用被框架覆盖
   Quart作为异步框架在WSGI兼容模式下、Flask在调试模式下，会在业务代码返回后重新生成响应头列表，如果你注入头时操作的列表不是框架最终发送时持有的引用，append的自定义头会被直接丢弃。

修复后可运行代码 #

from datetime import datetime
from werkzeug.wrappers import Request, Response

class SampleWSGI():    
    def __init__(self, app, config=None):
        self.app = app
        self.app_id = config["app_id"]
        self.secret_key = config["secret_key"]

    def __call__(self, environ, start_response):
        request = Request(environ)
        apikey = request.args.get("apikey")
        # 修复favicon分支，符合WSGI返回要求
        if "favicon.ico" in request.path:
            favicon_resp = Response(status=204)
            return favicon_resp(environ, start_response)

        auth_result = self.authorize_user(apikey=apikey, client_path=request.path)
        if auth_result.status_code >= 400:
            fail_resp = Response(u'Authorization failed', mimetype= 'text/plain', status=auth_result.status_code)
            return fail_resp(environ, start_response)

        # 正常请求：计时+头注入+日志打点
        start_time = datetime.utcnow().timestamp()
        def injecting_start_response(status, headers, exc_info=None):
            end_time = datetime.utcnow().timestamp()
            time_taken = str(end_time - start_time)
            # 去重处理，避免重复头被框架过滤
            for idx, (header_name, header_val) in enumerate(headers):
                if header_name.lower() == "x-time-taken":
                    headers[idx] = ("X-Time-Taken", time_taken)
                    break
            else:
                headers.append(("X-Time-Taken", time_taken))
            # 此时已拿到最终响应头，再执行日志记录
            self._process_response(self.app_id, headers, request)
            return start_response(status, headers, exc_info)
        
        return self.app(environ, injecting_start_response)

    @staticmethod
    def _process_response(app_id, resp_headers, request):
        print(f"final response headers: {resp_headers}")

注意事项 #

• WSGI中间件开发中绝对不能在所有头处理逻辑完成前提前触发start_response，否则所有头修改操作都会失效
• 所有请求分支必须满足WSGI协议要求：调用start_response传入状态和头，同时返回可迭代的响应体
• 注入自定义头时优先做去重处理，避免WSGI服务器或框架因为重复头字段自动清理自定义内容

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