100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
"""A worker is what its token says, and a superuser can take the token back."""
|
|
|
|
from datetime import timedelta
|
|
|
|
import jwt
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
from fluksio.core import security
|
|
from fluksio.core.config import settings
|
|
from fluksio.flow.remote import PROTOCOL
|
|
|
|
PREFIX = f"{settings.API_V1_STR}/workers"
|
|
MISSING = "sha256:" + "0" * 64
|
|
|
|
|
|
def mint(client: TestClient, headers: dict[str, str]) -> dict[str, str]:
|
|
answer = client.post(
|
|
f"{PREFIX}/tokens", headers=headers, json={"name": "gpu-box", "labels": ["gpu"]}
|
|
)
|
|
assert answer.status_code == 200
|
|
return answer.json()
|
|
|
|
|
|
def hello(*labels: str) -> dict[str, object]:
|
|
return {"op": "hello", "protocol": PROTOCOL, "name": "elsewhere", "labels": labels}
|
|
|
|
|
|
def test_the_token_names_the_worker_and_its_labels(
|
|
client: TestClient, superuser_token_headers: dict[str, str]
|
|
) -> None:
|
|
issued = mint(client, superuser_token_headers)
|
|
with client.websocket_connect(f"{PREFIX}/attach?token={issued['token']}") as ws:
|
|
ws.send_json(hello())
|
|
assert ws.receive_json()["name"] == "gpu-box"
|
|
attached = client.get(PREFIX, headers=superuser_token_headers).json()
|
|
assert [(w["name"], w["labels"]) for w in attached] == [("gpu-box", ["gpu"])]
|
|
|
|
|
|
def test_a_worker_cannot_claim_a_label_its_token_lacks(
|
|
client: TestClient, superuser_token_headers: dict[str, str]
|
|
) -> None:
|
|
issued = mint(client, superuser_token_headers)
|
|
with client.websocket_connect(f"{PREFIX}/attach?token={issued['token']}") as ws:
|
|
ws.send_json(hello("gpu", "vault"))
|
|
assert ws.receive_json()["op"] == "refused"
|
|
|
|
|
|
def test_a_revoked_token_is_dropped_and_refused(
|
|
client: TestClient,
|
|
superuser_token_headers: dict[str, str],
|
|
normal_user_token_headers: dict[str, str],
|
|
) -> None:
|
|
issued = mint(client, superuser_token_headers)
|
|
url = f"{PREFIX}/attach?token={issued['token']}"
|
|
bearer = {"Authorization": f"Bearer {issued['token']}"}
|
|
assert (
|
|
client.get(
|
|
f"{settings.API_V1_STR}/artifacts/{MISSING}", headers=bearer
|
|
).status_code
|
|
== 404
|
|
)
|
|
revoke = f"{PREFIX}/tokens/{issued['jti']}/revoke"
|
|
assert client.post(revoke, headers=normal_user_token_headers).status_code == 403
|
|
|
|
with client.websocket_connect(url) as ws:
|
|
ws.send_json(hello("gpu"))
|
|
ws.receive_json()
|
|
assert client.post(revoke, headers=superuser_token_headers).json()["revoked_at"]
|
|
with pytest.raises(WebSocketDisconnect):
|
|
ws.receive_json()
|
|
|
|
with pytest.raises(WebSocketDisconnect), client.websocket_connect(url):
|
|
pass
|
|
assert (
|
|
client.get(
|
|
f"{settings.API_V1_STR}/artifacts/{MISSING}", headers=bearer
|
|
).status_code
|
|
== 401
|
|
)
|
|
listed = client.get(f"{PREFIX}/tokens", headers=superuser_token_headers).json()
|
|
assert [row for row in listed if row["jti"] == issued["jti"]][0]["revoked_at"]
|
|
assert all("token" not in row for row in listed)
|
|
|
|
|
|
def test_a_token_minted_before_revocation_is_refused(client: TestClient) -> None:
|
|
claims = security.decode_worker_token(
|
|
security.create_worker_token("old", timedelta(days=1))
|
|
)
|
|
del claims["jti"], claims["labels"]
|
|
legacy = jwt.encode(
|
|
claims, security.oauth_key(), algorithm=security.OAUTH_ALGORITHM
|
|
)
|
|
with (
|
|
pytest.raises(WebSocketDisconnect),
|
|
client.websocket_connect(f"{PREFIX}/attach?token={legacy}"),
|
|
):
|
|
pass
|