amoghrajesh commented on code in PR #69178:
URL: https://github.com/apache/airflow/pull/69178#discussion_r3505960090
##########
task-sdk/tests/task_sdk/api/test_client.py:
##########
@@ -1831,6 +1831,34 @@ def handle_request(request: httpx.Request) ->
httpx.Response:
assert isinstance(result, TaskStateStoreResponse)
assert result.value == "spark_app_001"
+ def test_get_url_quotes_key_as_single_path_segment(self):
+ """Test that task state store keys cannot escape the store API path."""
+ requests_seen = []
+ crafted_key = "../../../variables/secret_key"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ requests_seen.append(request)
+ if request.url.path == "/variables/secret_key":
+ return httpx.Response(
+ status_code=200,
+ json={"key": "secret_key", "value": "super-secret-value"},
+ )
+ return httpx.Response(
+ status_code=404,
+ json={"detail": {"reason": "not_found", "message": "Task state
key not found"}},
+ )
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ result = client.task_state_store.get(ti_id=self.TI_ID, key=crafted_key)
+
+ assert (
+ requests_seen[0].url.raw_path
+ ==
f"/store/ti/{self.TI_ID}/..%2F..%2F..%2Fvariables%2Fsecret_key".encode()
+ )
+ assert requests_seen[0].url.path != "/variables/secret_key"
+ assert isinstance(result, ErrorResponse)
+ assert "super-secret-value" not in str(result)
Review Comment:
Doesn't have to be so complex, can simply be similar to:
```python
def test_get_url_encodes_key_with_slash(self):
requests_seen = []
def handle_request(request: httpx.Request) -> httpx.Response:
requests_seen.append(request)
return httpx.Response(
status_code=200,
json={"value": "spark_app_001"},
)
client = make_client(transport=httpx.MockTransport(handle_request))
client.task_state_store.get(ti_id=self.TI_ID, key="spark/job_id")
assert b"%2F" in requests_seen[0].url.raw_path
assert requests_seen[0].url.raw_path ==
f"/store/ti/{self.TI_ID}/spark%2Fjob_id".encode()
```
##########
task-sdk/tests/task_sdk/api/test_client.py:
##########
@@ -1831,6 +1831,34 @@ def handle_request(request: httpx.Request) ->
httpx.Response:
assert isinstance(result, TaskStateStoreResponse)
assert result.value == "spark_app_001"
+ def test_get_url_quotes_key_as_single_path_segment(self):
+ """Test that task state store keys cannot escape the store API path."""
+ requests_seen = []
+ crafted_key = "../../../variables/secret_key"
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ requests_seen.append(request)
+ if request.url.path == "/variables/secret_key":
+ return httpx.Response(
+ status_code=200,
+ json={"key": "secret_key", "value": "super-secret-value"},
+ )
+ return httpx.Response(
+ status_code=404,
+ json={"detail": {"reason": "not_found", "message": "Task state
key not found"}},
+ )
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ result = client.task_state_store.get(ti_id=self.TI_ID, key=crafted_key)
+
+ assert (
+ requests_seen[0].url.raw_path
+ ==
f"/store/ti/{self.TI_ID}/..%2F..%2F..%2Fvariables%2Fsecret_key".encode()
+ )
+ assert requests_seen[0].url.path != "/variables/secret_key"
+ assert isinstance(result, ErrorResponse)
+ assert "super-secret-value" not in str(result)
Review Comment:
Same for the other tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
