stroblme
3a9fc43707
The hello carries the fluksio-worker version. An agent refused for an older protocol stays listed as Needs upgrade until it attaches again. Docs gain an Upgrading workers section.
121 lines
4.3 KiB
Python
121 lines
4.3 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_the_list_names_each_version_and_who_is_behind(
|
|
client: TestClient, superuser_token_headers: dict[str, str]
|
|
) -> None:
|
|
issued = mint(client, superuser_token_headers)
|
|
url = f"{PREFIX}/attach?token={issued['token']}"
|
|
|
|
def listed() -> list[tuple[object, ...]]:
|
|
answer = client.get(PREFIX, headers=superuser_token_headers).json()
|
|
return [(w["protocol"], w["version"], w["behind"]) for w in answer]
|
|
|
|
with client.websocket_connect(url) as ws:
|
|
ws.send_json({**hello(), "protocol": PROTOCOL - 1, "version": "0.1.5"})
|
|
assert ws.receive_json()["op"] == "refused"
|
|
assert listed() == [(PROTOCOL - 1, "0.1.5", True)]
|
|
|
|
with client.websocket_connect(url) as ws:
|
|
ws.send_json({**hello(), "version": "0.2.0"})
|
|
ws.receive_json()
|
|
assert listed() == [(PROTOCOL, "0.2.0", False)]
|
|
|
|
|
|
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
|