Added a new test for making sure the proxy displays a session_id when using link auth. (#2039)

Added pytest to check correctness of the link authentication pipeline.

Context: this PR is the first step towards refactoring the link authentication pipeline to use https (instead of psql) to send the db info to the proxy. There was a test missing for this pipeline in this repo, so this PR adds that test as preparation for the actual change of psql -> https.
Co-authored-by: Bojan Serafimov <bojan.serafimov7@gmail.com>
Co-authored-by: Dmitry Rodionov <dmitry@neon.tech>
Co-authored-by: Stas Kelvic <stas@neon.tech>
Co-authored-by: Dimitrii Ivanov <dima@neon.tech>

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}/<psql_session_id>'."
+    # 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.

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 <postgress-instance>'."
+        """
         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."""