feat(apigw_manager/plugin): 新增 response_rewrite 和 redirect 插件的 build 函数 (#223)

Author: F-cq

sdks/apigw-manager/src/apigw_manager/plugin/config.py

@@ -565,6 +565,123 @@ def build_request_validation(
     }
 
 
+@dataclass
+class HeadersConfig:
+    add: Dict[str, str]
+    set: Dict[str, str]
+    remove: List[str]
+
+
+def build_response_rewrite_validation(
+        status_code: int,
+        body: str,
+        vars: Optional[str],
+        body_base64: bool = False,
+        headers: Optional[HeadersConfig] = None,
+) -> Dict[str, str]:
+    """generate response-rewrite plugin config
+
+    Args:
+        status_code (int): 修改上游返回状态码，默认保留原始响应代码。
+        body (str):
+            - 修改上游返回的 body 内容，如果设置了新内容，header 里面的 content-length 字段也会被去掉。
+            - 注意，这个字段只允许对插件配置中传递的主体进行解码，并不对上游响应进行解码。
+        vars (str, Optional):
+            - vars 是一个表达式列表，只有满足条件的请求和响应才会修改 body 和 header 信息，来自 lua-resty-expr。
+            - 如果 vars 字段为空，那么所有的重写动作都会被无条件的执行。
+        body_base64 (bool, Optional): 当设置时，在写给客户端之前，在body中传递的主体将被解码，这在一些图像和 Protobuffer 场景中使用。
+        headers (HeadersConfig, optional):
+            - add (Dict[str, str]): 添加新的响应头。格式为 {"name": "value", ...}。这个值能够以 $var 的格式包含 NGINX 变量。
+            - set (Dict[str, str]): 改写响应头。格式为 {"name": "value", ...}。这个值能够以 $var 的格式包含 NGINX 变量。
+            - remove (List[str]): 移除响应头。格式为 ["name", ...]。
+
+    Raises:
+        ValueError: status_code must be between 200 and 598
+        ValueError: key {} can not contain ':'
+
+    Returns:
+        {
+            "type": "response-rewrite",
+            "yaml": "body: ''\nbody_base64: false\nheaders:\n  add:\n  - key: key1:value1\n  remove:\n  - key: key1\n  set:\n  - key: key1\n    value: value1\nstatus_code: 200\nvars: ''\n"
+        }
+    """
+
+    if not (200 <= status_code <= 598):
+        raise ValueError("status_code must be between 200 and 598")
+
+    if vars:
+        _check_vars(vars, "response_rewrite")
+
+    add_data = []
+    for k, v in headers.add.items():
+        if ":" in k:
+            raise ValueError(f"key {k} can not contain ':'")
+        add_data.append({"key": "{}:{}".format(k, v)})
+
+    set_data = []
+    for k, v in headers.set.items():
+        if ":" in k:
+            raise ValueError(f"key {k} can not contain ':'")
+        set_data.append({"key": k, "value": v})
+
+    remove_data = []
+    for k in headers.remove:
+        if ":" in k:
+            raise ValueError(f"key {k} can not contain ':'")
+        remove_data.append({"key": k})
+
+    return {
+        "type": "response-rewrite",
+        "yaml": yaml_dump(
+            {
+                "status_code": status_code,
+                "body_base64": body_base64,
+                "body": body,
+                "vars": vars,
+                "headers": {
+                    "add": add_data,
+                    "set": set_data,
+                    "remove": remove_data,
+                }
+            }
+        ),
+    }
+
+
+def build_redirect_validation(
+    uri: str,
+    ret_code: int = 302,
+) -> Dict[str, str]:
+    """generate redirect plugin config
+
+    Args:
+        uri (str): 要重定向到的 URI，可以包含 NGINX 变量。
+        ret_code (int, Optional): HTTP 响应码。
+
+    Raises:
+        ValueError: ret_code cannot be less than 200
+
+    Returns:
+        {
+            "type": "redirect",
+            "yaml": "ret_code: 200\nuri: ''\n"
+        }
+    """
+
+    if ret_code and ret_code < 200:
+        raise ValueError("ret_code cannot be less than 200")
+
+    return {
+        "type": "redirect",
+        "yaml": yaml_dump(
+            {
+                "uri": uri,
+                "ret_code": ret_code,
+            }
+        ),
+    }
+
+
 def _check_percentage(percentage: int, location: str):
     if percentage and not (0 < percentage <= 100):
         raise ValueError(f"The percentage of {location} must be greater than 0 and less than or equal to 100")

sdks/apigw-manager/tests/apigw_manager/plugin/test_config.py

@@ -18,11 +18,14 @@
     build_api_breaker,
     build_fault_injection,
     build_request_validation,
+    build_response_rewrite_validation,
+    build_redirect_validation,
     build_stage_plugin_config_for_definition_yaml,
     UnhealthyConfig,
     HealthyConfig,
     AbortConfig,
     DelayConfig,
+    HeadersConfig,
 )
 
 
@@ -479,6 +482,129 @@ def test_build_request_validation(
 
         assert build_request_validation(body_schema, header_schema, rejected_msg, rejected_code) == expected
 
+    @pytest.mark.parametrize(
+        "status_code, body, vars, body_base64, headers, will_error, expected",
+        [
+            (
+                200,
+                "",
+                "",
+                False,
+                HeadersConfig(
+                    add={"key1": "value1"},
+                    set={"key1": "value1"},
+                    remove=["key1"],
+                ),
+                False,
+                {
+                    "type": "response-rewrite",
+                    "yaml": "body: ''\n"
+                            "body_base64: false\n"
+                            "headers:\n"
+                            "  add:\n"
+                            "  - key: key1:value1\n"
+                            "  remove:\n"
+                            "  - key: key1\n"
+                            "  set:\n"
+                            "  - key: key1\n    value: value1\n"
+                            "status_code: 200\n"
+                            "vars: ''\n"
+                },
+            ),
+            (
+                100,
+                "",
+                "",
+                False,
+                HeadersConfig(
+                    add={"key1": "value1"},
+                    set={"key1": "value1"},
+                    remove=["key1"],
+                ),
+                True,
+                None,
+            ),
+            (
+                200,
+                "",
+                "",
+                False,
+                HeadersConfig(
+                    add={"key1::": "value1"},
+                    set={},
+                    remove=[],
+                ),
+                True,
+                None,
+            ),
+            (
+                200,
+                "",
+                "",
+                False,
+                HeadersConfig(
+                    add={},
+                    set={"key1::": "value1"},
+                    remove=[],
+                ),
+                True,
+                None,
+            ),
+            (
+                200,
+                "",
+                "",
+                False,
+                HeadersConfig(
+                    add={},
+                    set={},
+                    remove=["key1::"],
+                ),
+                True,
+                None,
+            ),
+        ]
+    )
+    def test_build_response_rewrite_validation(
+        self, status_code, body, vars, body_base64, headers, will_error, expected
+    ):
+        if will_error:
+            with pytest.raises(ValueError):
+                build_response_rewrite_validation(status_code, body, vars, body_base64, headers)
+            return
+
+        assert build_response_rewrite_validation(status_code, body, vars, body_base64, headers) == expected
+
+    @pytest.mark.parametrize(
+        "uri, ret_code, will_error, expected",
+        [
+            (
+                "",
+                200,
+                False,
+                {
+                    "type": "redirect",
+                    "yaml": "ret_code: 200\nuri: ''\n"
+                },
+            ),
+            (
+                "",
+                100,
+                True,
+                None,
+            ),
+        ]
+    )
+    def test_build_redirect_validation(
+        self, uri, ret_code, will_error, expected
+    ):
+        if will_error:
+            with pytest.raises(ValueError):
+                build_redirect_validation(uri, ret_code)
+            return
+
+        assert build_redirect_validation(uri, ret_code) == expected
+
 
 class TestBuildStagePluginConfigForDefinitionYaml:
     def test_build_stage_plugin_config_for_definition_yaml(self):