diff --git a/poetry.lock b/poetry.lock index e1f2e576eb..6bce17008e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -622,8 +622,8 @@ six = ">=1.4.0" websocket-client = ">=0.32.0" [package.extras] +tls = ["idna (>=2.0.0)", "cryptography (>=1.3.4)", "pyOpenSSL (>=17.5.0)"] ssh = ["paramiko (>=2.4.2)"] -tls = ["pyOpenSSL (>=17.5.0)", "cryptography (>=1.3.4)", "idna (>=2.0.0)"] [[package]] name = "ecdsa" @@ -1055,8 +1055,8 @@ optional = false python-versions = ">=3.7" [package.extras] -docs = ["furo (>=2021.7.5b38)", "proselint (>=0.10.2)", "sphinx-autodoc-typehints (>=1.12)", "sphinx (>=4)"] -test = ["appdirs (==1.4.4)", "pytest-cov (>=2.7)", "pytest-mock (>=3.6)", "pytest (>=6)"] +test = ["pytest (>=6)", "pytest-mock (>=3.6)", "pytest-cov (>=2.7)", "appdirs (==1.4.4)"] +docs = ["sphinx (>=4)", "sphinx-autodoc-typehints (>=1.12)", "proselint (>=0.10.2)", "furo (>=2021.7.5b38)"] [[package]] name = "pluggy" @@ -1067,8 +1067,8 @@ optional = false python-versions = ">=3.6" [package.extras] -dev = ["pre-commit", "tox"] -testing = ["pytest", "pytest-benchmark"] +testing = ["pytest-benchmark", "pytest"] +dev = ["tox", "pre-commit"] [[package]] name = "prometheus-client" @@ -1197,6 +1197,20 @@ toml = "*" [package.extras] testing = ["argcomplete", "hypothesis (>=3.56)", "mock", "nose", "requests", "xmlschema"] +[[package]] +name = "pytest-asyncio" +version = "0.19.0" +description = "Pytest support for asyncio" +category = "main" +optional = false +python-versions = ">=3.7" + +[package.dependencies] +pytest = ">=6.1.0" + +[package.extras] +testing = ["pytest-trio (>=0.7.0)", "mypy (>=0.931)", "flaky (>=3.5.0)", "hypothesis (>=5.7.1)", "coverage (>=6.2)"] + [[package]] name = "pytest-forked" version = "1.4.0" @@ -1537,7 +1551,7 @@ testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-flake8", "pytest- [metadata] lock-version = "1.1" python-versions = "^3.9" -content-hash = "2112382a6723ed3b77d242db926c7445fa809fafcf11da127b5292565d2ba798" +content-hash = "badfeff521c68277b10555ab4174847b7315d82818ef5841e600299fb6128698" [metadata.files] aiopg = [ @@ -2076,6 +2090,10 @@ pytest = [ {file = "pytest-6.2.5-py3-none-any.whl", hash = "sha256:7310f8d27bc79ced999e760ca304d69f6ba6c6649c0b60fb0e04a4a77cacc134"}, {file = "pytest-6.2.5.tar.gz", hash = "sha256:131b36680866a76e6781d13f101efb86cf674ebb9762eb70d3082b6f29889e89"}, ] +pytest-asyncio = [ + {file = "pytest-asyncio-0.19.0.tar.gz", hash = "sha256:ac4ebf3b6207259750bc32f4c1d8fcd7e79739edbc67ad0c58dd150b1d072fed"}, + {file = "pytest_asyncio-0.19.0-py3-none-any.whl", hash = "sha256:7a97e37cfe1ed296e2e84941384bdd37c376453912d397ed39293e0916f521fa"}, +] pytest-forked = [ {file = "pytest-forked-1.4.0.tar.gz", hash = "sha256:8b67587c8f98cbbadfdd804539ed5455b6ed03802203485dd2f53c1422d7440e"}, {file = "pytest_forked-1.4.0-py3-none-any.whl", hash = "sha256:bbbb6717efc886b9d64537b41fb1497cfaf3c9601276be8da2cccfea5a3c8ad8"}, diff --git a/pyproject.toml b/pyproject.toml index d648d1050a..2c9270934d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ pytest-timeout = "^2.1.0" Werkzeug = "2.1.2" pytest-order = "^1.0.1" allure-pytest = "^2.9.45" +pytest-asyncio = "^0.19.0" [tool.poetry.dev-dependencies] flake8 = "^5.0.4" diff --git a/test_runner/batch_others/test_proxy.py b/test_runner/batch_others/test_proxy.py index dcff177044..4ffd458b22 100644 --- a/test_runner/batch_others/test_proxy.py +++ b/test_runner/batch_others/test_proxy.py @@ -1,5 +1,11 @@ +import json +import subprocess +from urllib.parse import urlparse + import psycopg2 import pytest +from fixtures.log_helper import log +from fixtures.neon_fixtures import PSQL, NeonProxy, VanillaPostgres def test_proxy_select_1(static_proxy): @@ -23,6 +29,121 @@ def test_password_hack(static_proxy): static_proxy.safe_psql("select 1", sslsni=0, user=user, password=magic) +def get_session_id_from_uri_line(uri_prefix, uri_line): + assert uri_prefix in uri_line + + url_parts = urlparse(uri_line) + psql_session_id = url_parts.path[1:] + assert psql_session_id.isalnum(), "session_id should only contain alphanumeric chars." + link_auth_uri_prefix = uri_line[: -len(url_parts.path)] + # invariant: the prefix must match the uri_prefix. + assert ( + link_auth_uri_prefix == uri_prefix + ), f"Line='{uri_line}' should contain a http auth link of form '{uri_prefix}/'." + # invariant: the entire link_auth_uri should be on its own line, module spaces. + assert " ".join(uri_line.split(" ")) == f"{uri_prefix}/{psql_session_id}" + + return psql_session_id + + +def create_and_send_db_info(local_vanilla_pg, psql_session_id, mgmt_port): + pg_user = "proxy" + pg_password = "password" + + local_vanilla_pg.start() + query = f"create user {pg_user} with login superuser password '{pg_password}'" + local_vanilla_pg.safe_psql(query) + + port = local_vanilla_pg.default_options["port"] + host = local_vanilla_pg.default_options["host"] + dbname = local_vanilla_pg.default_options["dbname"] + + db_info_dict = { + "session_id": psql_session_id, + "result": { + "Success": { + "host": host, + "port": port, + "dbname": dbname, + "user": pg_user, + "password": pg_password, + } + }, + } + db_info_str = json.dumps(db_info_dict) + cmd_args = [ + "psql", + "-h", + "127.0.0.1", # localhost + "-p", + f"{mgmt_port}", + "-c", + db_info_str, + ] + + log.info(f"Sending to proxy the user and db info: {' '.join(cmd_args)}") + p = subprocess.Popen(cmd_args, stdout=subprocess.PIPE) + out, err = p.communicate() + assert "ok" in str(out) + + +async def get_uri_line_from_process_welcome_notice(link_auth_uri_prefix, proc): + """ + Returns the line from the welcome notice from proc containing link_auth_uri_prefix. + :param link_auth_uri_prefix: the uri prefix used to indicate the line of interest + :param proc: the process to read the welcome message from. + :return: a line containing the full link authentication uri. + """ + max_num_lines_of_welcome_message = 15 + for attempt in range(max_num_lines_of_welcome_message): + raw_line = await proc.stderr.readline() + line = raw_line.decode("utf-8").strip() + if link_auth_uri_prefix in line: + return line + assert False, f"did not find line containing '{link_auth_uri_prefix}'" + + +@pytest.mark.asyncio +async def test_psql_session_id(vanilla_pg: VanillaPostgres, link_proxy: NeonProxy): + """ + Test copied and modified from: test_project_psql_link_auth test from cloud/tests_e2e/tests/test_project.py + Step 1. establish connection to the proxy + Step 2. retrieve session_id: + Step 2.1: read welcome message + Step 2.2: parse session_id + Step 3. create a vanilla_pg and send user and db info via command line (using Popen) a psql query via mgmt port to proxy. + Step 4. assert that select 1 has been executed correctly. + """ + + # Step 1. + psql = PSQL( + host=link_proxy.host, + port=link_proxy.proxy_port, + ) + proc = await psql.run("select 1") + + # Step 2.1 + uri_prefix = link_proxy.link_auth_uri_prefix + line_str = await get_uri_line_from_process_welcome_notice(uri_prefix, proc) + + # step 2.2 + psql_session_id = get_session_id_from_uri_line(uri_prefix, line_str) + log.info(f"Parsed psql_session_id='{psql_session_id}' from Neon welcome message.") + + # Step 3. + create_and_send_db_info(vanilla_pg, psql_session_id, link_proxy.mgmt_port) + + # Step 4. + # Expecting proxy output:: + # b' ?column? \n' + # b'----------\n' + # b' 1\n' + # b'(1 row)\n' + out_bytes = await proc.stdout.read() + expected_out_bytes = b" ?column? \n----------\n 1\n(1 row)\n\n" + assert out_bytes == expected_out_bytes + + # Pass extra options to the server. # # Currently, proxy eats the extra connection options, so this fails. diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index f1cffbe5ef..3af0cf4dcb 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -1,6 +1,7 @@ from __future__ import annotations import abc +import asyncio import enum import filecmp import json @@ -1716,21 +1717,58 @@ def remote_pg(test_output_dir: Path) -> Iterator[RemotePostgres]: yield remote_pg +class PSQL: + """ + Helper class to make it easier to run psql in the proxy tests. + Copied and modified from PSQL from cloud/tests_e2e/common/psql.py + """ + + path: str + database_url: str + + def __init__( + self, + path: str = "psql", + host: str = "127.0.0.1", + port: int = 5432, + ): + assert shutil.which(path) + + self.path = path + self.database_url = f"postgres://{host}:{port}/main?options=project%3Dgeneric-project-name" + + async def run(self, query=None): + run_args = [self.path, self.database_url] + run_args += ["--command", query] if query is not None else [] + + cmd_line = subprocess.list2cmdline(run_args) + log.info(f"Run psql: {cmd_line}") + return await asyncio.create_subprocess_exec( + *run_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + + class NeonProxy(PgProtocol): - def __init__(self, proxy_port: int, http_port: int, auth_endpoint: str): + def __init__(self, proxy_port: int, http_port: int, auth_endpoint=None, mgmt_port=None): super().__init__(dsn=auth_endpoint, port=proxy_port) self.host = "127.0.0.1" self.http_port = http_port self.proxy_port = proxy_port + self.mgmt_port = mgmt_port self.auth_endpoint = auth_endpoint self._popen: Optional[subprocess.Popen[bytes]] = None + self.link_auth_uri_prefix = "http://dummy-uri" def start(self) -> None: + """ + Starts a proxy with option '--auth-backend postgres' and a postgres instance already provided though '--auth-endpoint '." + """ assert self._popen is None + assert self.auth_endpoint is not None # Start proxy args = [ - os.path.join(str(neon_binpath), "proxy"), + os.path.join(neon_binpath, "proxy"), *["--http", f"{self.host}:{self.http_port}"], *["--proxy", f"{self.host}:{self.proxy_port}"], *["--auth-backend", "postgres"], @@ -1739,6 +1777,25 @@ class NeonProxy(PgProtocol): self._popen = subprocess.Popen(args) self._wait_until_ready() + def start_with_link_auth(self) -> None: + """ + Starts a proxy with option '--auth-backend link' and a dummy authentication link '--uri dummy-auth-link'." + """ + assert self._popen is None + + # Start proxy + bin_proxy = os.path.join(str(neon_binpath), "proxy") + args = [bin_proxy] + args.extend(["--http", f"{self.host}:{self.http_port}"]) + args.extend(["--proxy", f"{self.host}:{self.proxy_port}"]) + args.extend(["--mgmt", f"{self.host}:{self.mgmt_port}"]) + args.extend(["--auth-backend", "link"]) + args.extend(["--uri", self.link_auth_uri_prefix]) + arg_str = " ".join(args) + log.info(f"starting proxy with command line ::: {arg_str}") + self._popen = subprocess.Popen(args, stdout=subprocess.PIPE) + self._wait_until_ready() + @backoff.on_exception(backoff.expo, requests.exceptions.RequestException, max_time=10) def _wait_until_ready(self): requests.get(f"http://{self.host}:{self.http_port}/v1/status") @@ -1753,6 +1810,17 @@ class NeonProxy(PgProtocol): self._popen.kill() +@pytest.fixture(scope="function") +def link_proxy(port_distributor) -> Iterator[NeonProxy]: + """Neon proxy that routes through link auth.""" + http_port = port_distributor.get_port() + proxy_port = port_distributor.get_port() + mgmt_port = port_distributor.get_port() + with NeonProxy(proxy_port, http_port, mgmt_port=mgmt_port) as proxy: + proxy.start_with_link_auth() + yield proxy + + @pytest.fixture(scope="function") def static_proxy(vanilla_pg, port_distributor) -> Iterator[NeonProxy]: """Neon proxy that routes directly to vanilla postgres."""