From f04efbfe7136c253979f6310d190dad3ed04d143 Mon Sep 17 00:00:00 2001
From: Half-Shot <will@half-shot.uk>
Date: Tue, 22 Aug 2023 16:56:22 +0100
Subject: [PATCH] Allowed headers on errors to be conditional based on config
 setting.

---
 synapse/api/errors.py          | 38 +++++++++++++++++++++-------------
 synapse/config/experimental.py |  3 +++
 synapse/http/server.py         | 10 +++++----
 3 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/synapse/api/errors.py b/synapse/api/errors.py
index 4badfcbb9c4c..41501d9defb2 100644
--- a/synapse/api/errors.py
+++ b/synapse/api/errors.py
@@ -127,14 +127,12 @@ class CodeMessageException(RuntimeError):
     Attributes:
         code: HTTP error code
         msg: string describing the error
-        headers: optional response headers to send
     """
 
     def __init__(
         self,
         code: Union[int, HTTPStatus],
         msg: str,
-        headers: Optional[Dict[str, str]] = None,
     ):
         super().__init__("%d: %s" % (code, msg))
 
@@ -146,7 +144,11 @@ def __init__(
         # To eliminate this behaviour, we convert them to their integer equivalents here.
         self.code = int(code)
         self.msg = msg
-        self.headers = headers
+
+    def headers_dict(
+        self, config: Optional["HomeServerConfig"]
+    ) -> Optional[Dict[str, str]]:
+        return None
 
 
 class RedirectException(CodeMessageException):
@@ -192,7 +194,6 @@ def __init__(
         msg: str,
         errcode: str = Codes.UNKNOWN,
         additional_fields: Optional[Dict] = None,
-        headers: Optional[Dict[str, str]] = None,
     ):
         """Constructs a synapse error.
 
@@ -201,7 +202,7 @@ def __init__(
             msg: The human-readable error message.
             errcode: The matrix error code e.g 'M_FORBIDDEN'
         """
-        super().__init__(code, msg, headers)
+        super().__init__(code, msg)
         self.errcode = errcode
         if additional_fields is None:
             self._additional_fields: Dict = {}
@@ -360,11 +361,14 @@ def __init__(
         self,
         required_scopes: List[str],
     ):
-        headers = {
+        self.required_scopes = required_scopes
+        super().__init__(401, "Insufficient scope", Codes.FORBIDDEN, None)
+
+    def headers_dict(self, config: Optional["HomeServerConfig"]) -> Dict[str, str]:
+        return {
             "WWW-Authenticate": 'Bearer error="insufficient_scope", scope="%s"'
-            % (" ".join(required_scopes))
+            % (" ".join(self.required_scopes))
         }
-        super().__init__(401, "Insufficient scope", Codes.FORBIDDEN, None, headers)
 
 
 class UnstableSpecAuthError(AuthError):
@@ -511,17 +515,23 @@ def __init__(
         retry_after_ms: Optional[int] = None,
         errcode: str = Codes.LIMIT_EXCEEDED,
     ):
-        headers = (
-            None
-            if retry_after_ms is None
-            else {"Retry-After": str(math.ceil(retry_after_ms / 1000))}
-        )
-        super().__init__(code, msg, errcode, headers=headers)
+        super().__init__(code, msg, errcode)
         self.retry_after_ms = retry_after_ms
 
     def error_dict(self, config: Optional["HomeServerConfig"]) -> "JsonDict":
         return cs_error(self.msg, self.errcode, retry_after_ms=self.retry_after_ms)
 
+    def headers_dict(
+        self, config: Optional["HomeServerConfig"]
+    ) -> Optional[Dict[str, str]]:
+        if (
+            self.retry_after_ms is not None
+            and config
+            and config.experimental.msc4041_enabled
+        ):
+            return {"Retry-After": str(math.ceil(self.retry_after_ms / 1000))}
+        return None
+
 
 class RoomKeysVersionError(SynapseError):
     """A client has tried to upload to a non-current version of the room_keys store"""
diff --git a/synapse/config/experimental.py b/synapse/config/experimental.py
index ac9449b18f70..508133c70bc0 100644
--- a/synapse/config/experimental.py
+++ b/synapse/config/experimental.py
@@ -398,3 +398,6 @@ def read_config(self, config: JsonDict, **kwargs: Any) -> None:
         self.msc4010_push_rules_account_data = experimental.get(
             "msc4010_push_rules_account_data", False
         )
+
+        # MSC4041: Use http header Retry-After to enable library-assisted retry handling
+        self.msc4041_enabled: bool = experimental.get("msc4041_enabled", False)
diff --git a/synapse/http/server.py b/synapse/http/server.py
index 5109cec983c9..c4af2eca44a2 100644
--- a/synapse/http/server.py
+++ b/synapse/http/server.py
@@ -112,8 +112,9 @@ def return_json_error(
         exc: SynapseError = f.value
         error_code = exc.code
         error_dict = exc.error_dict(config)
-        if exc.headers is not None:
-            for header, value in exc.headers.items():
+        headers_dict = exc.headers_dict(config)
+        if headers_dict is not None:
+            for header, value in headers_dict.items():
                 request.setHeader(header, value)
         logger.info("%s SynapseError: %s - %s", request, error_code, exc.msg)
     elif f.check(CancelledError):
@@ -176,8 +177,9 @@ def return_html_error(
         cme: CodeMessageException = f.value
         code = cme.code
         msg = cme.msg
-        if cme.headers is not None:
-            for header, value in cme.headers.items():
+        headers = cme.headers_dict(None)
+        if headers is not None:
+            for header, value in headers.items():
                 request.setHeader(header, value)
 
         if isinstance(cme, RedirectException):