rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-01-03 19:42:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
KlimentSerafimov
2022-08-22 20:02:45 -04:00
committed by GitHub
parent 9dd19ec397
commit b98fa5d6b0
4 changed files with 216 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
poetry.lock generated
View File

@@ -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"},

1
pyproject.toml
View File

@@ -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"

121
test_runner/batch_others/test_proxy.py
View File

@@ -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.

72
test_runner/fixtures/neon_fixtures.py
View File

@@ -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."""