From 3208220456bc3847f649881fa982c35961c458f4 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Tue, 29 Oct 2024 10:36:45 +0000 Subject: [PATCH] parameterize proxy test fixtures --- test_runner/fixtures/neon_fixtures.py | 64 ++++++++++++------- test_runner/regress/test_auth_broker.py | 85 +++++++++++++++++++++++-- 2 files changed, 121 insertions(+), 28 deletions(-) diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index afd7fc2c4c..90d59750d4 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -3590,16 +3590,41 @@ def neon_authorize_jwk() -> jwk.JWK: @pytest.fixture(scope="function") -def static_auth_broker( - port_distributor: PortDistributor, - neon_binpath: Path, - test_output_dir: Path, +def cplane_endpoint_jwks( httpserver: HTTPServer, neon_authorize_jwk: jwk.JWK, - http2_echoserver: H2Server, -) -> Iterable[NeonAuthBroker]: - """Neon Auth Broker that routes to a mocked local_proxy and a mocked cplane HTTP API.""" + role_names: list[str], + audience: str | None, +) -> jwk.JWK: + # return static fixture jwks. + jwk = neon_authorize_jwk.export_public(as_dict=True) + httpserver.expect_request("/authorize/jwks.json").respond_with_json({"keys": [jwk]}) + id = str(uuid.uuid4()) + + # return jwks mock addr on GetEndpointJwks + httpserver.expect_request(re.compile("^/cplane/endpoints/.+/jwks$")).respond_with_json( + { + "jwks": [ + { + "id": id, + "jwks_url": httpserver.url_for("/authorize/jwks.json"), + "provider_name": "test", + "jwt_audience": audience, + "role_names": role_names, + } + ] + } + ) + + return neon_authorize_jwk + + +@pytest.fixture(scope="function") +def cplane_wake_compute_local_proxy( + httpserver: HTTPServer, + http2_echoserver: H2Server, +) -> None: local_proxy_addr = f"{http2_echoserver.host}:{http2_echoserver.port}" # return local_proxy addr on ProxyWakeCompute. @@ -3614,24 +3639,15 @@ def static_auth_broker( } ) - # return jwks mock addr on GetEndpointJwks - httpserver.expect_request(re.compile("^/cplane/endpoints/.+/jwks$")).respond_with_json( - { - "jwks": [ - { - "id": "foo", - "jwks_url": httpserver.url_for("/authorize/jwks.json"), - "provider_name": "test", - "jwt_audience": None, - "role_names": ["anonymous", "authenticated"], - } - ] - } - ) - # return static fixture jwks. - jwk = neon_authorize_jwk.export_public(as_dict=True) - httpserver.expect_request("/authorize/jwks.json").respond_with_json({"keys": [jwk]}) +@pytest.fixture(scope="function") +def static_auth_broker( + port_distributor: PortDistributor, + neon_binpath: Path, + test_output_dir: Path, + httpserver: HTTPServer, +) -> Iterable[NeonAuthBroker]: + """Neon Auth Broker that routes to a mocked local_proxy and a mocked cplane HTTP API.""" mgmt_port = port_distributor.get_port() http_port = port_distributor.get_port() diff --git a/test_runner/regress/test_auth_broker.py b/test_runner/regress/test_auth_broker.py index 11dc7d56b5..eac27fda6a 100644 --- a/test_runner/regress/test_auth_broker.py +++ b/test_runner/regress/test_auth_broker.py @@ -6,19 +6,24 @@ from jwcrypto import jwk, jwt @pytest.mark.asyncio +@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]]) +@pytest.mark.parametrize("audience", [None]) async def test_auth_broker_happy( static_auth_broker: NeonAuthBroker, - neon_authorize_jwk: jwk.JWK, + cplane_wake_compute_local_proxy: None, + cplane_endpoint_jwks: jwk.JWK, ): """ Signs a JWT and uses it to authorize a query to local_proxy. """ token = jwt.JWT( - header={"kid": neon_authorize_jwk.key_id, "alg": "RS256"}, claims={"sub": "user1"} + header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"}, claims={"sub": "user1"} + ) + token.make_signed_token(cplane_endpoint_jwks) + res = await static_auth_broker.query( + "foo", ["arg1"], user="anonymous", token=token.serialize(), expected_code=200 ) - token.make_signed_token(neon_authorize_jwk) - res = await static_auth_broker.query("foo", ["arg1"], user="anonymous", token=token.serialize()) # local proxy mock just echos back the request # check that we forward the correct data @@ -35,3 +40,75 @@ async def test_auth_broker_happy( "query": "foo", "params": ["arg1"], }, "Query body should be forwarded" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]]) +@pytest.mark.parametrize("audience", [None]) +async def test_auth_broker_incorrect_role( + static_auth_broker: NeonAuthBroker, + cplane_wake_compute_local_proxy: None, + cplane_endpoint_jwks: jwk.JWK, +): + """ + Connects to auth broker with the wrong role associated with the JWKs + """ + + token = jwt.JWT( + header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"}, claims={"sub": "user1"} + ) + token.make_signed_token(cplane_endpoint_jwks) + res = await static_auth_broker.query( + "foo", ["arg1"], user="wrong_role", token=token.serialize(), expected_code=400 + ) + + # if the user is wrong, we announce that the jwk was not found. + assert "jwk not found" in res["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]]) +@pytest.mark.parametrize("audience", ["neon"]) +async def test_auth_broker_incorrect_aud( + static_auth_broker: NeonAuthBroker, + cplane_wake_compute_local_proxy: None, + cplane_endpoint_jwks: jwk.JWK, +): + """ + Connects to auth broker with the wrong audience associated with the JWKs + """ + + token = jwt.JWT( + header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"}, + claims={"sub": "user1", "aud": "wrong_aud"}, + ) + token.make_signed_token(cplane_endpoint_jwks) + res = await static_auth_broker.query( + "foo", ["arg1"], user="anonymous", token=token.serialize(), expected_code=400 + ) + + assert "invalid JWT token audience" in res["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]]) +@pytest.mark.parametrize("audience", ["neon"]) +async def test_auth_broker_missing_aud( + static_auth_broker: NeonAuthBroker, + cplane_wake_compute_local_proxy: None, + cplane_endpoint_jwks: jwk.JWK, +): + """ + Connects to auth broker with no audience + """ + + token = jwt.JWT( + header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"}, + claims={"sub": "user1"}, + ) + token.make_signed_token(cplane_endpoint_jwks) + res = await static_auth_broker.query( + "foo", ["arg1"], user="anonymous", token=token.serialize(), expected_code=400 + ) + + assert "invalid JWT token audience" in res["message"]