Skip to content

middleware

Utilities for middleware response handling.

JsonResponseMiddleware

Bases: ABC

Base class for middleware that transforms JSON response bodies.

Source code in src/stac_auth_proxy/utils/middleware.py 
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class JsonResponseMiddleware(ABC):
    """Base class for middleware that transforms JSON response bodies."""

    app: ASGIApp

    # Expected data type for JSON responses. Only responses matching this type will be transformed.
    # If None, all JSON responses will be transformed regardless of type.
    expected_data_type: Optional[type] = dict

    @abstractmethod
    def should_transform_response(
        self, request: Request, scope: Scope
    ) -> bool:  # mypy: ignore
        """
        Determine if this response should be transformed. At a minimum, this
        should check the request's path and content type.

        Returns
        -------
            bool: True if the response should be transformed

        """
        ...

    @abstractmethod
    def transform_json(self, data: Any, request: Request) -> Any:
        """
        Transform the JSON data.

        Args:
            data: The parsed JSON data
            request: The HTTP request object

        Returns:
        -------
            The transformed JSON data

        """
        ...

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Process the request/response."""
        if scope["type"] != "http":
            return await self.app(scope, receive, send)

        start_message: Optional[Message] = None
        body = b""

        async def transform_response(message: Message) -> None:
            nonlocal start_message
            nonlocal body

            start_message = start_message or message
            headers = MutableHeaders(scope=start_message)
            request = Request(scope)

            if not self.should_transform_response(
                request=request,
                scope=start_message,
            ):
                # For non-JSON responses, send the start message immediately
                await send(message)
                return

            # Delay sending start message until we've processed the body
            if message["type"] == "http.response.start":
                return

            body += message["body"]

            # Skip body chunks until all chunks have been received
            if message.get("more_body"):
                return

            # Transform the JSON body
            if body:
                try:
                    data = json.loads(body)
                except json.JSONDecodeError as e:
                    logger.error("Error parsing JSON: %s", e)
                    logger.error("Body: %s", body)
                    logger.error("Response scope: %s", scope)
                    response = JSONResponse(
                        {"error": "Received invalid JSON from upstream server"},
                        status_code=502,
                    )
                    await response(scope, receive, send)
                    return

                if self.expected_data_type is None or isinstance(
                    data, self.expected_data_type
                ):
                    transformed = self.transform_json(data, request=request)
                    body = json.dumps(transformed).encode()
                else:
                    logger.warning(
                        "Received JSON response with unexpected data type %r from upstream server (%r %r), "
                        "skipping transformation (expected: %r)",
                        type(data).__name__,
                        request.method,
                        request.url,
                        self.expected_data_type.__name__,
                    )

            # Update content-length header
            headers["content-length"] = str(len(body))
            start_message["headers"] = [
                (key.encode(), value.encode()) for key, value in headers.items()
            ]

            # Send response
            await send(start_message)
            await send(
                {
                    "type": "http.response.body",
                    "body": body,
                    "more_body": False,
                }
            )

        return await self.app(scope, receive, transform_response)

__call__(scope: Scope, receive: Receive, send: Send) -> None async

Process the request/response.

Source code in src/stac_auth_proxy/utils/middleware.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
    """Process the request/response."""
    if scope["type"] != "http":
        return await self.app(scope, receive, send)

    start_message: Optional[Message] = None
    body = b""

    async def transform_response(message: Message) -> None:
        nonlocal start_message
        nonlocal body

        start_message = start_message or message
        headers = MutableHeaders(scope=start_message)
        request = Request(scope)

        if not self.should_transform_response(
            request=request,
            scope=start_message,
        ):
            # For non-JSON responses, send the start message immediately
            await send(message)
            return

        # Delay sending start message until we've processed the body
        if message["type"] == "http.response.start":
            return

        body += message["body"]

        # Skip body chunks until all chunks have been received
        if message.get("more_body"):
            return

        # Transform the JSON body
        if body:
            try:
                data = json.loads(body)
            except json.JSONDecodeError as e:
                logger.error("Error parsing JSON: %s", e)
                logger.error("Body: %s", body)
                logger.error("Response scope: %s", scope)
                response = JSONResponse(
                    {"error": "Received invalid JSON from upstream server"},
                    status_code=502,
                )
                await response(scope, receive, send)
                return

            if self.expected_data_type is None or isinstance(
                data, self.expected_data_type
            ):
                transformed = self.transform_json(data, request=request)
                body = json.dumps(transformed).encode()
            else:
                logger.warning(
                    "Received JSON response with unexpected data type %r from upstream server (%r %r), "
                    "skipping transformation (expected: %r)",
                    type(data).__name__,
                    request.method,
                    request.url,
                    self.expected_data_type.__name__,
                )

        # Update content-length header
        headers["content-length"] = str(len(body))
        start_message["headers"] = [
            (key.encode(), value.encode()) for key, value in headers.items()
        ]

        # Send response
        await send(start_message)
        await send(
            {
                "type": "http.response.body",
                "body": body,
                "more_body": False,
            }
        )

    return await self.app(scope, receive, transform_response)

should_transform_response(request: Request, scope: Scope) -> bool abstractmethod

Determine if this response should be transformed. At a minimum, this should check the request's path and content type.

Returns
bool: True if the response should be transformed
Source code in src/stac_auth_proxy/utils/middleware.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
@abstractmethod
def should_transform_response(
    self, request: Request, scope: Scope
) -> bool:  # mypy: ignore
    """
    Determine if this response should be transformed. At a minimum, this
    should check the request's path and content type.

    Returns
    -------
        bool: True if the response should be transformed

    """
    ...

transform_json(data: Any, request: Request) -> Any abstractmethod

Transform the JSON data.

Parameters:

Name Type Description Default
data Any

The parsed JSON data

 required
request Request

The HTTP request object

 required 
The transformed JSON data
Source code in src/stac_auth_proxy/utils/middleware.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
@abstractmethod
def transform_json(self, data: Any, request: Request) -> Any:
    """
    Transform the JSON data.

    Args:
        data: The parsed JSON data
        request: The HTTP request object

    Returns:
    -------
        The transformed JSON data

    """
    ...

required_conformance(*conformances: str, attr_name: str = '__required_conformances__')

Register required conformance classes with a middleware class.

Source code in src/stac_auth_proxy/utils/middleware.py 
139
140
141
142
143
144
145
146
147
148
149
def required_conformance(
    *conformances: str,
    attr_name: str = "__required_conformances__",
):
    """Register required conformance classes with a middleware class."""

    def decorator(middleware):
        setattr(middleware, attr_name, list(conformances))
        return middleware

    return decorator