v3-0008-Add-pytest-suite-for-OAuth.patch
text/x-patch
Filename: v3-0008-Add-pytest-suite-for-OAuth.patch
Type: text/x-patch
Part: 7
Patch
Same data as JSON:
GET /api/v1/attachments/:id/patch
the parsed metadata as JSON — format, series position, per-file stats; never the diff bytes.
API reference →
Format: format-patch
Series: patch v3-0008
Subject: Add pytest suite for OAuth
| File | + | − |
|---|---|---|
| src/test/python/client/conftest.py | 126 | 0 |
| src/test/python/client/__init__.py | 0 | 0 |
| src/test/python/client/test_client.py | 180 | 0 |
| src/test/python/client/test_oauth.py | 936 | 0 |
| src/test/python/.gitignore | 2 | 0 |
| src/test/python/Makefile | 38 | 0 |
| src/test/python/pq3.py | 727 | 0 |
| src/test/python/pytest.ini | 4 | 0 |
| src/test/python/README | 54 | 0 |
| src/test/python/requirements.txt | 7 | 0 |
| src/test/python/server/conftest.py | 45 | 0 |
| src/test/python/server/__init__.py | 0 | 0 |
| src/test/python/server/test_oauth.py | 1012 | 0 |
| src/test/python/server/test_server.py | 21 | 0 |
| src/test/python/server/validate_bearer.py | 101 | 0 |
| src/test/python/server/validate_reflect.py | 34 | 0 |
| src/test/python/test_internals.py | 138 | 0 |
| src/test/python/test_pq3.py | 558 | 0 |
| src/test/python/tls.py | 195 | 0 |
From 1c24881d4b1e8777ce176d2c276fe8120bd6e648 Mon Sep 17 00:00:00 2001
From: Jacob Champion <pchampion@vmware.com>
Date: Fri, 4 Jun 2021 09:06:38 -0700
Subject: [PATCH v3 8/9] Add pytest suite for OAuth
Requires Python 3; on the first run of `make installcheck` the
dependencies will be installed into ./venv for you. See the README for
more details.
---
src/test/python/.gitignore | 2 +
src/test/python/Makefile | 38 +
src/test/python/README | 54 ++
src/test/python/client/__init__.py | 0
src/test/python/client/conftest.py | 126 +++
src/test/python/client/test_client.py | 180 ++++
src/test/python/client/test_oauth.py | 936 ++++++++++++++++++
src/test/python/pq3.py | 727 ++++++++++++++
src/test/python/pytest.ini | 4 +
src/test/python/requirements.txt | 7 +
src/test/python/server/__init__.py | 0
src/test/python/server/conftest.py | 45 +
src/test/python/server/test_oauth.py | 1012 ++++++++++++++++++++
src/test/python/server/test_server.py | 21 +
src/test/python/server/validate_bearer.py | 101 ++
src/test/python/server/validate_reflect.py | 34 +
src/test/python/test_internals.py | 138 +++
src/test/python/test_pq3.py | 558 +++++++++++
src/test/python/tls.py | 195 ++++
19 files changed, 4178 insertions(+)
create mode 100644 src/test/python/.gitignore
create mode 100644 src/test/python/Makefile
create mode 100644 src/test/python/README
create mode 100644 src/test/python/client/__init__.py
create mode 100644 src/test/python/client/conftest.py
create mode 100644 src/test/python/client/test_client.py
create mode 100644 src/test/python/client/test_oauth.py
create mode 100644 src/test/python/pq3.py
create mode 100644 src/test/python/pytest.ini
create mode 100644 src/test/python/requirements.txt
create mode 100644 src/test/python/server/__init__.py
create mode 100644 src/test/python/server/conftest.py
create mode 100644 src/test/python/server/test_oauth.py
create mode 100644 src/test/python/server/test_server.py
create mode 100755 src/test/python/server/validate_bearer.py
create mode 100755 src/test/python/server/validate_reflect.py
create mode 100644 src/test/python/test_internals.py
create mode 100644 src/test/python/test_pq3.py
create mode 100644 src/test/python/tls.py
diff --git a/src/test/python/.gitignore b/src/test/python/.gitignore
new file mode 100644
index 0000000000..0e8f027b2e
--- /dev/null
+++ b/src/test/python/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+/venv/
diff --git a/src/test/python/Makefile b/src/test/python/Makefile
new file mode 100644
index 0000000000..b0695b6287
--- /dev/null
+++ b/src/test/python/Makefile
@@ -0,0 +1,38 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+# Only Python 3 is supported, but if it's named something different on your
+# system you can override it with the PYTHON3 variable.
+PYTHON3 := python3
+
+# All dependencies are placed into this directory. The default is .gitignored
+# for you, but you can override it if you'd like.
+VENV := ./venv
+
+override VBIN := $(VENV)/bin
+override PIP := $(VBIN)/pip
+override PYTEST := $(VBIN)/py.test
+override ISORT := $(VBIN)/isort
+override BLACK := $(VBIN)/black
+
+.PHONY: installcheck indent
+
+installcheck: $(PYTEST)
+ $(PYTEST) -v -rs
+
+indent: $(ISORT) $(BLACK)
+ $(ISORT) --profile black *.py client/*.py server/*.py
+ $(BLACK) *.py client/*.py server/*.py
+
+$(PYTEST) $(ISORT) $(BLACK) &: requirements.txt | $(PIP)
+ $(PIP) install --force-reinstall -r $<
+
+$(PIP):
+ $(PYTHON3) -m venv $(VENV)
+
+# A convenience recipe to rebuild psycopg2 against the local libpq.
+.PHONY: rebuild-psycopg2
+rebuild-psycopg2: | $(PIP)
+ $(PIP) install --force-reinstall --no-binary :all: $(shell grep psycopg2 requirements.txt)
diff --git a/src/test/python/README b/src/test/python/README
new file mode 100644
index 0000000000..0bda582c4b
--- /dev/null
+++ b/src/test/python/README
@@ -0,0 +1,54 @@
+A test suite for exercising both the libpq client and the server backend at the
+protocol level, based on pytest and Construct.
+
+The test suite currently assumes that the standard PG* environment variables
+point to the database under test and are sufficient to log in a superuser on
+that system. In other words, a bare `psql` needs to Just Work before the test
+suite can do its thing. For a newly built dev cluster, typically all that I need
+to do is a
+
+ export PGDATABASE=postgres
+
+but you can adjust as needed for your setup.
+
+## Requirements
+
+A supported version (3.6+) of Python.
+
+The first run of
+
+ make installcheck
+
+will install a local virtual environment and all needed dependencies. During
+development, if libpq changes incompatibly, you can issue
+
+ $ make rebuild-psycopg2
+
+to force a rebuild of the client library.
+
+## Hacking
+
+The code style is enforced by a _very_ opinionated autoformatter. Running the
+
+ make indent
+
+recipe will invoke it for you automatically. Don't fight the tool; part of the
+zen is in knowing that if the formatter makes your code ugly, there's probably a
+cleaner way to write your code.
+
+## Advanced Usage
+
+The Makefile is there for convenience, but you don't have to use it. Activate
+the virtualenv to be able to use pytest directly:
+
+ $ source venv/bin/activate
+ $ py.test -k oauth
+ ...
+ $ py.test ./server/test_server.py
+ ...
+ $ deactivate # puts the PATH et al back the way it was before
+
+To make quick smoke tests possible, slow tests have been marked explicitly. You
+can skip them by saying e.g.
+
+ $ py.test -m 'not slow'
diff --git a/src/test/python/client/__init__.py b/src/test/python/client/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/test/python/client/conftest.py b/src/test/python/client/conftest.py
new file mode 100644
index 0000000000..f38da7a138
--- /dev/null
+++ b/src/test/python/client/conftest.py
@@ -0,0 +1,126 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import socket
+import sys
+import threading
+
+import psycopg2
+import pytest
+
+import pq3
+
+BLOCKING_TIMEOUT = 2 # the number of seconds to wait for blocking calls
+
+
+@pytest.fixture
+def server_socket(unused_tcp_port_factory):
+ """
+ Returns a listening socket bound to an ephemeral port.
+ """
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+ s.bind(("127.0.0.1", unused_tcp_port_factory()))
+ s.listen(1)
+ s.settimeout(BLOCKING_TIMEOUT)
+ yield s
+
+
+class ClientHandshake(threading.Thread):
+ """
+ A thread that connects to a local Postgres server using psycopg2. Once the
+ opening handshake completes, the connection will be immediately closed.
+ """
+
+ def __init__(self, *, port, **kwargs):
+ super().__init__()
+
+ kwargs["port"] = port
+ self._kwargs = kwargs
+
+ self.exception = None
+
+ def run(self):
+ try:
+ conn = psycopg2.connect(host="127.0.0.1", **self._kwargs)
+ conn.close()
+ except Exception as e:
+ self.exception = e
+
+ def check_completed(self, timeout=BLOCKING_TIMEOUT):
+ """
+ Joins the client thread. Raises an exception if the thread could not be
+ joined, or if it threw an exception itself. (The exception will be
+ cleared, so future calls to check_completed will succeed.)
+ """
+ self.join(timeout)
+
+ if self.is_alive():
+ raise TimeoutError("client thread did not handshake within the timeout")
+ elif self.exception:
+ e = self.exception
+ self.exception = None
+ raise e
+
+
+@pytest.fixture
+def accept(server_socket):
+ """
+ Returns a factory function that, when called, returns a pair (sock, client)
+ where sock is a server socket that has accepted a connection from client,
+ and client is an instance of ClientHandshake. Clients will complete their
+ handshakes and cleanly disconnect.
+
+ The default connstring options may be extended or overridden by passing
+ arbitrary keyword arguments. Keep in mind that you generally should not
+ override the host or port, since they point to the local test server.
+
+ For situations where a client needs to connect more than once to complete a
+ handshake, the accept function may be called more than once. (The client
+ returned for subsequent calls will always be the same client that was
+ returned for the first call.)
+
+ Tests must either complete the handshake so that the client thread can be
+ automatically joined during teardown, or else call client.check_completed()
+ and manually handle any expected errors.
+ """
+ _, port = server_socket.getsockname()
+
+ client = None
+ default_opts = dict(
+ port=port,
+ user=pq3.pguser(),
+ sslmode="disable",
+ )
+
+ def factory(**kwargs):
+ nonlocal client
+
+ if client is None:
+ opts = dict(default_opts)
+ opts.update(kwargs)
+
+ # The server_socket is already listening, so the client thread can
+ # be safely started; it'll block on the connection until we accept.
+ client = ClientHandshake(**opts)
+ client.start()
+
+ sock, _ = server_socket.accept()
+ return sock, client
+
+ yield factory
+ client.check_completed()
+
+
+@pytest.fixture
+def conn(accept):
+ """
+ Returns an accepted, wrapped pq3 connection to a psycopg2 client. The socket
+ will be closed when the test finishes, and the client will be checked for a
+ cleanly completed handshake.
+ """
+ sock, client = accept()
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ yield conn
diff --git a/src/test/python/client/test_client.py b/src/test/python/client/test_client.py
new file mode 100644
index 0000000000..c4c946fda4
--- /dev/null
+++ b/src/test/python/client/test_client.py
@@ -0,0 +1,180 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import base64
+import sys
+
+import psycopg2
+import pytest
+from cryptography.hazmat.primitives import hashes, hmac
+
+import pq3
+
+
+def finish_handshake(conn):
+ """
+ Sends the AuthenticationOK message and the standard opening salvo of server
+ messages, then asserts that the client immediately sends a Terminate message
+ to close the connection cleanly.
+ """
+ pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.OK)
+ pq3.send(conn, pq3.types.ParameterStatus, name=b"client_encoding", value=b"UTF-8")
+ pq3.send(conn, pq3.types.ParameterStatus, name=b"DateStyle", value=b"ISO, MDY")
+ pq3.send(conn, pq3.types.ReadyForQuery, status=b"I")
+
+ pkt = pq3.recv1(conn)
+ assert pkt.type == pq3.types.Terminate
+
+
+def test_handshake(conn):
+ startup = pq3.recv1(conn, cls=pq3.Startup)
+ assert startup.proto == pq3.protocol(3, 0)
+
+ finish_handshake(conn)
+
+
+def test_aborted_connection(accept):
+ """
+ Make sure the client correctly reports an early close during handshakes.
+ """
+ sock, client = accept()
+ sock.close()
+
+ expected = "server closed the connection unexpectedly"
+ with pytest.raises(psycopg2.OperationalError, match=expected):
+ client.check_completed()
+
+
+#
+# SCRAM-SHA-256 (see RFC 5802: https://tools.ietf.org/html/rfc5802)
+#
+
+
+@pytest.fixture
+def password():
+ """
+ Returns a password for use by both client and server.
+ """
+ # TODO: parameterize this with passwords that require SASLprep.
+ return "secret"
+
+
+@pytest.fixture
+def pwconn(accept, password):
+ """
+ Like the conn fixture, but uses a password in the connection.
+ """
+ sock, client = accept(password=password)
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ yield conn
+
+
+def sha256(data):
+ """The H(str) function from Section 2.2."""
+ digest = hashes.Hash(hashes.SHA256())
+ digest.update(data)
+ return digest.finalize()
+
+
+def hmac_256(key, data):
+ """The HMAC(key, str) function from Section 2.2."""
+ h = hmac.HMAC(key, hashes.SHA256())
+ h.update(data)
+ return h.finalize()
+
+
+def xor(a, b):
+ """The XOR operation from Section 2.2."""
+ res = bytearray(a)
+ for i, byte in enumerate(b):
+ res[i] ^= byte
+ return bytes(res)
+
+
+def h_i(data, salt, i):
+ """The Hi(str, salt, i) function from Section 2.2."""
+ assert i > 0
+
+ acc = hmac_256(data, salt + b"\x00\x00\x00\x01")
+ last = acc
+ i -= 1
+
+ while i:
+ u = hmac_256(data, last)
+ acc = xor(acc, u)
+
+ last = u
+ i -= 1
+
+ return acc
+
+
+def test_scram(pwconn, password):
+ startup = pq3.recv1(pwconn, cls=pq3.Startup)
+ assert startup.proto == pq3.protocol(3, 0)
+
+ pq3.send(
+ pwconn,
+ pq3.types.AuthnRequest,
+ type=pq3.authn.SASL,
+ body=[b"SCRAM-SHA-256", b""],
+ )
+
+ # Get the client-first-message.
+ pkt = pq3.recv1(pwconn)
+ assert pkt.type == pq3.types.PasswordMessage
+
+ initial = pq3.SASLInitialResponse.parse(pkt.payload)
+ assert initial.name == b"SCRAM-SHA-256"
+
+ c_bind, authzid, c_name, c_nonce = initial.data.split(b",")
+ assert c_bind == b"n" # no channel bindings on a plaintext connection
+ assert authzid == b"" # we don't support authzid currently
+ assert c_name == b"n=" # libpq doesn't honor the GS2 username
+ assert c_nonce.startswith(b"r=")
+
+ # Send the server-first-message.
+ salt = b"12345"
+ iterations = 2
+
+ s_nonce = c_nonce + b"somenonce"
+ s_salt = b"s=" + base64.b64encode(salt)
+ s_iterations = b"i=%d" % iterations
+
+ msg = b",".join([s_nonce, s_salt, s_iterations])
+ pq3.send(pwconn, pq3.types.AuthnRequest, type=pq3.authn.SASLContinue, body=msg)
+
+ # Get the client-final-message.
+ pkt = pq3.recv1(pwconn)
+ assert pkt.type == pq3.types.PasswordMessage
+
+ c_bind_final, c_nonce_final, c_proof = pkt.payload.split(b",")
+ assert c_bind_final == b"c=" + base64.b64encode(c_bind + b"," + authzid + b",")
+ assert c_nonce_final == s_nonce
+
+ # Calculate what the client proof should be.
+ salted_password = h_i(password.encode("ascii"), salt, iterations)
+ client_key = hmac_256(salted_password, b"Client Key")
+ stored_key = sha256(client_key)
+
+ auth_message = b",".join(
+ [c_name, c_nonce, s_nonce, s_salt, s_iterations, c_bind_final, c_nonce_final]
+ )
+ client_signature = hmac_256(stored_key, auth_message)
+ client_proof = xor(client_key, client_signature)
+
+ expected = b"p=" + base64.b64encode(client_proof)
+ assert c_proof == expected
+
+ # Send the correct server signature.
+ server_key = hmac_256(salted_password, b"Server Key")
+ server_signature = hmac_256(server_key, auth_message)
+
+ s_verify = b"v=" + base64.b64encode(server_signature)
+ pq3.send(pwconn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal, body=s_verify)
+
+ # Done!
+ finish_handshake(pwconn)
diff --git a/src/test/python/client/test_oauth.py b/src/test/python/client/test_oauth.py
new file mode 100644
index 0000000000..a754a9c0b6
--- /dev/null
+++ b/src/test/python/client/test_oauth.py
@@ -0,0 +1,936 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import base64
+import http.server
+import json
+import secrets
+import sys
+import threading
+import time
+import urllib.parse
+
+import psycopg2
+import pytest
+
+import pq3
+
+from .conftest import BLOCKING_TIMEOUT
+
+
+def finish_handshake(conn):
+ """
+ Sends the AuthenticationOK message and the standard opening salvo of server
+ messages, then asserts that the client immediately sends a Terminate message
+ to close the connection cleanly.
+ """
+ pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.OK)
+ pq3.send(conn, pq3.types.ParameterStatus, name=b"client_encoding", value=b"UTF-8")
+ pq3.send(conn, pq3.types.ParameterStatus, name=b"DateStyle", value=b"ISO, MDY")
+ pq3.send(conn, pq3.types.ReadyForQuery, status=b"I")
+
+ pkt = pq3.recv1(conn)
+ assert pkt.type == pq3.types.Terminate
+
+
+#
+# OAUTHBEARER (see RFC 7628: https://tools.ietf.org/html/rfc7628)
+#
+
+
+def start_oauth_handshake(conn):
+ """
+ Negotiates an OAUTHBEARER SASL challenge. Returns the client's initial
+ response data.
+ """
+ startup = pq3.recv1(conn, cls=pq3.Startup)
+ assert startup.proto == pq3.protocol(3, 0)
+
+ pq3.send(
+ conn, pq3.types.AuthnRequest, type=pq3.authn.SASL, body=[b"OAUTHBEARER", b""]
+ )
+
+ pkt = pq3.recv1(conn)
+ assert pkt.type == pq3.types.PasswordMessage
+
+ initial = pq3.SASLInitialResponse.parse(pkt.payload)
+ assert initial.name == b"OAUTHBEARER"
+
+ return initial.data
+
+
+def get_auth_value(initial):
+ """
+ Finds the auth value (e.g. "Bearer somedata..." in the client's initial SASL
+ response.
+ """
+ kvpairs = initial.split(b"\x01")
+ assert kvpairs[0] == b"n,," # no channel binding or authzid
+ assert kvpairs[2] == b"" # ends with an empty kvpair
+ assert kvpairs[3] == b"" # ...and there's nothing after it
+ assert len(kvpairs) == 4
+
+ key, value = kvpairs[1].split(b"=", 2)
+ assert key == b"auth"
+
+ return value
+
+
+def xtest_oauth_success(conn): # TODO
+ initial = start_oauth_handshake(conn)
+
+ auth = get_auth_value(initial)
+ assert auth.startswith(b"Bearer ")
+
+ # Accept the token. TODO actually validate
+ pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+ finish_handshake(conn)
+
+
+class OpenIDProvider(threading.Thread):
+ """
+ A thread that runs a mock OpenID provider server.
+ """
+
+ def __init__(self, *, port):
+ super().__init__()
+
+ self.exception = None
+
+ addr = ("", port)
+ self.server = self._Server(addr, self._Handler)
+
+ # TODO: allow HTTPS only, somehow
+ oauth = self._OAuthState()
+ oauth.host = f"localhost:{port}"
+ oauth.issuer = f"http://localhost:{port}"
+
+ # The following endpoints are required to be advertised by providers,
+ # even though our chosen client implementation does not actually make
+ # use of them.
+ oauth.register_endpoint(
+ "authorization_endpoint", "POST", "/authorize", self._authorization_handler
+ )
+ oauth.register_endpoint("jwks_uri", "GET", "/keys", self._jwks_handler)
+
+ self.server.oauth = oauth
+
+ def run(self):
+ try:
+ self.server.serve_forever()
+ except Exception as e:
+ self.exception = e
+
+ def stop(self, timeout=BLOCKING_TIMEOUT):
+ """
+ Shuts down the server and joins its thread. Raises an exception if the
+ thread could not be joined, or if it threw an exception itself. Must
+ only be called once, after start().
+ """
+ self.server.shutdown()
+ self.join(timeout)
+
+ if self.is_alive():
+ raise TimeoutError("client thread did not handshake within the timeout")
+ elif self.exception:
+ e = self.exception
+ raise e
+
+ class _OAuthState(object):
+ def __init__(self):
+ self.endpoint_paths = {}
+ self._endpoints = {}
+
+ def register_endpoint(self, name, method, path, func):
+ if method not in self._endpoints:
+ self._endpoints[method] = {}
+
+ self._endpoints[method][path] = func
+ self.endpoint_paths[name] = path
+
+ def endpoint(self, method, path):
+ if method not in self._endpoints:
+ return None
+
+ return self._endpoints[method].get(path)
+
+ class _Server(http.server.HTTPServer):
+ def handle_error(self, request, addr):
+ self.shutdown_request(request)
+ raise
+
+ @staticmethod
+ def _jwks_handler(headers, params):
+ return 200, {"keys": []}
+
+ @staticmethod
+ def _authorization_handler(headers, params):
+ # We don't actually want this to be called during these tests -- we
+ # should be using the device authorization endpoint instead.
+ assert (
+ False
+ ), "authorization handler called instead of device authorization handler"
+
+ class _Handler(http.server.BaseHTTPRequestHandler):
+ timeout = BLOCKING_TIMEOUT
+
+ def _discovery_handler(self, headers, params):
+ oauth = self.server.oauth
+
+ doc = {
+ "issuer": oauth.issuer,
+ "response_types_supported": ["token"],
+ "subject_types_supported": ["public"],
+ "id_token_signing_alg_values_supported": ["RS256"],
+ }
+
+ for name, path in oauth.endpoint_paths.items():
+ doc[name] = oauth.issuer + path
+
+ return 200, doc
+
+ def _handle(self, *, params=None, handler=None):
+ oauth = self.server.oauth
+ assert self.headers["Host"] == oauth.host
+
+ if handler is None:
+ handler = oauth.endpoint(self.command, self.path)
+ assert (
+ handler is not None
+ ), f"no registered endpoint for {self.command} {self.path}"
+
+ code, resp = handler(self.headers, params)
+
+ self.send_response(code)
+ self.send_header("Content-Type", "application/json")
+ self.end_headers()
+
+ resp = json.dumps(resp)
+ resp = resp.encode("utf-8")
+ self.wfile.write(resp)
+
+ self.close_connection = True
+
+ def do_GET(self):
+ if self.path == "/.well-known/openid-configuration":
+ self._handle(handler=self._discovery_handler)
+ return
+
+ self._handle()
+
+ def _request_body(self):
+ length = self.headers["Content-Length"]
+
+ # Handle only an explicit content-length.
+ assert length is not None
+ length = int(length)
+
+ return self.rfile.read(length).decode("utf-8")
+
+ def do_POST(self):
+ assert self.headers["Content-Type"] == "application/x-www-form-urlencoded"
+
+ body = self._request_body()
+ params = urllib.parse.parse_qs(body)
+
+ self._handle(params=params)
+
+
+@pytest.fixture
+def openid_provider(unused_tcp_port_factory):
+ """
+ A fixture that returns the OAuth state of a running OpenID provider server. The
+ server will be stopped when the fixture is torn down.
+ """
+ thread = OpenIDProvider(port=unused_tcp_port_factory())
+ thread.start()
+
+ try:
+ yield thread.server.oauth
+ finally:
+ thread.stop()
+
+
+@pytest.mark.parametrize("secret", [None, "", "hunter2"])
+@pytest.mark.parametrize("scope", [None, "", "openid email"])
+@pytest.mark.parametrize("retries", [0, 1])
+def test_oauth_with_explicit_issuer(
+ capfd, accept, openid_provider, retries, scope, secret
+):
+ client_id = secrets.token_hex()
+
+ sock, client = accept(
+ oauth_issuer=openid_provider.issuer,
+ oauth_client_id=client_id,
+ oauth_client_secret=secret,
+ oauth_scope=scope,
+ )
+
+ device_code = secrets.token_hex()
+ user_code = f"{secrets.token_hex(2)}-{secrets.token_hex(2)}"
+ verification_url = "https://example.com/device"
+
+ access_token = secrets.token_urlsafe()
+
+ def check_client_authn(headers, params):
+ if not secret:
+ assert params["client_id"] == [client_id]
+ return
+
+ # Require the client to use Basic authn; request-body credentials are
+ # NOT RECOMMENDED (RFC 6749, Sec. 2.3.1).
+ assert "Authorization" in headers
+
+ method, creds = headers["Authorization"].split()
+ assert method == "Basic"
+
+ expected = f"{client_id}:{secret}"
+ assert base64.b64decode(creds) == expected.encode("ascii")
+
+ # Set up our provider callbacks.
+ # NOTE that these callbacks will be called on a background thread. Don't do
+ # any unprotected state mutation here.
+
+ def authorization_endpoint(headers, params):
+ check_client_authn(headers, params)
+
+ if scope:
+ assert params["scope"] == [scope]
+ else:
+ assert "scope" not in params
+
+ resp = {
+ "device_code": device_code,
+ "user_code": user_code,
+ "interval": 0,
+ "verification_uri": verification_url,
+ "expires_in": 5,
+ }
+
+ return 200, resp
+
+ openid_provider.register_endpoint(
+ "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+ )
+
+ attempts = 0
+ retry_lock = threading.Lock()
+
+ def token_endpoint(headers, params):
+ check_client_authn(headers, params)
+
+ assert params["grant_type"] == ["urn:ietf:params:oauth:grant-type:device_code"]
+ assert params["device_code"] == [device_code]
+
+ now = time.monotonic()
+
+ with retry_lock:
+ nonlocal attempts
+
+ # If the test wants to force the client to retry, return an
+ # authorization_pending response and decrement the retry count.
+ if attempts < retries:
+ attempts += 1
+ return 400, {"error": "authorization_pending"}
+
+ # Successfully finish the request by sending the access bearer token.
+ resp = {
+ "access_token": access_token,
+ "token_type": "bearer",
+ }
+
+ return 200, resp
+
+ openid_provider.register_endpoint(
+ "token_endpoint", "POST", "/token", token_endpoint
+ )
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ # Initiate a handshake, which should result in the above endpoints
+ # being called.
+ initial = start_oauth_handshake(conn)
+
+ # Validate and accept the token.
+ auth = get_auth_value(initial)
+ assert auth == f"Bearer {access_token}".encode("ascii")
+
+ pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+ finish_handshake(conn)
+
+ if retries:
+ # Finally, make sure that the client prompted the user with the expected
+ # authorization URL and user code.
+ expected = f"Visit {verification_url} and enter the code: {user_code}"
+ _, stderr = capfd.readouterr()
+ assert expected in stderr
+
+
+def test_oauth_requires_client_id(accept, openid_provider):
+ sock, client = accept(
+ oauth_issuer=openid_provider.issuer,
+ # Do not set a client ID; this should cause a client error after the
+ # server asks for OAUTHBEARER and the client tries to contact the
+ # issuer.
+ )
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ # Initiate a handshake.
+ startup = pq3.recv1(conn, cls=pq3.Startup)
+ assert startup.proto == pq3.protocol(3, 0)
+
+ pq3.send(
+ conn,
+ pq3.types.AuthnRequest,
+ type=pq3.authn.SASL,
+ body=[b"OAUTHBEARER", b""],
+ )
+
+ # The client should disconnect at this point.
+ assert not conn.read()
+
+ expected_error = "no oauth_client_id is set"
+ with pytest.raises(psycopg2.OperationalError, match=expected_error):
+ client.check_completed()
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("error_code", ["authorization_pending", "slow_down"])
+@pytest.mark.parametrize("retries", [1, 2])
+def test_oauth_retry_interval(accept, openid_provider, retries, error_code):
+ sock, client = accept(
+ oauth_issuer=openid_provider.issuer,
+ oauth_client_id="some-id",
+ )
+
+ expected_retry_interval = 1
+ access_token = secrets.token_urlsafe()
+
+ # Set up our provider callbacks.
+ # NOTE that these callbacks will be called on a background thread. Don't do
+ # any unprotected state mutation here.
+
+ def authorization_endpoint(headers, params):
+ resp = {
+ "device_code": "my-device-code",
+ "user_code": "my-user-code",
+ "interval": expected_retry_interval,
+ "verification_uri": "https://example.com",
+ "expires_in": 5,
+ }
+
+ return 200, resp
+
+ openid_provider.register_endpoint(
+ "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+ )
+
+ attempts = 0
+ last_retry = None
+ retry_lock = threading.Lock()
+
+ def token_endpoint(headers, params):
+ now = time.monotonic()
+
+ with retry_lock:
+ nonlocal attempts, last_retry, expected_retry_interval
+
+ # Make sure the retry interval is being respected by the client.
+ if last_retry is not None:
+ interval = now - last_retry
+ assert interval >= expected_retry_interval
+
+ last_retry = now
+
+ # If the test wants to force the client to retry, return the desired
+ # error response and decrement the retry count.
+ if attempts < retries:
+ attempts += 1
+
+ # A slow_down code requires the client to additionally increase
+ # its interval by five seconds.
+ if error_code == "slow_down":
+ expected_retry_interval += 5
+
+ return 400, {"error": error_code}
+
+ # Successfully finish the request by sending the access bearer token.
+ resp = {
+ "access_token": access_token,
+ "token_type": "bearer",
+ }
+
+ return 200, resp
+
+ openid_provider.register_endpoint(
+ "token_endpoint", "POST", "/token", token_endpoint
+ )
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ # Initiate a handshake, which should result in the above endpoints
+ # being called.
+ initial = start_oauth_handshake(conn)
+
+ # Validate and accept the token.
+ auth = get_auth_value(initial)
+ assert auth == f"Bearer {access_token}".encode("ascii")
+
+ pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+ finish_handshake(conn)
+
+
+@pytest.mark.parametrize(
+ "failure_mode, error_pattern",
+ [
+ pytest.param(
+ {
+ "error": "invalid_client",
+ "error_description": "client authentication failed",
+ },
+ r"client authentication failed \(invalid_client\)",
+ id="authentication failure with description",
+ ),
+ pytest.param(
+ {"error": "invalid_request"},
+ r"\(invalid_request\)",
+ id="invalid request without description",
+ ),
+ pytest.param(
+ {},
+ r"failed to obtain device authorization",
+ id="broken error response",
+ ),
+ ],
+)
+def test_oauth_device_authorization_failures(
+ accept, openid_provider, failure_mode, error_pattern
+):
+ client_id = secrets.token_hex()
+
+ sock, client = accept(
+ oauth_issuer=openid_provider.issuer,
+ oauth_client_id=client_id,
+ )
+
+ # Set up our provider callbacks.
+ # NOTE that these callbacks will be called on a background thread. Don't do
+ # any unprotected state mutation here.
+
+ def authorization_endpoint(headers, params):
+ return 400, failure_mode
+
+ openid_provider.register_endpoint(
+ "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+ )
+
+ def token_endpoint(headers, params):
+ assert False, "token endpoint was invoked unexpectedly"
+
+ openid_provider.register_endpoint(
+ "token_endpoint", "POST", "/token", token_endpoint
+ )
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ # Initiate a handshake, which should result in the above endpoints
+ # being called.
+ startup = pq3.recv1(conn, cls=pq3.Startup)
+ assert startup.proto == pq3.protocol(3, 0)
+
+ pq3.send(
+ conn,
+ pq3.types.AuthnRequest,
+ type=pq3.authn.SASL,
+ body=[b"OAUTHBEARER", b""],
+ )
+
+ # The client should not continue the connection due to the hardcoded
+ # provider failure; we disconnect here.
+
+ # Now make sure the client correctly failed.
+ with pytest.raises(psycopg2.OperationalError, match=error_pattern):
+ client.check_completed()
+
+
+@pytest.mark.parametrize(
+ "failure_mode, error_pattern",
+ [
+ pytest.param(
+ {
+ "error": "expired_token",
+ "error_description": "the device code has expired",
+ },
+ r"the device code has expired \(expired_token\)",
+ id="expired token with description",
+ ),
+ pytest.param(
+ {"error": "access_denied"},
+ r"\(access_denied\)",
+ id="access denied without description",
+ ),
+ pytest.param(
+ {},
+ r"OAuth token retrieval failed",
+ id="broken error response",
+ ),
+ ],
+)
+@pytest.mark.parametrize("retries", [0, 1])
+def test_oauth_token_failures(
+ accept, openid_provider, retries, failure_mode, error_pattern
+):
+ client_id = secrets.token_hex()
+
+ sock, client = accept(
+ oauth_issuer=openid_provider.issuer,
+ oauth_client_id=client_id,
+ )
+
+ device_code = secrets.token_hex()
+ user_code = f"{secrets.token_hex(2)}-{secrets.token_hex(2)}"
+
+ # Set up our provider callbacks.
+ # NOTE that these callbacks will be called on a background thread. Don't do
+ # any unprotected state mutation here.
+
+ def authorization_endpoint(headers, params):
+ assert params["client_id"] == [client_id]
+
+ resp = {
+ "device_code": device_code,
+ "user_code": user_code,
+ "interval": 0,
+ "verification_uri": "https://example.com/device",
+ "expires_in": 5,
+ }
+
+ return 200, resp
+
+ openid_provider.register_endpoint(
+ "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+ )
+
+ retry_lock = threading.Lock()
+
+ def token_endpoint(headers, params):
+ with retry_lock:
+ nonlocal retries
+
+ # If the test wants to force the client to retry, return an
+ # authorization_pending response and decrement the retry count.
+ if retries > 0:
+ retries -= 1
+ return 400, {"error": "authorization_pending"}
+
+ return 400, failure_mode
+
+ openid_provider.register_endpoint(
+ "token_endpoint", "POST", "/token", token_endpoint
+ )
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ # Initiate a handshake, which should result in the above endpoints
+ # being called.
+ startup = pq3.recv1(conn, cls=pq3.Startup)
+ assert startup.proto == pq3.protocol(3, 0)
+
+ pq3.send(
+ conn,
+ pq3.types.AuthnRequest,
+ type=pq3.authn.SASL,
+ body=[b"OAUTHBEARER", b""],
+ )
+
+ # The client should not continue the connection due to the hardcoded
+ # provider failure; we disconnect here.
+
+ # Now make sure the client correctly failed.
+ with pytest.raises(psycopg2.OperationalError, match=error_pattern):
+ client.check_completed()
+
+
+@pytest.mark.parametrize("scope", [None, "openid email"])
+@pytest.mark.parametrize(
+ "base_response",
+ [
+ {"status": "invalid_token"},
+ {"extra_object": {"key": "value"}, "status": "invalid_token"},
+ {"extra_object": {"status": 1}, "status": "invalid_token"},
+ ],
+)
+def test_oauth_discovery(accept, openid_provider, base_response, scope):
+ sock, client = accept(oauth_client_id=secrets.token_hex())
+
+ device_code = secrets.token_hex()
+ user_code = f"{secrets.token_hex(2)}-{secrets.token_hex(2)}"
+ verification_url = "https://example.com/device"
+
+ access_token = secrets.token_urlsafe()
+
+ # Set up our provider callbacks.
+ # NOTE that these callbacks will be called on a background thread. Don't do
+ # any unprotected state mutation here.
+
+ def authorization_endpoint(headers, params):
+ if scope:
+ assert params["scope"] == [scope]
+ else:
+ assert "scope" not in params
+
+ resp = {
+ "device_code": device_code,
+ "user_code": user_code,
+ "interval": 0,
+ "verification_uri": verification_url,
+ "expires_in": 5,
+ }
+
+ return 200, resp
+
+ openid_provider.register_endpoint(
+ "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+ )
+
+ def token_endpoint(headers, params):
+ assert params["grant_type"] == ["urn:ietf:params:oauth:grant-type:device_code"]
+ assert params["device_code"] == [device_code]
+
+ # Successfully finish the request by sending the access bearer token.
+ resp = {
+ "access_token": access_token,
+ "token_type": "bearer",
+ }
+
+ return 200, resp
+
+ openid_provider.register_endpoint(
+ "token_endpoint", "POST", "/token", token_endpoint
+ )
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ initial = start_oauth_handshake(conn)
+
+ # For discovery, the client should send an empty auth header. See
+ # RFC 7628, Sec. 4.3.
+ auth = get_auth_value(initial)
+ assert auth == b""
+
+ # We will fail the first SASL exchange. First return a link to the
+ # discovery document, pointing to the test provider server.
+ resp = dict(base_response)
+
+ discovery_uri = f"{openid_provider.issuer}/.well-known/openid-configuration"
+ resp["openid-configuration"] = discovery_uri
+
+ if scope:
+ resp["scope"] = scope
+
+ resp = json.dumps(resp)
+
+ pq3.send(
+ conn,
+ pq3.types.AuthnRequest,
+ type=pq3.authn.SASLContinue,
+ body=resp.encode("ascii"),
+ )
+
+ # Per RFC, the client is required to send a dummy ^A response.
+ pkt = pq3.recv1(conn)
+ assert pkt.type == pq3.types.PasswordMessage
+ assert pkt.payload == b"\x01"
+
+ # Now fail the SASL exchange.
+ pq3.send(
+ conn,
+ pq3.types.ErrorResponse,
+ fields=[
+ b"SFATAL",
+ b"C28000",
+ b"Mdoesn't matter",
+ b"",
+ ],
+ )
+
+ # The client will connect to us a second time, using the parameters we sent
+ # it.
+ sock, _ = accept()
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ initial = start_oauth_handshake(conn)
+
+ # Validate and accept the token.
+ auth = get_auth_value(initial)
+ assert auth == f"Bearer {access_token}".encode("ascii")
+
+ pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+ finish_handshake(conn)
+
+
+@pytest.mark.parametrize(
+ "response,expected_error",
+ [
+ pytest.param(
+ "abcde",
+ 'Token "abcde" is invalid',
+ id="bad JSON: invalid syntax",
+ ),
+ pytest.param(
+ '"abcde"',
+ "top-level element must be an object",
+ id="bad JSON: top-level element is a string",
+ ),
+ pytest.param(
+ "[]",
+ "top-level element must be an object",
+ id="bad JSON: top-level element is an array",
+ ),
+ pytest.param(
+ "{}",
+ "server sent error response without a status",
+ id="bad JSON: no status member",
+ ),
+ pytest.param(
+ '{ "status": null }',
+ 'field "status" must be a string',
+ id="bad JSON: null status member",
+ ),
+ pytest.param(
+ '{ "status": 0 }',
+ 'field "status" must be a string',
+ id="bad JSON: int status member",
+ ),
+ pytest.param(
+ '{ "status": [ "bad" ] }',
+ 'field "status" must be a string',
+ id="bad JSON: array status member",
+ ),
+ pytest.param(
+ '{ "status": { "bad": "bad" } }',
+ 'field "status" must be a string',
+ id="bad JSON: object status member",
+ ),
+ pytest.param(
+ '{ "nested": { "status": "bad" } }',
+ "server sent error response without a status",
+ id="bad JSON: nested status",
+ ),
+ pytest.param(
+ '{ "status": "invalid_token" ',
+ "The input string ended unexpectedly",
+ id="bad JSON: unterminated object",
+ ),
+ pytest.param(
+ '{ "status": "invalid_token" } { }',
+ 'Expected end of input, but found "{"',
+ id="bad JSON: trailing data",
+ ),
+ pytest.param(
+ '{ "status": "invalid_token", "openid-configuration": 1 }',
+ 'field "openid-configuration" must be a string',
+ id="bad JSON: int openid-configuration member",
+ ),
+ pytest.param(
+ '{ "status": "invalid_token", "openid-configuration": 1 }',
+ 'field "openid-configuration" must be a string',
+ id="bad JSON: int openid-configuration member",
+ ),
+ pytest.param(
+ '{ "status": "invalid_token", "scope": 1 }',
+ 'field "scope" must be a string',
+ id="bad JSON: int scope member",
+ ),
+ ],
+)
+def test_oauth_discovery_server_error(accept, response, expected_error):
+ sock, client = accept(oauth_client_id=secrets.token_hex())
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ initial = start_oauth_handshake(conn)
+
+ # Fail the SASL exchange with an invalid JSON response.
+ pq3.send(
+ conn,
+ pq3.types.AuthnRequest,
+ type=pq3.authn.SASLContinue,
+ body=response.encode("utf-8"),
+ )
+
+ # The client should disconnect, so the socket is closed here. (If
+ # the client doesn't disconnect, it will report a different error
+ # below and the test will fail.)
+
+ with pytest.raises(psycopg2.OperationalError, match=expected_error):
+ client.check_completed()
+
+
+@pytest.mark.parametrize(
+ "sasl_err,resp_type,resp_payload,expected_error",
+ [
+ pytest.param(
+ {"status": "invalid_request"},
+ pq3.types.ErrorResponse,
+ dict(
+ fields=[b"SFATAL", b"C28000", b"Mexpected error message", b""],
+ ),
+ "expected error message",
+ id="standard server error: invalid_request",
+ ),
+ pytest.param(
+ {"status": "invalid_token"},
+ pq3.types.ErrorResponse,
+ dict(
+ fields=[b"SFATAL", b"C28000", b"Mexpected error message", b""],
+ ),
+ "expected error message",
+ id="standard server error: invalid_token without discovery URI",
+ ),
+ pytest.param(
+ {"status": "invalid_request"},
+ pq3.types.AuthnRequest,
+ dict(type=pq3.authn.SASLContinue, body=b""),
+ "server sent additional OAuth data",
+ id="broken server: additional challenge after error",
+ ),
+ pytest.param(
+ {"status": "invalid_request"},
+ pq3.types.AuthnRequest,
+ dict(type=pq3.authn.SASLFinal),
+ "server sent additional OAuth data",
+ id="broken server: SASL success after error",
+ ),
+ ],
+)
+def test_oauth_server_error(accept, sasl_err, resp_type, resp_payload, expected_error):
+ sock, client = accept()
+
+ with sock:
+ with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+ start_oauth_handshake(conn)
+
+ # Ignore the client data. Return an error "challenge".
+ resp = json.dumps(sasl_err)
+ resp = resp.encode("utf-8")
+
+ pq3.send(
+ conn, pq3.types.AuthnRequest, type=pq3.authn.SASLContinue, body=resp
+ )
+
+ # Per RFC, the client is required to send a dummy ^A response.
+ pkt = pq3.recv1(conn)
+ assert pkt.type == pq3.types.PasswordMessage
+ assert pkt.payload == b"\x01"
+
+ # Now fail the SASL exchange (in either a valid way, or an invalid
+ # one, depending on the test).
+ pq3.send(conn, resp_type, **resp_payload)
+
+ with pytest.raises(psycopg2.OperationalError, match=expected_error):
+ client.check_completed()
diff --git a/src/test/python/pq3.py b/src/test/python/pq3.py
new file mode 100644
index 0000000000..3a22dad0b6
--- /dev/null
+++ b/src/test/python/pq3.py
@@ -0,0 +1,727 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import contextlib
+import getpass
+import io
+import os
+import ssl
+import sys
+import textwrap
+
+from construct import *
+
+import tls
+
+
+def protocol(major, minor):
+ """
+ Returns the protocol version, in integer format, corresponding to the given
+ major and minor version numbers.
+ """
+ return (major << 16) | minor
+
+
+# Startup
+
+StringList = GreedyRange(NullTerminated(GreedyBytes))
+
+
+class KeyValueAdapter(Adapter):
+ """
+ Turns a key-value store into a null-terminated list of null-terminated
+ strings, as presented on the wire in the startup packet.
+ """
+
+ def _encode(self, obj, context, path):
+ if isinstance(obj, list):
+ return obj
+
+ l = []
+
+ for k, v in obj.items():
+ if isinstance(k, str):
+ k = k.encode("utf-8")
+ l.append(k)
+
+ if isinstance(v, str):
+ v = v.encode("utf-8")
+ l.append(v)
+
+ l.append(b"")
+ return l
+
+ def _decode(self, obj, context, path):
+ # TODO: turn a list back into a dict
+ return obj
+
+
+KeyValues = KeyValueAdapter(StringList)
+
+_startup_payload = Switch(
+ this.proto,
+ {
+ protocol(3, 0): KeyValues,
+ },
+ default=GreedyBytes,
+)
+
+
+def _default_protocol(this):
+ try:
+ if isinstance(this.payload, (list, dict)):
+ return protocol(3, 0)
+ except AttributeError:
+ pass # no payload passed during build
+
+ return 0
+
+
+def _startup_payload_len(this):
+ """
+ The payload field has a fixed size based on the length of the packet. But
+ if the caller hasn't supplied an explicit length at build time, we have to
+ build the payload to figure out how long it is, which requires us to know
+ the length first... This function exists solely to break the cycle.
+ """
+ assert this._building, "_startup_payload_len() cannot be called during parsing"
+
+ try:
+ payload = this.payload
+ except AttributeError:
+ return 0 # no payload
+
+ if isinstance(payload, bytes):
+ # already serialized; just use the given length
+ return len(payload)
+
+ try:
+ proto = this.proto
+ except AttributeError:
+ proto = _default_protocol(this)
+
+ data = _startup_payload.build(payload, proto=proto)
+ return len(data)
+
+
+Startup = Struct(
+ "len" / Default(Int32sb, lambda this: _startup_payload_len(this) + 8),
+ "proto" / Default(Hex(Int32sb), _default_protocol),
+ "payload" / FixedSized(this.len - 8, Default(_startup_payload, b"")),
+)
+
+# Pq3
+
+# Adapted from construct.core.EnumIntegerString
+class EnumNamedByte:
+ def __init__(self, val, name):
+ self._val = val
+ self._name = name
+
+ def __int__(self):
+ return ord(self._val)
+
+ def __str__(self):
+ return "(enum) %s %r" % (self._name, self._val)
+
+ def __repr__(self):
+ return "EnumNamedByte(%r)" % self._val
+
+ def __eq__(self, other):
+ if isinstance(other, EnumNamedByte):
+ other = other._val
+ if not isinstance(other, bytes):
+ return NotImplemented
+
+ return self._val == other
+
+ def __hash__(self):
+ return hash(self._val)
+
+
+# Adapted from construct.core.Enum
+class ByteEnum(Adapter):
+ def __init__(self, **mapping):
+ super(ByteEnum, self).__init__(Byte)
+ self.namemapping = {k: EnumNamedByte(v, k) for k, v in mapping.items()}
+ self.decmapping = {v: EnumNamedByte(v, k) for k, v in mapping.items()}
+
+ def __getattr__(self, name):
+ if name in self.namemapping:
+ return self.decmapping[self.namemapping[name]]
+ raise AttributeError
+
+ def _decode(self, obj, context, path):
+ b = bytes([obj])
+ try:
+ return self.decmapping[b]
+ except KeyError:
+ return EnumNamedByte(b, "(unknown)")
+
+ def _encode(self, obj, context, path):
+ if isinstance(obj, int):
+ return obj
+ elif isinstance(obj, bytes):
+ return ord(obj)
+ return int(obj)
+
+
+types = ByteEnum(
+ ErrorResponse=b"E",
+ ReadyForQuery=b"Z",
+ Query=b"Q",
+ EmptyQueryResponse=b"I",
+ AuthnRequest=b"R",
+ PasswordMessage=b"p",
+ BackendKeyData=b"K",
+ CommandComplete=b"C",
+ ParameterStatus=b"S",
+ DataRow=b"D",
+ Terminate=b"X",
+)
+
+
+authn = Enum(
+ Int32ub,
+ OK=0,
+ SASL=10,
+ SASLContinue=11,
+ SASLFinal=12,
+)
+
+
+_authn_body = Switch(
+ this.type,
+ {
+ authn.OK: Terminated,
+ authn.SASL: StringList,
+ },
+ default=GreedyBytes,
+)
+
+
+def _data_len(this):
+ assert this._building, "_data_len() cannot be called during parsing"
+
+ if not hasattr(this, "data") or this.data is None:
+ return -1
+
+ return len(this.data)
+
+
+# The protocol reuses the PasswordMessage for several authentication response
+# types, and there's no good way to figure out which is which without keeping
+# state for the entire stream. So this is a separate Construct that can be
+# explicitly parsed/built by code that knows it's needed.
+SASLInitialResponse = Struct(
+ "name" / NullTerminated(GreedyBytes),
+ "len" / Default(Int32sb, lambda this: _data_len(this)),
+ "data"
+ / IfThenElse(
+ # Allow tests to explicitly pass an incorrect length during testing, by
+ # not enforcing a FixedSized during build. (The len calculation above
+ # defaults to the correct size.)
+ this._building,
+ Optional(GreedyBytes),
+ If(this.len != -1, Default(FixedSized(this.len, GreedyBytes), b"")),
+ ),
+ Terminated, # make sure the entire response is consumed
+)
+
+
+_column = FocusedSeq(
+ "data",
+ "len" / Default(Int32sb, lambda this: _data_len(this)),
+ "data" / If(this.len != -1, FixedSized(this.len, GreedyBytes)),
+)
+
+
+_payload_map = {
+ types.ErrorResponse: Struct("fields" / StringList),
+ types.ReadyForQuery: Struct("status" / Bytes(1)),
+ types.Query: Struct("query" / NullTerminated(GreedyBytes)),
+ types.EmptyQueryResponse: Terminated,
+ types.AuthnRequest: Struct("type" / authn, "body" / Default(_authn_body, b"")),
+ types.BackendKeyData: Struct("pid" / Int32ub, "key" / Hex(Int32ub)),
+ types.CommandComplete: Struct("tag" / NullTerminated(GreedyBytes)),
+ types.ParameterStatus: Struct(
+ "name" / NullTerminated(GreedyBytes), "value" / NullTerminated(GreedyBytes)
+ ),
+ types.DataRow: Struct("columns" / Default(PrefixedArray(Int16sb, _column), b"")),
+ types.Terminate: Terminated,
+}
+
+
+_payload = FocusedSeq(
+ "_payload",
+ "_payload"
+ / Switch(
+ this._.type,
+ _payload_map,
+ default=GreedyBytes,
+ ),
+ Terminated, # make sure every payload consumes the entire packet
+)
+
+
+def _payload_len(this):
+ """
+ See _startup_payload_len() for an explanation.
+ """
+ assert this._building, "_payload_len() cannot be called during parsing"
+
+ try:
+ payload = this.payload
+ except AttributeError:
+ return 0 # no payload
+
+ if isinstance(payload, bytes):
+ # already serialized; just use the given length
+ return len(payload)
+
+ data = _payload.build(payload, type=this.type)
+ return len(data)
+
+
+Pq3 = Struct(
+ "type" / types,
+ "len" / Default(Int32ub, lambda this: _payload_len(this) + 4),
+ "payload" / FixedSized(this.len - 4, Default(_payload, b"")),
+)
+
+
+# Environment
+
+
+def pghost():
+ return os.environ.get("PGHOST", default="localhost")
+
+
+def pgport():
+ return int(os.environ.get("PGPORT", default=5432))
+
+
+def pguser():
+ try:
+ return os.environ["PGUSER"]
+ except KeyError:
+ return getpass.getuser()
+
+
+def pgdatabase():
+ return os.environ.get("PGDATABASE", default="postgres")
+
+
+# Connections
+
+
+def _hexdump_translation_map():
+ """
+ For hexdumps. Translates any unprintable or non-ASCII bytes into '.'.
+ """
+ input = bytearray()
+
+ for i in range(128):
+ c = chr(i)
+
+ if not c.isprintable():
+ input += bytes([i])
+
+ input += bytes(range(128, 256))
+
+ return bytes.maketrans(input, b"." * len(input))
+
+
+class _DebugStream(object):
+ """
+ Wraps a file-like object and adds hexdumps of the read and write data. Call
+ end_packet() on a _DebugStream to write the accumulated hexdumps to the
+ output stream, along with the packet that was sent.
+ """
+
+ _translation_map = _hexdump_translation_map()
+
+ def __init__(self, stream, out=sys.stdout):
+ """
+ Creates a new _DebugStream wrapping the given stream (which must have
+ been created by wrap()). All attributes not provided by the _DebugStream
+ are delegated to the wrapped stream. out is the text stream to which
+ hexdumps are written.
+ """
+ self.raw = stream
+ self._out = out
+ self._rbuf = io.BytesIO()
+ self._wbuf = io.BytesIO()
+
+ def __getattr__(self, name):
+ return getattr(self.raw, name)
+
+ def __setattr__(self, name, value):
+ if name in ("raw", "_out", "_rbuf", "_wbuf"):
+ return object.__setattr__(self, name, value)
+
+ setattr(self.raw, name, value)
+
+ def read(self, *args, **kwargs):
+ buf = self.raw.read(*args, **kwargs)
+
+ self._rbuf.write(buf)
+ return buf
+
+ def write(self, b):
+ self._wbuf.write(b)
+ return self.raw.write(b)
+
+ def recv(self, *args):
+ buf = self.raw.recv(*args)
+
+ self._rbuf.write(buf)
+ return buf
+
+ def _flush(self, buf, prefix):
+ width = 16
+ hexwidth = width * 3 - 1
+
+ count = 0
+ buf.seek(0)
+
+ while True:
+ line = buf.read(16)
+
+ if not line:
+ if count:
+ self._out.write("\n") # separate the output block with a newline
+ return
+
+ self._out.write("%s %04X:\t" % (prefix, count))
+ self._out.write("%*s\t" % (-hexwidth, line.hex(" ")))
+ self._out.write(line.translate(self._translation_map).decode("ascii"))
+ self._out.write("\n")
+
+ count += 16
+
+ def print_debug(self, obj, *, prefix=""):
+ contents = ""
+ if obj is not None:
+ contents = str(obj)
+
+ for line in contents.splitlines():
+ self._out.write("%s%s\n" % (prefix, line))
+
+ self._out.write("\n")
+
+ def flush_debug(self, *, prefix=""):
+ self._flush(self._rbuf, prefix + "<")
+ self._rbuf = io.BytesIO()
+
+ self._flush(self._wbuf, prefix + ">")
+ self._wbuf = io.BytesIO()
+
+ def end_packet(self, pkt, *, read=False, prefix="", indent=" "):
+ """
+ Marks the end of a logical "packet" of data. A string representation of
+ pkt will be printed, and the debug buffers will be flushed with an
+ indent. All lines can be optionally prefixed.
+
+ If read is True, the packet representation is written after the debug
+ buffers; otherwise the default of False (meaning write) causes the
+ packet representation to be dumped first. This is meant to capture the
+ logical flow of layer translation.
+ """
+ write = not read
+
+ if write:
+ self.print_debug(pkt, prefix=prefix + "> ")
+
+ self.flush_debug(prefix=prefix + indent)
+
+ if read:
+ self.print_debug(pkt, prefix=prefix + "< ")
+
+
+@contextlib.contextmanager
+def wrap(socket, *, debug_stream=None):
+ """
+ Transforms a raw socket into a connection that can be used for Construct
+ building and parsing. The return value is a context manager and can be used
+ in a with statement.
+ """
+ # It is critical that buffering be disabled here, so that we can still
+ # manipulate the raw socket without desyncing the stream.
+ with socket.makefile("rwb", buffering=0) as sfile:
+ # Expose the original socket's recv() on the SocketIO object we return.
+ def recv(self, *args):
+ return socket.recv(*args)
+
+ sfile.recv = recv.__get__(sfile)
+
+ conn = sfile
+ if debug_stream:
+ conn = _DebugStream(conn, debug_stream)
+
+ try:
+ yield conn
+ finally:
+ if debug_stream:
+ conn.flush_debug(prefix="? ")
+
+
+def _send(stream, cls, obj):
+ debugging = hasattr(stream, "flush_debug")
+ out = io.BytesIO()
+
+ # Ideally we would build directly to the passed stream, but because we need
+ # to reparse the generated output for the debugging case, build to an
+ # intermediate BytesIO and send it instead.
+ cls.build_stream(obj, out)
+ buf = out.getvalue()
+
+ stream.write(buf)
+ if debugging:
+ pkt = cls.parse(buf)
+ stream.end_packet(pkt)
+
+ stream.flush()
+
+
+def send(stream, packet_type, payload_data=None, **payloadkw):
+ """
+ Sends a packet on the given pq3 connection. type is the pq3.types member
+ that should be assigned to the packet. If payload_data is given, it will be
+ used as the packet payload; otherwise the key/value pairs in payloadkw will
+ be the payload contents.
+ """
+ data = payloadkw
+
+ if payload_data is not None:
+ if payloadkw:
+ raise ValueError(
+ "payload_data and payload keywords may not be used simultaneously"
+ )
+
+ data = payload_data
+
+ _send(stream, Pq3, dict(type=packet_type, payload=data))
+
+
+def send_startup(stream, proto=None, **kwargs):
+ """
+ Sends a startup packet on the given pq3 connection. In most cases you should
+ use the handshake functions instead, which will do this for you.
+
+ By default, a protocol version 3 packet will be sent. This can be overridden
+ with the proto parameter.
+ """
+ pkt = {}
+
+ if proto is not None:
+ pkt["proto"] = proto
+ if kwargs:
+ pkt["payload"] = kwargs
+
+ _send(stream, Startup, pkt)
+
+
+def recv1(stream, *, cls=Pq3):
+ """
+ Receives a single pq3 packet from the given stream and returns it.
+ """
+ resp = cls.parse_stream(stream)
+
+ debugging = hasattr(stream, "flush_debug")
+ if debugging:
+ stream.end_packet(resp, read=True)
+
+ return resp
+
+
+def handshake(stream, **kwargs):
+ """
+ Performs a libpq v3 startup handshake. kwargs should contain the key/value
+ parameters to send to the server in the startup packet.
+ """
+ # Send our startup parameters.
+ send_startup(stream, **kwargs)
+
+ # Receive and dump packets until the server indicates it's ready for our
+ # first query.
+ while True:
+ resp = recv1(stream)
+ if resp is None:
+ raise RuntimeError("server closed connection during handshake")
+
+ if resp.type == types.ReadyForQuery:
+ return
+ elif resp.type == types.ErrorResponse:
+ raise RuntimeError(
+ f"received error response from peer: {resp.payload.fields!r}"
+ )
+
+
+# TLS
+
+
+class _TLSStream(object):
+ """
+ A file-like object that performs TLS encryption/decryption on a wrapped
+ stream. Differs from ssl.SSLSocket in that we have full visibility and
+ control over the TLS layer.
+ """
+
+ def __init__(self, stream, context):
+ self._stream = stream
+ self._debugging = hasattr(stream, "flush_debug")
+
+ self._in = ssl.MemoryBIO()
+ self._out = ssl.MemoryBIO()
+ self._ssl = context.wrap_bio(self._in, self._out)
+
+ def handshake(self):
+ try:
+ self._pump(lambda: self._ssl.do_handshake())
+ finally:
+ self._flush_debug(prefix="? ")
+
+ def read(self, *args):
+ return self._pump(lambda: self._ssl.read(*args))
+
+ def write(self, *args):
+ return self._pump(lambda: self._ssl.write(*args))
+
+ def _decode(self, buf):
+ """
+ Attempts to decode a buffer of TLS data into a packet representation
+ that can be printed.
+
+ TODO: handle buffers (and record fragments) that don't align with packet
+ boundaries.
+ """
+ end = len(buf)
+ bio = io.BytesIO(buf)
+
+ ret = io.StringIO()
+
+ while bio.tell() < end:
+ record = tls.Plaintext.parse_stream(bio)
+
+ if ret.tell() > 0:
+ ret.write("\n")
+ ret.write("[Record] ")
+ ret.write(str(record))
+ ret.write("\n")
+
+ if record.type == tls.ContentType.handshake:
+ record_cls = tls.Handshake
+ else:
+ continue
+
+ innerlen = len(record.fragment)
+ inner = io.BytesIO(record.fragment)
+
+ while inner.tell() < innerlen:
+ msg = record_cls.parse_stream(inner)
+
+ indented = "[Message] " + str(msg)
+ indented = textwrap.indent(indented, " ")
+
+ ret.write("\n")
+ ret.write(indented)
+ ret.write("\n")
+
+ return ret.getvalue()
+
+ def flush(self):
+ if not self._out.pending:
+ self._stream.flush()
+ return
+
+ buf = self._out.read()
+ self._stream.write(buf)
+
+ if self._debugging:
+ pkt = self._decode(buf)
+ self._stream.end_packet(pkt, prefix=" ")
+
+ self._stream.flush()
+
+ def _pump(self, operation):
+ while True:
+ try:
+ return operation()
+ except (ssl.SSLWantReadError, ssl.SSLWantWriteError) as e:
+ want = e
+ self._read_write(want)
+
+ def _recv(self, maxsize):
+ buf = self._stream.recv(4096)
+ if not buf:
+ self._in.write_eof()
+ return
+
+ self._in.write(buf)
+
+ if not self._debugging:
+ return
+
+ pkt = self._decode(buf)
+ self._stream.end_packet(pkt, read=True, prefix=" ")
+
+ def _read_write(self, want):
+ # XXX This needs work. So many corner cases yet to handle. For one,
+ # doing blocking writes in flush may lead to distributed deadlock if the
+ # peer is already blocking on its writes.
+
+ if isinstance(want, ssl.SSLWantWriteError):
+ assert self._out.pending, "SSL backend wants write without data"
+
+ self.flush()
+
+ if isinstance(want, ssl.SSLWantReadError):
+ self._recv(4096)
+
+ def _flush_debug(self, prefix):
+ if not self._debugging:
+ return
+
+ self._stream.flush_debug(prefix=prefix)
+
+
+@contextlib.contextmanager
+def tls_handshake(stream, context):
+ """
+ Performs a TLS handshake over the given stream (which must have been created
+ via a call to wrap()), and returns a new stream which transparently tunnels
+ data over the TLS connection.
+
+ If the passed stream has debugging enabled, the returned stream will also
+ have debugging, using the same output IO.
+ """
+ debugging = hasattr(stream, "flush_debug")
+
+ # Send our startup parameters.
+ send_startup(stream, proto=protocol(1234, 5679))
+
+ # Look at the SSL response.
+ resp = stream.read(1)
+ if debugging:
+ stream.flush_debug(prefix=" ")
+
+ if resp == b"N":
+ raise RuntimeError("server does not support SSLRequest")
+ if resp != b"S":
+ raise RuntimeError(f"unexpected response of type {resp!r} during TLS startup")
+
+ tls = _TLSStream(stream, context)
+ tls.handshake()
+
+ if debugging:
+ tls = _DebugStream(tls, stream._out)
+
+ try:
+ yield tls
+ # TODO: teardown/unwrap the connection?
+ finally:
+ if debugging:
+ tls.flush_debug(prefix="? ")
diff --git a/src/test/python/pytest.ini b/src/test/python/pytest.ini
new file mode 100644
index 0000000000..ab7a6e7fb9
--- /dev/null
+++ b/src/test/python/pytest.ini
@@ -0,0 +1,4 @@
+[pytest]
+
+markers =
+ slow: mark test as slow
diff --git a/src/test/python/requirements.txt b/src/test/python/requirements.txt
new file mode 100644
index 0000000000..32f105ea84
--- /dev/null
+++ b/src/test/python/requirements.txt
@@ -0,0 +1,7 @@
+black
+cryptography~=3.4.6
+construct~=2.10.61
+isort~=5.6
+psycopg2~=2.8.6
+pytest~=6.1
+pytest-asyncio~=0.14.0
diff --git a/src/test/python/server/__init__.py b/src/test/python/server/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/test/python/server/conftest.py b/src/test/python/server/conftest.py
new file mode 100644
index 0000000000..ba7342a453
--- /dev/null
+++ b/src/test/python/server/conftest.py
@@ -0,0 +1,45 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import contextlib
+import socket
+import sys
+
+import pytest
+
+import pq3
+
+
+@pytest.fixture
+def connect():
+ """
+ A factory fixture that, when called, returns a socket connected to a
+ Postgres server, wrapped in a pq3 connection. The calling test will be
+ skipped automatically if a server is not running at PGHOST:PGPORT, so it's
+ best to connect as soon as possible after the test case begins, to avoid
+ doing unnecessary work.
+ """
+ # Set up an ExitStack to handle safe cleanup of all of the moving pieces.
+ with contextlib.ExitStack() as stack:
+
+ def conn_factory():
+ addr = (pq3.pghost(), pq3.pgport())
+
+ try:
+ sock = socket.create_connection(addr, timeout=2)
+ except ConnectionError as e:
+ pytest.skip(f"unable to connect to {addr}: {e}")
+
+ # Have ExitStack close our socket.
+ stack.enter_context(sock)
+
+ # Wrap the connection in a pq3 layer and have ExitStack clean it up
+ # too.
+ wrap_ctx = pq3.wrap(sock, debug_stream=sys.stdout)
+ conn = stack.enter_context(wrap_ctx)
+
+ return conn
+
+ yield conn_factory
diff --git a/src/test/python/server/test_oauth.py b/src/test/python/server/test_oauth.py
new file mode 100644
index 0000000000..cb5ca7fa23
--- /dev/null
+++ b/src/test/python/server/test_oauth.py
@@ -0,0 +1,1012 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import base64
+import contextlib
+import json
+import os
+import pathlib
+import secrets
+import shlex
+import shutil
+import socket
+import struct
+from multiprocessing import shared_memory
+
+import psycopg2
+import pytest
+from psycopg2 import sql
+
+import pq3
+
+MAX_SASL_MESSAGE_LENGTH = 65535
+
+INVALID_AUTHORIZATION_ERRCODE = b"28000"
+PROTOCOL_VIOLATION_ERRCODE = b"08P01"
+FEATURE_NOT_SUPPORTED_ERRCODE = b"0A000"
+
+SHARED_MEM_NAME = "oauth-pytest"
+MAX_TOKEN_SIZE = 4096
+MAX_UINT16 = 2 ** 16 - 1
+
+
+def skip_if_no_postgres():
+ """
+ Used by the oauth_ctx fixture to skip this test module if no Postgres server
+ is running.
+
+ This logic is nearly duplicated with the conn fixture. Ideally oauth_ctx
+ would depend on that, but a module-scope fixture can't depend on a
+ test-scope fixture, and we haven't reached the rule of three yet.
+ """
+ addr = (pq3.pghost(), pq3.pgport())
+
+ try:
+ with socket.create_connection(addr, timeout=2):
+ pass
+ except ConnectionError as e:
+ pytest.skip(f"unable to connect to {addr}: {e}")
+
+
+@contextlib.contextmanager
+def prepend_file(path, lines):
+ """
+ A context manager that prepends a file on disk with the desired lines of
+ text. When the context manager is exited, the file will be restored to its
+ original contents.
+ """
+ # First make a backup of the original file.
+ bak = path + ".bak"
+ shutil.copy2(path, bak)
+
+ try:
+ # Write the new lines, followed by the original file content.
+ with open(path, "w") as new, open(bak, "r") as orig:
+ new.writelines(lines)
+ shutil.copyfileobj(orig, new)
+
+ # Return control to the calling code.
+ yield
+
+ finally:
+ # Put the backup back into place.
+ os.replace(bak, path)
+
+
+@pytest.fixture(scope="module")
+def oauth_ctx():
+ """
+ Creates a database and user that use the oauth auth method. The context
+ object contains the dbname and user attributes as strings to be used during
+ connection, as well as the issuer and scope that have been set in the HBA
+ configuration.
+
+ This fixture assumes that the standard PG* environment variables point to a
+ server running on a local machine, and that the PGUSER has rights to create
+ databases and roles.
+ """
+ skip_if_no_postgres() # don't bother running these tests without a server
+
+ id = secrets.token_hex(4)
+
+ class Context:
+ dbname = "oauth_test_" + id
+
+ user = "oauth_user_" + id
+ map_user = "oauth_map_user_" + id
+ authz_user = "oauth_authz_user_" + id
+
+ issuer = "https://example.com/" + id
+ scope = "openid " + id
+
+ ctx = Context()
+ hba_lines = (
+ f'host {ctx.dbname} {ctx.map_user} samehost oauth issuer="{ctx.issuer}" scope="{ctx.scope}" map=oauth\n',
+ f'host {ctx.dbname} {ctx.authz_user} samehost oauth issuer="{ctx.issuer}" scope="{ctx.scope}" trust_validator_authz=1\n',
+ f'host {ctx.dbname} all samehost oauth issuer="{ctx.issuer}" scope="{ctx.scope}"\n',
+ )
+ ident_lines = (r"oauth /^(.*)@example\.com$ \1",)
+
+ conn = psycopg2.connect("")
+ conn.autocommit = True
+
+ with contextlib.closing(conn):
+ c = conn.cursor()
+
+ # Create our roles and database.
+ user = sql.Identifier(ctx.user)
+ map_user = sql.Identifier(ctx.map_user)
+ authz_user = sql.Identifier(ctx.authz_user)
+ dbname = sql.Identifier(ctx.dbname)
+
+ c.execute(sql.SQL("CREATE ROLE {} LOGIN;").format(user))
+ c.execute(sql.SQL("CREATE ROLE {} LOGIN;").format(map_user))
+ c.execute(sql.SQL("CREATE ROLE {} LOGIN;").format(authz_user))
+ c.execute(sql.SQL("CREATE DATABASE {};").format(dbname))
+
+ # Make this test script the server's oauth_validator.
+ path = pathlib.Path(__file__).parent / "validate_bearer.py"
+ path = str(path.absolute())
+
+ cmd = f"{shlex.quote(path)} {SHARED_MEM_NAME} <&%f"
+ c.execute("ALTER SYSTEM SET oauth_validator_command TO %s;", (cmd,))
+
+ # Replace pg_hba and pg_ident.
+ c.execute("SHOW hba_file;")
+ hba = c.fetchone()[0]
+
+ c.execute("SHOW ident_file;")
+ ident = c.fetchone()[0]
+
+ with prepend_file(hba, hba_lines), prepend_file(ident, ident_lines):
+ c.execute("SELECT pg_reload_conf();")
+
+ # Use the new database and user.
+ yield ctx
+
+ # Put things back the way they were.
+ c.execute("SELECT pg_reload_conf();")
+
+ c.execute("ALTER SYSTEM RESET oauth_validator_command;")
+ c.execute(sql.SQL("DROP DATABASE {};").format(dbname))
+ c.execute(sql.SQL("DROP ROLE {};").format(authz_user))
+ c.execute(sql.SQL("DROP ROLE {};").format(map_user))
+ c.execute(sql.SQL("DROP ROLE {};").format(user))
+
+
+@pytest.fixture()
+def conn(oauth_ctx, connect):
+ """
+ A convenience wrapper for connect(). The main purpose of this fixture is to
+ make sure oauth_ctx runs its setup code before the connection is made.
+ """
+ return connect()
+
+
+@pytest.fixture(scope="module", autouse=True)
+def authn_id_extension(oauth_ctx):
+ """
+ Performs a `CREATE EXTENSION authn_id` in the test database. This fixture is
+ autoused, so tests don't need to rely on it.
+ """
+ conn = psycopg2.connect(database=oauth_ctx.dbname)
+ conn.autocommit = True
+
+ with contextlib.closing(conn):
+ c = conn.cursor()
+ c.execute("CREATE EXTENSION authn_id;")
+
+
+@pytest.fixture(scope="session")
+def shared_mem():
+ """
+ Yields a shared memory segment that can be used for communication between
+ the bearer_token fixture and ./validate_bearer.py.
+ """
+ size = MAX_TOKEN_SIZE + 2 # two byte length prefix
+ mem = shared_memory.SharedMemory(SHARED_MEM_NAME, create=True, size=size)
+
+ try:
+ with contextlib.closing(mem):
+ yield mem
+ finally:
+ mem.unlink()
+
+
+@pytest.fixture()
+def bearer_token(shared_mem):
+ """
+ Returns a factory function that, when called, will store a Bearer token in
+ shared_mem. If token is None (the default), a new token will be generated
+ using secrets.token_urlsafe() and returned; otherwise the passed token will
+ be used as-is.
+
+ When token is None, the generated token size in bytes may be specified as an
+ argument; if unset, a small 16-byte token will be generated. The token size
+ may not exceed MAX_TOKEN_SIZE in any case.
+
+ The return value is the token, converted to a bytes object.
+
+ As a special case for testing failure modes, accept_any may be set to True.
+ This signals to the validator command that any bearer token should be
+ accepted. The returned token in this case may be used or discarded as needed
+ by the test.
+ """
+
+ def set_token(token=None, *, size=16, accept_any=False):
+ if token is not None:
+ size = len(token)
+
+ if size > MAX_TOKEN_SIZE:
+ raise ValueError(f"token size {size} exceeds maximum size {MAX_TOKEN_SIZE}")
+
+ if token is None:
+ if size % 4:
+ raise ValueError(f"requested token size {size} is not a multiple of 4")
+
+ token = secrets.token_urlsafe(size // 4 * 3)
+ assert len(token) == size
+
+ try:
+ token = token.encode("ascii")
+ except AttributeError:
+ pass # already encoded
+
+ if accept_any:
+ # Two-byte magic value.
+ shared_mem.buf[:2] = struct.pack("H", MAX_UINT16)
+ else:
+ # Two-byte length prefix, then the token data.
+ shared_mem.buf[:2] = struct.pack("H", len(token))
+ shared_mem.buf[2 : size + 2] = token
+
+ return token
+
+ return set_token
+
+
+def begin_oauth_handshake(conn, oauth_ctx, *, user=None):
+ if user is None:
+ user = oauth_ctx.authz_user
+
+ pq3.send_startup(conn, user=user, database=oauth_ctx.dbname)
+
+ resp = pq3.recv1(conn)
+ assert resp.type == pq3.types.AuthnRequest
+
+ # The server should advertise exactly one mechanism.
+ assert resp.payload.type == pq3.authn.SASL
+ assert resp.payload.body == [b"OAUTHBEARER", b""]
+
+
+def send_initial_response(conn, *, auth=None, bearer=None):
+ """
+ Sends the OAUTHBEARER initial response on the connection, using the given
+ bearer token. Alternatively to a bearer token, the initial response's auth
+ field may be explicitly specified to test corner cases.
+ """
+ if bearer is not None and auth is not None:
+ raise ValueError("exactly one of the auth and bearer kwargs must be set")
+
+ if bearer is not None:
+ auth = b"Bearer " + bearer
+
+ if auth is None:
+ raise ValueError("exactly one of the auth and bearer kwargs must be set")
+
+ initial = pq3.SASLInitialResponse.build(
+ dict(
+ name=b"OAUTHBEARER",
+ data=b"n,,\x01auth=" + auth + b"\x01\x01",
+ )
+ )
+ pq3.send(conn, pq3.types.PasswordMessage, initial)
+
+
+def expect_handshake_success(conn):
+ """
+ Validates that the server responds with an AuthnOK message, and then drains
+ the connection until a ReadyForQuery message is received.
+ """
+ resp = pq3.recv1(conn)
+
+ assert resp.type == pq3.types.AuthnRequest
+ assert resp.payload.type == pq3.authn.OK
+ assert not resp.payload.body
+
+ receive_until(conn, pq3.types.ReadyForQuery)
+
+
+def expect_handshake_failure(conn, oauth_ctx):
+ """
+ Performs the OAUTHBEARER SASL failure "handshake" and validates the server's
+ side of the conversation, including the final ErrorResponse.
+ """
+
+ # We expect a discovery "challenge" back from the server before the authn
+ # failure message.
+ resp = pq3.recv1(conn)
+ assert resp.type == pq3.types.AuthnRequest
+
+ req = resp.payload
+ assert req.type == pq3.authn.SASLContinue
+
+ body = json.loads(req.body)
+ assert body["status"] == "invalid_token"
+ assert body["scope"] == oauth_ctx.scope
+
+ expected_config = oauth_ctx.issuer + "/.well-known/openid-configuration"
+ assert body["openid-configuration"] == expected_config
+
+ # Send the dummy response to complete the failed handshake.
+ pq3.send(conn, pq3.types.PasswordMessage, b"\x01")
+ resp = pq3.recv1(conn)
+
+ err = ExpectedError(INVALID_AUTHORIZATION_ERRCODE, "bearer authentication failed")
+ err.match(resp)
+
+
+def receive_until(conn, type):
+ """
+ receive_until pulls packets off the pq3 connection until a packet with the
+ desired type is found, or an error response is received.
+ """
+ while True:
+ pkt = pq3.recv1(conn)
+
+ if pkt.type == type:
+ return pkt
+ elif pkt.type == pq3.types.ErrorResponse:
+ raise RuntimeError(
+ f"received error response from peer: {pkt.payload.fields!r}"
+ )
+
+
+@pytest.mark.parametrize("token_len", [16, 1024, 4096])
+@pytest.mark.parametrize(
+ "auth_prefix",
+ [
+ b"Bearer ",
+ b"bearer ",
+ b"Bearer ",
+ ],
+)
+def test_oauth(conn, oauth_ctx, bearer_token, auth_prefix, token_len):
+ begin_oauth_handshake(conn, oauth_ctx)
+
+ # Generate our bearer token with the desired length.
+ token = bearer_token(size=token_len)
+ auth = auth_prefix + token
+
+ send_initial_response(conn, auth=auth)
+ expect_handshake_success(conn)
+
+ # Make sure that the server has not set an authenticated ID.
+ pq3.send(conn, pq3.types.Query, query=b"SELECT authn_id();")
+ resp = receive_until(conn, pq3.types.DataRow)
+
+ row = resp.payload
+ assert row.columns == [None]
+
+
+@pytest.mark.parametrize(
+ "token_value",
+ [
+ "abcdzA==",
+ "123456M=",
+ "x-._~+/x",
+ ],
+)
+def test_oauth_bearer_corner_cases(conn, oauth_ctx, bearer_token, token_value):
+ begin_oauth_handshake(conn, oauth_ctx)
+
+ send_initial_response(conn, bearer=bearer_token(token_value))
+
+ expect_handshake_success(conn)
+
+
+@pytest.mark.parametrize(
+ "user,authn_id,should_succeed",
+ [
+ pytest.param(
+ lambda ctx: ctx.user,
+ lambda ctx: ctx.user,
+ True,
+ id="validator authn: succeeds when authn_id == username",
+ ),
+ pytest.param(
+ lambda ctx: ctx.user,
+ lambda ctx: None,
+ False,
+ id="validator authn: fails when authn_id is not set",
+ ),
+ pytest.param(
+ lambda ctx: ctx.user,
+ lambda ctx: ctx.authz_user,
+ False,
+ id="validator authn: fails when authn_id != username",
+ ),
+ pytest.param(
+ lambda ctx: ctx.map_user,
+ lambda ctx: ctx.map_user + "@example.com",
+ True,
+ id="validator with map: succeeds when authn_id matches map",
+ ),
+ pytest.param(
+ lambda ctx: ctx.map_user,
+ lambda ctx: None,
+ False,
+ id="validator with map: fails when authn_id is not set",
+ ),
+ pytest.param(
+ lambda ctx: ctx.map_user,
+ lambda ctx: ctx.map_user + "@example.net",
+ False,
+ id="validator with map: fails when authn_id doesn't match map",
+ ),
+ pytest.param(
+ lambda ctx: ctx.authz_user,
+ lambda ctx: None,
+ True,
+ id="validator authz: succeeds with no authn_id",
+ ),
+ pytest.param(
+ lambda ctx: ctx.authz_user,
+ lambda ctx: "",
+ True,
+ id="validator authz: succeeds with empty authn_id",
+ ),
+ pytest.param(
+ lambda ctx: ctx.authz_user,
+ lambda ctx: "postgres",
+ True,
+ id="validator authz: succeeds with basic username",
+ ),
+ pytest.param(
+ lambda ctx: ctx.authz_user,
+ lambda ctx: "me@example.com",
+ True,
+ id="validator authz: succeeds with email address",
+ ),
+ ],
+)
+def test_oauth_authn_id(conn, oauth_ctx, bearer_token, user, authn_id, should_succeed):
+ token = None
+
+ authn_id = authn_id(oauth_ctx)
+ if authn_id is not None:
+ authn_id = authn_id.encode("ascii")
+
+ # As a hack to get the validator to reflect arbitrary output from this
+ # test, encode the desired output as a base64 token. The validator will
+ # key on the leading "output=" to differentiate this from the random
+ # tokens generated by secrets.token_urlsafe().
+ output = b"output=" + authn_id + b"\n"
+ token = base64.urlsafe_b64encode(output)
+
+ token = bearer_token(token)
+ username = user(oauth_ctx)
+
+ begin_oauth_handshake(conn, oauth_ctx, user=username)
+ send_initial_response(conn, bearer=token)
+
+ if not should_succeed:
+ expect_handshake_failure(conn, oauth_ctx)
+ return
+
+ expect_handshake_success(conn)
+
+ # Check the reported authn_id.
+ pq3.send(conn, pq3.types.Query, query=b"SELECT authn_id();")
+ resp = receive_until(conn, pq3.types.DataRow)
+
+ row = resp.payload
+ assert row.columns == [authn_id]
+
+
+class ExpectedError(object):
+ def __init__(self, code, msg=None, detail=None):
+ self.code = code
+ self.msg = msg
+ self.detail = detail
+
+ # Protect against the footgun of an accidental empty string, which will
+ # "match" anything. If you don't want to match message or detail, just
+ # don't pass them.
+ if self.msg == "":
+ raise ValueError("msg must be non-empty or None")
+ if self.detail == "":
+ raise ValueError("detail must be non-empty or None")
+
+ def _getfield(self, resp, type):
+ """
+ Searches an ErrorResponse for a single field of the given type (e.g.
+ "M", "C", "D") and returns its value. Asserts if it doesn't find exactly
+ one field.
+ """
+ prefix = type.encode("ascii")
+ fields = [f for f in resp.payload.fields if f.startswith(prefix)]
+
+ assert len(fields) == 1
+ return fields[0][1:] # strip off the type byte
+
+ def match(self, resp):
+ """
+ Checks that the given response matches the expected code, message, and
+ detail (if given). The error code must match exactly. The expected
+ message and detail must be contained within the actual strings.
+ """
+ assert resp.type == pq3.types.ErrorResponse
+
+ code = self._getfield(resp, "C")
+ assert code == self.code
+
+ if self.msg:
+ msg = self._getfield(resp, "M")
+ expected = self.msg.encode("utf-8")
+ assert expected in msg
+
+ if self.detail:
+ detail = self._getfield(resp, "D")
+ expected = self.detail.encode("utf-8")
+ assert expected in detail
+
+
+def test_oauth_rejected_bearer(conn, oauth_ctx, bearer_token):
+ # Generate a new bearer token, which we will proceed not to use.
+ _ = bearer_token()
+
+ begin_oauth_handshake(conn, oauth_ctx)
+
+ # Send a bearer token that doesn't match what the validator expects. It
+ # should fail the connection.
+ send_initial_response(conn, bearer=b"xxxxxx")
+
+ expect_handshake_failure(conn, oauth_ctx)
+
+
+@pytest.mark.parametrize(
+ "bad_bearer",
+ [
+ b"Bearer ",
+ b"Bearer a===b",
+ b"Bearer hello!",
+ b"Bearer me@example.com",
+ b'OAuth realm="Example"',
+ b"",
+ ],
+)
+def test_oauth_invalid_bearer(conn, oauth_ctx, bearer_token, bad_bearer):
+ # Tell the validator to accept any token. This ensures that the invalid
+ # bearer tokens are rejected before the validation step.
+ _ = bearer_token(accept_any=True)
+
+ begin_oauth_handshake(conn, oauth_ctx)
+ send_initial_response(conn, auth=bad_bearer)
+
+ expect_handshake_failure(conn, oauth_ctx)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize(
+ "resp_type,resp,err",
+ [
+ pytest.param(
+ None,
+ None,
+ None,
+ marks=pytest.mark.slow,
+ id="no response (expect timeout)",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ b"hello",
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "did not send a kvsep response",
+ ),
+ id="bad dummy response",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ b"\x01\x01",
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "did not send a kvsep response",
+ ),
+ id="multiple kvseps",
+ ),
+ pytest.param(
+ pq3.types.Query,
+ dict(query=b""),
+ ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "expected SASL response"),
+ id="bad response message type",
+ ),
+ ],
+)
+def test_oauth_bad_response_to_error_challenge(conn, oauth_ctx, resp_type, resp, err):
+ begin_oauth_handshake(conn, oauth_ctx)
+
+ # Send an empty auth initial response, which will force an authn failure.
+ send_initial_response(conn, auth=b"")
+
+ # We expect a discovery "challenge" back from the server before the authn
+ # failure message.
+ pkt = pq3.recv1(conn)
+ assert pkt.type == pq3.types.AuthnRequest
+
+ req = pkt.payload
+ assert req.type == pq3.authn.SASLContinue
+
+ body = json.loads(req.body)
+ assert body["status"] == "invalid_token"
+
+ if resp_type is None:
+ # Do not send the dummy response. We should time out and not get a
+ # response from the server.
+ with pytest.raises(socket.timeout):
+ conn.read(1)
+
+ # Done with the test.
+ return
+
+ # Send the bad response.
+ pq3.send(conn, resp_type, resp)
+
+ # Make sure the server fails the connection correctly.
+ pkt = pq3.recv1(conn)
+ err.match(pkt)
+
+
+@pytest.mark.parametrize(
+ "type,payload,err",
+ [
+ pytest.param(
+ pq3.types.ErrorResponse,
+ dict(fields=[b""]),
+ ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "expected SASL response"),
+ id="error response in initial message",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ b"x" * (MAX_SASL_MESSAGE_LENGTH + 1),
+ ExpectedError(
+ INVALID_AUTHORIZATION_ERRCODE, "bearer authentication failed"
+ ),
+ id="overlong initial response data",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"SCRAM-SHA-256")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE, "invalid SASL authentication mechanism"
+ ),
+ id="bad SASL mechanism selection",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", len=2, data=b"x")),
+ ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "insufficient data"),
+ id="SASL data underflow",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", len=0, data=b"x")),
+ ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "invalid message format"),
+ id="SASL data overflow",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "message is empty",
+ ),
+ id="empty",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"n,,\x01auth=\x01\x01\0")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "length does not match input length",
+ ),
+ id="contains null byte",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"\x01")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "Unexpected channel-binding flag", # XXX this is a bit strange
+ ),
+ id="initial error response",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"p=tls-server-end-point,,\x01")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "server does not support channel binding",
+ ),
+ id="uses channel binding",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"x,,\x01")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "Unexpected channel-binding flag",
+ ),
+ id="invalid channel binding specifier",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "Comma expected",
+ ),
+ id="bad GS2 header: missing channel binding terminator",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,a")),
+ ExpectedError(
+ FEATURE_NOT_SUPPORTED_ERRCODE,
+ "client uses authorization identity",
+ ),
+ id="bad GS2 header: authzid in use",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,b,")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "Unexpected attribute",
+ ),
+ id="bad GS2 header: extra attribute",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "Unexpected attribute 0x00", # XXX this is a bit strange
+ ),
+ id="bad GS2 header: missing authzid terminator",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,,")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "Key-value separator expected",
+ ),
+ id="missing initial kvsep",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,,")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "Key-value separator expected",
+ ),
+ id="missing initial kvsep",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"y,,\x01\x01")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "does not contain an auth value",
+ ),
+ id="missing auth value: empty key-value list",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"y,,\x01host=example.com\x01\x01")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "does not contain an auth value",
+ ),
+ id="missing auth value: other keys present",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"y,,\x01host=example.com")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "unterminated key/value pair",
+ ),
+ id="missing value terminator",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,,\x01")),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "did not contain a final terminator",
+ ),
+ id="missing list terminator: empty list",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"y,,\x01auth=Bearer 0\x01")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "did not contain a final terminator",
+ ),
+ id="missing list terminator: with auth value",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"y,,\x01auth=Bearer 0\x01\x01blah")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "additional data after the final terminator",
+ ),
+ id="additional key after terminator",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(name=b"OAUTHBEARER", data=b"y,,\x01key\x01\x01")
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "key without a value",
+ ),
+ id="key without value",
+ ),
+ pytest.param(
+ pq3.types.PasswordMessage,
+ pq3.SASLInitialResponse.build(
+ dict(
+ name=b"OAUTHBEARER",
+ data=b"y,,\x01auth=Bearer 0\x01auth=Bearer 1\x01\x01",
+ )
+ ),
+ ExpectedError(
+ PROTOCOL_VIOLATION_ERRCODE,
+ "malformed OAUTHBEARER message",
+ "contains multiple auth values",
+ ),
+ id="multiple auth values",
+ ),
+ ],
+)
+def test_oauth_bad_initial_response(conn, oauth_ctx, type, payload, err):
+ begin_oauth_handshake(conn, oauth_ctx)
+
+ # The server expects a SASL response; give it something else instead.
+ if not isinstance(payload, dict):
+ payload = dict(payload_data=payload)
+ pq3.send(conn, type, **payload)
+
+ resp = pq3.recv1(conn)
+ err.match(resp)
+
+
+def test_oauth_empty_initial_response(conn, oauth_ctx, bearer_token):
+ begin_oauth_handshake(conn, oauth_ctx)
+
+ # Send an initial response without data.
+ initial = pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER"))
+ pq3.send(conn, pq3.types.PasswordMessage, initial)
+
+ # The server should respond with an empty challenge so we can send the data
+ # it wants.
+ pkt = pq3.recv1(conn)
+
+ assert pkt.type == pq3.types.AuthnRequest
+ assert pkt.payload.type == pq3.authn.SASLContinue
+ assert not pkt.payload.body
+
+ # Now send the initial data.
+ data = b"n,,\x01auth=Bearer " + bearer_token() + b"\x01\x01"
+ pq3.send(conn, pq3.types.PasswordMessage, data)
+
+ # Server should now complete the handshake.
+ expect_handshake_success(conn)
+
+
+@pytest.fixture()
+def set_validator():
+ """
+ A per-test fixture that allows a test to override the setting of
+ oauth_validator_command for the cluster. The setting will be reverted during
+ teardown.
+
+ Passing None will perform an ALTER SYSTEM RESET.
+ """
+ conn = psycopg2.connect("")
+ conn.autocommit = True
+
+ with contextlib.closing(conn):
+ c = conn.cursor()
+
+ # Save the previous value.
+ c.execute("SHOW oauth_validator_command;")
+ prev_cmd = c.fetchone()[0]
+
+ def setter(cmd):
+ c.execute("ALTER SYSTEM SET oauth_validator_command TO %s;", (cmd,))
+ c.execute("SELECT pg_reload_conf();")
+
+ yield setter
+
+ # Restore the previous value.
+ c.execute("ALTER SYSTEM SET oauth_validator_command TO %s;", (prev_cmd,))
+ c.execute("SELECT pg_reload_conf();")
+
+
+def test_oauth_no_validator(oauth_ctx, set_validator, connect, bearer_token):
+ # Clear out our validator command, then establish a new connection.
+ set_validator("")
+ conn = connect()
+
+ begin_oauth_handshake(conn, oauth_ctx)
+ send_initial_response(conn, bearer=bearer_token())
+
+ # The server should fail the connection.
+ expect_handshake_failure(conn, oauth_ctx)
+
+
+def test_oauth_validator_role(oauth_ctx, set_validator, connect):
+ # Switch the validator implementation. This validator will reflect the
+ # PGUSER as the authenticated identity.
+ path = pathlib.Path(__file__).parent / "validate_reflect.py"
+ path = str(path.absolute())
+
+ set_validator(f"{shlex.quote(path)} '%r' <&%f")
+ conn = connect()
+
+ # Log in. Note that the reflection validator ignores the bearer token.
+ begin_oauth_handshake(conn, oauth_ctx, user=oauth_ctx.user)
+ send_initial_response(conn, bearer=b"dontcare")
+ expect_handshake_success(conn)
+
+ # Check the user identity.
+ pq3.send(conn, pq3.types.Query, query=b"SELECT authn_id();")
+ resp = receive_until(conn, pq3.types.DataRow)
+
+ row = resp.payload
+ expected = oauth_ctx.user.encode("utf-8")
+ assert row.columns == [expected]
+
+
+def test_oauth_role_with_shell_unsafe_characters(oauth_ctx, set_validator, connect):
+ """
+ XXX This test pins undesirable behavior. We should be able to handle any
+ valid Postgres role name.
+ """
+ # Switch the validator implementation. This validator will reflect the
+ # PGUSER as the authenticated identity.
+ path = pathlib.Path(__file__).parent / "validate_reflect.py"
+ path = str(path.absolute())
+
+ set_validator(f"{shlex.quote(path)} '%r' <&%f")
+ conn = connect()
+
+ unsafe_username = "hello'there"
+ begin_oauth_handshake(conn, oauth_ctx, user=unsafe_username)
+
+ # The server should reject the handshake.
+ send_initial_response(conn, bearer=b"dontcare")
+ expect_handshake_failure(conn, oauth_ctx)
diff --git a/src/test/python/server/test_server.py b/src/test/python/server/test_server.py
new file mode 100644
index 0000000000..02126dba79
--- /dev/null
+++ b/src/test/python/server/test_server.py
@@ -0,0 +1,21 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import pq3
+
+
+def test_handshake(connect):
+ """Basic sanity check."""
+ conn = connect()
+
+ pq3.handshake(conn, user=pq3.pguser(), database=pq3.pgdatabase())
+
+ pq3.send(conn, pq3.types.Query, query=b"")
+
+ resp = pq3.recv1(conn)
+ assert resp.type == pq3.types.EmptyQueryResponse
+
+ resp = pq3.recv1(conn)
+ assert resp.type == pq3.types.ReadyForQuery
diff --git a/src/test/python/server/validate_bearer.py b/src/test/python/server/validate_bearer.py
new file mode 100755
index 0000000000..2cc73ff154
--- /dev/null
+++ b/src/test/python/server/validate_bearer.py
@@ -0,0 +1,101 @@
+#! /usr/bin/env python3
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+# DO NOT USE THIS OAUTH VALIDATOR IN PRODUCTION. It doesn't actually validate
+# anything, and it logs the bearer token data, which is sensitive.
+#
+# This executable is used as an oauth_validator_command in concert with
+# test_oauth.py. Memory is shared and communicated from that test module's
+# bearer_token() fixture.
+#
+# This script must run under the Postgres server environment; keep the
+# dependency list fairly standard.
+
+import base64
+import binascii
+import contextlib
+import struct
+import sys
+from multiprocessing import shared_memory
+
+MAX_UINT16 = 2 ** 16 - 1
+
+
+def remove_shm_from_resource_tracker():
+ """
+ Monkey-patch multiprocessing.resource_tracker so SharedMemory won't be
+ tracked. Pulled from this thread, where there are more details:
+
+ https://bugs.python.org/issue38119
+
+ TL;DR: all clients of shared memory segments automatically destroy them on
+ process exit, which makes shared memory segments much less useful. This
+ monkeypatch removes that behavior so that we can defer to the test to manage
+ the segment lifetime.
+
+ Ideally a future Python patch will pull in this fix and then the entire
+ function can go away.
+ """
+ from multiprocessing import resource_tracker
+
+ def fix_register(name, rtype):
+ if rtype == "shared_memory":
+ return
+ return resource_tracker._resource_tracker.register(self, name, rtype)
+
+ resource_tracker.register = fix_register
+
+ def fix_unregister(name, rtype):
+ if rtype == "shared_memory":
+ return
+ return resource_tracker._resource_tracker.unregister(self, name, rtype)
+
+ resource_tracker.unregister = fix_unregister
+
+ if "shared_memory" in resource_tracker._CLEANUP_FUNCS:
+ del resource_tracker._CLEANUP_FUNCS["shared_memory"]
+
+
+def main(args):
+ remove_shm_from_resource_tracker() # XXX remove some day
+
+ # Get the expected token from the currently running test.
+ shared_mem_name = args[0]
+
+ mem = shared_memory.SharedMemory(shared_mem_name)
+ with contextlib.closing(mem):
+ # First two bytes are the token length.
+ size = struct.unpack("H", mem.buf[:2])[0]
+
+ if size == MAX_UINT16:
+ # Special case: the test wants us to accept any token.
+ sys.stderr.write("accepting token without validation\n")
+ return
+
+ # The remainder of the buffer contains the expected token.
+ assert size <= (mem.size - 2)
+ expected_token = mem.buf[2 : size + 2].tobytes()
+
+ mem.buf[:] = b"\0" * mem.size # scribble over the token
+
+ token = sys.stdin.buffer.read()
+ if token != expected_token:
+ sys.exit(f"failed to match Bearer token ({token!r} != {expected_token!r})")
+
+ # See if the test wants us to print anything. If so, it will have encoded
+ # the desired output in the token with an "output=" prefix.
+ try:
+ # altchars="-_" corresponds to the urlsafe alphabet.
+ data = base64.b64decode(token, altchars="-_", validate=True)
+
+ if data.startswith(b"output="):
+ sys.stdout.buffer.write(data[7:])
+
+ except binascii.Error:
+ pass
+
+
+if __name__ == "__main__":
+ main(sys.argv[1:])
diff --git a/src/test/python/server/validate_reflect.py b/src/test/python/server/validate_reflect.py
new file mode 100755
index 0000000000..24c3a7e715
--- /dev/null
+++ b/src/test/python/server/validate_reflect.py
@@ -0,0 +1,34 @@
+#! /usr/bin/env python3
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+# DO NOT USE THIS OAUTH VALIDATOR IN PRODUCTION. It ignores the bearer token
+# entirely and automatically logs the user in.
+#
+# This executable is used as an oauth_validator_command in concert with
+# test_oauth.py. It expects the user's desired role name as an argument; the
+# actual token will be discarded and the user will be logged in with the role
+# name as the authenticated identity.
+#
+# This script must run under the Postgres server environment; keep the
+# dependency list fairly standard.
+
+import sys
+
+
+def main(args):
+ # We have to read the entire token as our first action to unblock the
+ # server, but we won't actually use it.
+ _ = sys.stdin.buffer.read()
+
+ if len(args) != 1:
+ sys.exit("usage: ./validate_reflect.py ROLE")
+
+ # Log the user in as the provided role.
+ role = args[0]
+ print(role)
+
+
+if __name__ == "__main__":
+ main(sys.argv[1:])
diff --git a/src/test/python/test_internals.py b/src/test/python/test_internals.py
new file mode 100644
index 0000000000..dee4855fc0
--- /dev/null
+++ b/src/test/python/test_internals.py
@@ -0,0 +1,138 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import io
+
+from pq3 import _DebugStream
+
+
+def test_DebugStream_read():
+ under = io.BytesIO(b"abcdefghijklmnopqrstuvwxyz")
+ out = io.StringIO()
+
+ stream = _DebugStream(under, out)
+
+ res = stream.read(5)
+ assert res == b"abcde"
+
+ res = stream.read(16)
+ assert res == b"fghijklmnopqrstu"
+
+ stream.flush_debug()
+
+ res = stream.read()
+ assert res == b"vwxyz"
+
+ stream.flush_debug()
+
+ expected = (
+ "< 0000:\t61 62 63 64 65 66 67 68 69 6a 6b 6c 6d 6e 6f 70\tabcdefghijklmnop\n"
+ "< 0010:\t71 72 73 74 75 \tqrstu\n"
+ "\n"
+ "< 0000:\t76 77 78 79 7a \tvwxyz\n"
+ "\n"
+ )
+ assert out.getvalue() == expected
+
+
+def test_DebugStream_write():
+ under = io.BytesIO()
+ out = io.StringIO()
+
+ stream = _DebugStream(under, out)
+
+ stream.write(b"\x00\x01\x02")
+ stream.flush()
+
+ assert under.getvalue() == b"\x00\x01\x02"
+
+ stream.write(b"\xc0\xc1\xc2")
+ stream.flush()
+
+ assert under.getvalue() == b"\x00\x01\x02\xc0\xc1\xc2"
+
+ stream.flush_debug()
+
+ expected = "> 0000:\t00 01 02 c0 c1 c2 \t......\n\n"
+ assert out.getvalue() == expected
+
+
+def test_DebugStream_read_write():
+ under = io.BytesIO(b"abcdefghijklmnopqrstuvwxyz")
+ out = io.StringIO()
+ stream = _DebugStream(under, out)
+
+ res = stream.read(5)
+ assert res == b"abcde"
+
+ stream.write(b"xxxxx")
+ stream.flush()
+
+ assert under.getvalue() == b"abcdexxxxxklmnopqrstuvwxyz"
+
+ res = stream.read(5)
+ assert res == b"klmno"
+
+ stream.write(b"xxxxx")
+ stream.flush()
+
+ assert under.getvalue() == b"abcdexxxxxklmnoxxxxxuvwxyz"
+
+ stream.flush_debug()
+
+ expected = (
+ "< 0000:\t61 62 63 64 65 6b 6c 6d 6e 6f \tabcdeklmno\n"
+ "\n"
+ "> 0000:\t78 78 78 78 78 78 78 78 78 78 \txxxxxxxxxx\n"
+ "\n"
+ )
+ assert out.getvalue() == expected
+
+
+def test_DebugStream_end_packet():
+ under = io.BytesIO(b"abcdefghijklmnopqrstuvwxyz")
+ out = io.StringIO()
+ stream = _DebugStream(under, out)
+
+ stream.read(5)
+ stream.end_packet("read description", read=True, indent=" ")
+
+ stream.write(b"xxxxx")
+ stream.flush()
+ stream.end_packet("write description", indent=" ")
+
+ stream.read(5)
+ stream.write(b"xxxxx")
+ stream.flush()
+ stream.end_packet("read/write combo for read", read=True, indent=" ")
+
+ stream.read(5)
+ stream.write(b"xxxxx")
+ stream.flush()
+ stream.end_packet("read/write combo for write", indent=" ")
+
+ expected = (
+ " < 0000:\t61 62 63 64 65 \tabcde\n"
+ "\n"
+ "< read description\n"
+ "\n"
+ "> write description\n"
+ "\n"
+ " > 0000:\t78 78 78 78 78 \txxxxx\n"
+ "\n"
+ " < 0000:\t6b 6c 6d 6e 6f \tklmno\n"
+ "\n"
+ " > 0000:\t78 78 78 78 78 \txxxxx\n"
+ "\n"
+ "< read/write combo for read\n"
+ "\n"
+ "> read/write combo for write\n"
+ "\n"
+ " < 0000:\t75 76 77 78 79 \tuvwxy\n"
+ "\n"
+ " > 0000:\t78 78 78 78 78 \txxxxx\n"
+ "\n"
+ )
+ assert out.getvalue() == expected
diff --git a/src/test/python/test_pq3.py b/src/test/python/test_pq3.py
new file mode 100644
index 0000000000..e0c0e0568d
--- /dev/null
+++ b/src/test/python/test_pq3.py
@@ -0,0 +1,558 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import contextlib
+import getpass
+import io
+import struct
+import sys
+
+import pytest
+from construct import Container, PaddingError, StreamError, TerminatedError
+
+import pq3
+
+
+@pytest.mark.parametrize(
+ "raw,expected,extra",
+ [
+ pytest.param(
+ b"\x00\x00\x00\x10\x00\x04\x00\x00abcdefgh",
+ Container(len=16, proto=0x40000, payload=b"abcdefgh"),
+ b"",
+ id="8-byte payload",
+ ),
+ pytest.param(
+ b"\x00\x00\x00\x08\x00\x04\x00\x00",
+ Container(len=8, proto=0x40000, payload=b""),
+ b"",
+ id="no payload",
+ ),
+ pytest.param(
+ b"\x00\x00\x00\x09\x00\x04\x00\x00abcde",
+ Container(len=9, proto=0x40000, payload=b"a"),
+ b"bcde",
+ id="1-byte payload and extra padding",
+ ),
+ pytest.param(
+ b"\x00\x00\x00\x0B\x00\x03\x00\x00hi\x00",
+ Container(len=11, proto=pq3.protocol(3, 0), payload=[b"hi"]),
+ b"",
+ id="implied parameter list when using proto version 3.0",
+ ),
+ ],
+)
+def test_Startup_parse(raw, expected, extra):
+ with io.BytesIO(raw) as stream:
+ actual = pq3.Startup.parse_stream(stream)
+
+ assert actual == expected
+ assert stream.read() == extra
+
+
+@pytest.mark.parametrize(
+ "packet,expected_bytes",
+ [
+ pytest.param(
+ dict(),
+ b"\x00\x00\x00\x08\x00\x00\x00\x00",
+ id="nothing set",
+ ),
+ pytest.param(
+ dict(len=10, proto=0x12345678),
+ b"\x00\x00\x00\x0A\x12\x34\x56\x78\x00\x00",
+ id="len and proto set explicitly",
+ ),
+ pytest.param(
+ dict(proto=0x12345678),
+ b"\x00\x00\x00\x08\x12\x34\x56\x78",
+ id="implied len with no payload",
+ ),
+ pytest.param(
+ dict(proto=0x12345678, payload=b"abcd"),
+ b"\x00\x00\x00\x0C\x12\x34\x56\x78abcd",
+ id="implied len with payload",
+ ),
+ pytest.param(
+ dict(payload=[b""]),
+ b"\x00\x00\x00\x09\x00\x03\x00\x00\x00",
+ id="implied proto version 3 when sending parameters",
+ ),
+ pytest.param(
+ dict(payload=[b"hi", b""]),
+ b"\x00\x00\x00\x0C\x00\x03\x00\x00hi\x00\x00",
+ id="implied proto version 3 and len when sending more than one parameter",
+ ),
+ pytest.param(
+ dict(payload=dict(user="jsmith", database="postgres")),
+ b"\x00\x00\x00\x27\x00\x03\x00\x00user\x00jsmith\x00database\x00postgres\x00\x00",
+ id="auto-serialization of dict parameters",
+ ),
+ ],
+)
+def test_Startup_build(packet, expected_bytes):
+ actual = pq3.Startup.build(packet)
+ assert actual == expected_bytes
+
+
+@pytest.mark.parametrize(
+ "raw,expected,extra",
+ [
+ pytest.param(
+ b"*\x00\x00\x00\x08abcd",
+ dict(type=b"*", len=8, payload=b"abcd"),
+ b"",
+ id="4-byte payload",
+ ),
+ pytest.param(
+ b"*\x00\x00\x00\x04",
+ dict(type=b"*", len=4, payload=b""),
+ b"",
+ id="no payload",
+ ),
+ pytest.param(
+ b"*\x00\x00\x00\x05xabcd",
+ dict(type=b"*", len=5, payload=b"x"),
+ b"abcd",
+ id="1-byte payload with extra padding",
+ ),
+ pytest.param(
+ b"R\x00\x00\x00\x08\x00\x00\x00\x00",
+ dict(
+ type=pq3.types.AuthnRequest,
+ len=8,
+ payload=dict(type=pq3.authn.OK, body=None),
+ ),
+ b"",
+ id="AuthenticationOk",
+ ),
+ pytest.param(
+ b"R\x00\x00\x00\x12\x00\x00\x00\x0AEXTERNAL\x00\x00",
+ dict(
+ type=pq3.types.AuthnRequest,
+ len=18,
+ payload=dict(type=pq3.authn.SASL, body=[b"EXTERNAL", b""]),
+ ),
+ b"",
+ id="AuthenticationSASL",
+ ),
+ pytest.param(
+ b"R\x00\x00\x00\x0D\x00\x00\x00\x0B12345",
+ dict(
+ type=pq3.types.AuthnRequest,
+ len=13,
+ payload=dict(type=pq3.authn.SASLContinue, body=b"12345"),
+ ),
+ b"",
+ id="AuthenticationSASLContinue",
+ ),
+ pytest.param(
+ b"R\x00\x00\x00\x0D\x00\x00\x00\x0C12345",
+ dict(
+ type=pq3.types.AuthnRequest,
+ len=13,
+ payload=dict(type=pq3.authn.SASLFinal, body=b"12345"),
+ ),
+ b"",
+ id="AuthenticationSASLFinal",
+ ),
+ pytest.param(
+ b"p\x00\x00\x00\x0Bhunter2",
+ dict(
+ type=pq3.types.PasswordMessage,
+ len=11,
+ payload=b"hunter2",
+ ),
+ b"",
+ id="PasswordMessage",
+ ),
+ pytest.param(
+ b"K\x00\x00\x00\x0C\x00\x00\x00\x00\x12\x34\x56\x78",
+ dict(
+ type=pq3.types.BackendKeyData,
+ len=12,
+ payload=dict(pid=0, key=0x12345678),
+ ),
+ b"",
+ id="BackendKeyData",
+ ),
+ pytest.param(
+ b"C\x00\x00\x00\x08SET\x00",
+ dict(
+ type=pq3.types.CommandComplete,
+ len=8,
+ payload=dict(tag=b"SET"),
+ ),
+ b"",
+ id="CommandComplete",
+ ),
+ pytest.param(
+ b"E\x00\x00\x00\x11Mbad!\x00Mdog!\x00\x00",
+ dict(type=b"E", len=17, payload=dict(fields=[b"Mbad!", b"Mdog!", b""])),
+ b"",
+ id="ErrorResponse",
+ ),
+ pytest.param(
+ b"S\x00\x00\x00\x08a\x00b\x00",
+ dict(
+ type=pq3.types.ParameterStatus,
+ len=8,
+ payload=dict(name=b"a", value=b"b"),
+ ),
+ b"",
+ id="ParameterStatus",
+ ),
+ pytest.param(
+ b"Z\x00\x00\x00\x05x",
+ dict(type=b"Z", len=5, payload=dict(status=b"x")),
+ b"",
+ id="ReadyForQuery",
+ ),
+ pytest.param(
+ b"Q\x00\x00\x00\x06!\x00",
+ dict(type=pq3.types.Query, len=6, payload=dict(query=b"!")),
+ b"",
+ id="Query",
+ ),
+ pytest.param(
+ b"D\x00\x00\x00\x0B\x00\x01\x00\x00\x00\x01!",
+ dict(type=pq3.types.DataRow, len=11, payload=dict(columns=[b"!"])),
+ b"",
+ id="DataRow",
+ ),
+ pytest.param(
+ b"D\x00\x00\x00\x06\x00\x00extra",
+ dict(type=pq3.types.DataRow, len=6, payload=dict(columns=[])),
+ b"extra",
+ id="DataRow with extra data",
+ ),
+ pytest.param(
+ b"I\x00\x00\x00\x04",
+ dict(type=pq3.types.EmptyQueryResponse, len=4, payload=None),
+ b"",
+ id="EmptyQueryResponse",
+ ),
+ pytest.param(
+ b"I\x00\x00\x00\x04\xFF",
+ dict(type=b"I", len=4, payload=None),
+ b"\xFF",
+ id="EmptyQueryResponse with extra bytes",
+ ),
+ pytest.param(
+ b"X\x00\x00\x00\x04",
+ dict(type=pq3.types.Terminate, len=4, payload=None),
+ b"",
+ id="Terminate",
+ ),
+ ],
+)
+def test_Pq3_parse(raw, expected, extra):
+ with io.BytesIO(raw) as stream:
+ actual = pq3.Pq3.parse_stream(stream)
+
+ assert actual == expected
+ assert stream.read() == extra
+
+
+@pytest.mark.parametrize(
+ "fields,expected",
+ [
+ pytest.param(
+ dict(type=b"*", len=5),
+ b"*\x00\x00\x00\x05\x00",
+ id="type and len set explicitly",
+ ),
+ pytest.param(
+ dict(type=b"*"),
+ b"*\x00\x00\x00\x04",
+ id="implied len with no payload",
+ ),
+ pytest.param(
+ dict(type=b"*", payload=b"1234"),
+ b"*\x00\x00\x00\x081234",
+ id="implied len with payload",
+ ),
+ pytest.param(
+ dict(type=pq3.types.AuthnRequest, payload=dict(type=pq3.authn.OK)),
+ b"R\x00\x00\x00\x08\x00\x00\x00\x00",
+ id="implied len/type for AuthenticationOK",
+ ),
+ pytest.param(
+ dict(
+ type=pq3.types.AuthnRequest,
+ payload=dict(
+ type=pq3.authn.SASL,
+ body=[b"SCRAM-SHA-256-PLUS", b"SCRAM-SHA-256", b""],
+ ),
+ ),
+ b"R\x00\x00\x00\x2A\x00\x00\x00\x0ASCRAM-SHA-256-PLUS\x00SCRAM-SHA-256\x00\x00",
+ id="implied len/type for AuthenticationSASL",
+ ),
+ pytest.param(
+ dict(
+ type=pq3.types.AuthnRequest,
+ payload=dict(type=pq3.authn.SASLContinue, body=b"12345"),
+ ),
+ b"R\x00\x00\x00\x0D\x00\x00\x00\x0B12345",
+ id="implied len/type for AuthenticationSASLContinue",
+ ),
+ pytest.param(
+ dict(
+ type=pq3.types.AuthnRequest,
+ payload=dict(type=pq3.authn.SASLFinal, body=b"12345"),
+ ),
+ b"R\x00\x00\x00\x0D\x00\x00\x00\x0C12345",
+ id="implied len/type for AuthenticationSASLFinal",
+ ),
+ pytest.param(
+ dict(
+ type=pq3.types.PasswordMessage,
+ payload=b"hunter2",
+ ),
+ b"p\x00\x00\x00\x0Bhunter2",
+ id="implied len/type for PasswordMessage",
+ ),
+ pytest.param(
+ dict(type=pq3.types.BackendKeyData, payload=dict(pid=1, key=7)),
+ b"K\x00\x00\x00\x0C\x00\x00\x00\x01\x00\x00\x00\x07",
+ id="implied len/type for BackendKeyData",
+ ),
+ pytest.param(
+ dict(type=pq3.types.CommandComplete, payload=dict(tag=b"SET")),
+ b"C\x00\x00\x00\x08SET\x00",
+ id="implied len/type for CommandComplete",
+ ),
+ pytest.param(
+ dict(type=pq3.types.ErrorResponse, payload=dict(fields=[b"error", b""])),
+ b"E\x00\x00\x00\x0Berror\x00\x00",
+ id="implied len/type for ErrorResponse",
+ ),
+ pytest.param(
+ dict(type=pq3.types.ParameterStatus, payload=dict(name=b"a", value=b"b")),
+ b"S\x00\x00\x00\x08a\x00b\x00",
+ id="implied len/type for ParameterStatus",
+ ),
+ pytest.param(
+ dict(type=pq3.types.ReadyForQuery, payload=dict(status=b"I")),
+ b"Z\x00\x00\x00\x05I",
+ id="implied len/type for ReadyForQuery",
+ ),
+ pytest.param(
+ dict(type=pq3.types.Query, payload=dict(query=b"SELECT 1;")),
+ b"Q\x00\x00\x00\x0eSELECT 1;\x00",
+ id="implied len/type for Query",
+ ),
+ pytest.param(
+ dict(type=pq3.types.DataRow, payload=dict(columns=[b"abcd"])),
+ b"D\x00\x00\x00\x0E\x00\x01\x00\x00\x00\x04abcd",
+ id="implied len/type for DataRow",
+ ),
+ pytest.param(
+ dict(type=pq3.types.EmptyQueryResponse),
+ b"I\x00\x00\x00\x04",
+ id="implied len for EmptyQueryResponse",
+ ),
+ pytest.param(
+ dict(type=pq3.types.Terminate),
+ b"X\x00\x00\x00\x04",
+ id="implied len for Terminate",
+ ),
+ ],
+)
+def test_Pq3_build(fields, expected):
+ actual = pq3.Pq3.build(fields)
+ assert actual == expected
+
+
+@pytest.mark.parametrize(
+ "raw,expected,extra",
+ [
+ pytest.param(
+ b"\x00\x00",
+ dict(columns=[]),
+ b"",
+ id="no columns",
+ ),
+ pytest.param(
+ b"\x00\x01\x00\x00\x00\x04abcd",
+ dict(columns=[b"abcd"]),
+ b"",
+ id="one column",
+ ),
+ pytest.param(
+ b"\x00\x02\x00\x00\x00\x04abcd\x00\x00\x00\x01x",
+ dict(columns=[b"abcd", b"x"]),
+ b"",
+ id="multiple columns",
+ ),
+ pytest.param(
+ b"\x00\x02\x00\x00\x00\x00\x00\x00\x00\x01x",
+ dict(columns=[b"", b"x"]),
+ b"",
+ id="empty column value",
+ ),
+ pytest.param(
+ b"\x00\x02\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF",
+ dict(columns=[None, None]),
+ b"",
+ id="null columns",
+ ),
+ ],
+)
+def test_DataRow_parse(raw, expected, extra):
+ pkt = b"D" + struct.pack("!i", len(raw) + 4) + raw
+ with io.BytesIO(pkt) as stream:
+ actual = pq3.Pq3.parse_stream(stream)
+
+ assert actual.type == pq3.types.DataRow
+ assert actual.payload == expected
+ assert stream.read() == extra
+
+
+@pytest.mark.parametrize(
+ "fields,expected",
+ [
+ pytest.param(
+ dict(),
+ b"\x00\x00",
+ id="no columns",
+ ),
+ pytest.param(
+ dict(columns=[None, None]),
+ b"\x00\x02\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF",
+ id="null columns",
+ ),
+ ],
+)
+def test_DataRow_build(fields, expected):
+ actual = pq3.Pq3.build(dict(type=pq3.types.DataRow, payload=fields))
+
+ expected = b"D" + struct.pack("!i", len(expected) + 4) + expected
+ assert actual == expected
+
+
+@pytest.mark.parametrize(
+ "raw,expected,exception",
+ [
+ pytest.param(
+ b"EXTERNAL\x00\xFF\xFF\xFF\xFF",
+ dict(name=b"EXTERNAL", len=-1, data=None),
+ None,
+ id="no initial response",
+ ),
+ pytest.param(
+ b"EXTERNAL\x00\x00\x00\x00\x02me",
+ dict(name=b"EXTERNAL", len=2, data=b"me"),
+ None,
+ id="initial response",
+ ),
+ pytest.param(
+ b"EXTERNAL\x00\x00\x00\x00\x02meextra",
+ None,
+ TerminatedError,
+ id="extra data",
+ ),
+ pytest.param(
+ b"EXTERNAL\x00\x00\x00\x00\xFFme",
+ None,
+ StreamError,
+ id="underflow",
+ ),
+ ],
+)
+def test_SASLInitialResponse_parse(raw, expected, exception):
+ ctx = contextlib.nullcontext()
+ if exception:
+ ctx = pytest.raises(exception)
+
+ with ctx:
+ actual = pq3.SASLInitialResponse.parse(raw)
+ assert actual == expected
+
+
+@pytest.mark.parametrize(
+ "fields,expected",
+ [
+ pytest.param(
+ dict(name=b"EXTERNAL"),
+ b"EXTERNAL\x00\xFF\xFF\xFF\xFF",
+ id="no initial response",
+ ),
+ pytest.param(
+ dict(name=b"EXTERNAL", data=None),
+ b"EXTERNAL\x00\xFF\xFF\xFF\xFF",
+ id="no initial response (explicit None)",
+ ),
+ pytest.param(
+ dict(name=b"EXTERNAL", data=b""),
+ b"EXTERNAL\x00\x00\x00\x00\x00",
+ id="empty response",
+ ),
+ pytest.param(
+ dict(name=b"EXTERNAL", data=b"me@example.com"),
+ b"EXTERNAL\x00\x00\x00\x00\x0Eme@example.com",
+ id="initial response",
+ ),
+ pytest.param(
+ dict(name=b"EXTERNAL", len=2, data=b"me@example.com"),
+ b"EXTERNAL\x00\x00\x00\x00\x02me@example.com",
+ id="data overflow",
+ ),
+ pytest.param(
+ dict(name=b"EXTERNAL", len=14, data=b"me"),
+ b"EXTERNAL\x00\x00\x00\x00\x0Eme",
+ id="data underflow",
+ ),
+ ],
+)
+def test_SASLInitialResponse_build(fields, expected):
+ actual = pq3.SASLInitialResponse.build(fields)
+ assert actual == expected
+
+
+@pytest.mark.parametrize(
+ "version,expected_bytes",
+ [
+ pytest.param((3, 0), b"\x00\x03\x00\x00", id="version 3"),
+ pytest.param((1234, 5679), b"\x04\xd2\x16\x2f", id="SSLRequest"),
+ ],
+)
+def test_protocol(version, expected_bytes):
+ # Make sure the integer returned by protocol is correctly serialized on the
+ # wire.
+ assert struct.pack("!i", pq3.protocol(*version)) == expected_bytes
+
+
+@pytest.mark.parametrize(
+ "envvar,func,expected",
+ [
+ ("PGHOST", pq3.pghost, "localhost"),
+ ("PGPORT", pq3.pgport, 5432),
+ ("PGUSER", pq3.pguser, getpass.getuser()),
+ ("PGDATABASE", pq3.pgdatabase, "postgres"),
+ ],
+)
+def test_env_defaults(monkeypatch, envvar, func, expected):
+ monkeypatch.delenv(envvar, raising=False)
+
+ actual = func()
+ assert actual == expected
+
+
+@pytest.mark.parametrize(
+ "envvars,func,expected",
+ [
+ (dict(PGHOST="otherhost"), pq3.pghost, "otherhost"),
+ (dict(PGPORT="6789"), pq3.pgport, 6789),
+ (dict(PGUSER="postgres"), pq3.pguser, "postgres"),
+ (dict(PGDATABASE="template1"), pq3.pgdatabase, "template1"),
+ ],
+)
+def test_env(monkeypatch, envvars, func, expected):
+ for k, v in envvars.items():
+ monkeypatch.setenv(k, v)
+
+ actual = func()
+ assert actual == expected
diff --git a/src/test/python/tls.py b/src/test/python/tls.py
new file mode 100644
index 0000000000..075c02c1ca
--- /dev/null
+++ b/src/test/python/tls.py
@@ -0,0 +1,195 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+from construct import *
+
+#
+# TLS 1.3
+#
+# Most of the types below are transcribed from RFC 8446:
+#
+# https://tools.ietf.org/html/rfc8446
+#
+
+
+def _Vector(size_field, element):
+ return Prefixed(size_field, GreedyRange(element))
+
+
+# Alerts
+
+AlertLevel = Enum(
+ Byte,
+ warning=1,
+ fatal=2,
+)
+
+AlertDescription = Enum(
+ Byte,
+ close_notify=0,
+ unexpected_message=10,
+ bad_record_mac=20,
+ decryption_failed_RESERVED=21,
+ record_overflow=22,
+ decompression_failure=30,
+ handshake_failure=40,
+ no_certificate_RESERVED=41,
+ bad_certificate=42,
+ unsupported_certificate=43,
+ certificate_revoked=44,
+ certificate_expired=45,
+ certificate_unknown=46,
+ illegal_parameter=47,
+ unknown_ca=48,
+ access_denied=49,
+ decode_error=50,
+ decrypt_error=51,
+ export_restriction_RESERVED=60,
+ protocol_version=70,
+ insufficient_security=71,
+ internal_error=80,
+ user_canceled=90,
+ no_renegotiation=100,
+ unsupported_extension=110,
+)
+
+Alert = Struct(
+ "level" / AlertLevel,
+ "description" / AlertDescription,
+)
+
+
+# Extensions
+
+ExtensionType = Enum(
+ Int16ub,
+ server_name=0,
+ max_fragment_length=1,
+ status_request=5,
+ supported_groups=10,
+ signature_algorithms=13,
+ use_srtp=14,
+ heartbeat=15,
+ application_layer_protocol_negotiation=16,
+ signed_certificate_timestamp=18,
+ client_certificate_type=19,
+ server_certificate_type=20,
+ padding=21,
+ pre_shared_key=41,
+ early_data=42,
+ supported_versions=43,
+ cookie=44,
+ psk_key_exchange_modes=45,
+ certificate_authorities=47,
+ oid_filters=48,
+ post_handshake_auth=49,
+ signature_algorithms_cert=50,
+ key_share=51,
+)
+
+Extension = Struct(
+ "extension_type" / ExtensionType,
+ "extension_data" / Prefixed(Int16ub, GreedyBytes),
+)
+
+
+# ClientHello
+
+
+class _CipherSuiteAdapter(Adapter):
+ class _hextuple(tuple):
+ def __repr__(self):
+ return f"(0x{self[0]:02X}, 0x{self[1]:02X})"
+
+ def _encode(self, obj, context, path):
+ return bytes(obj)
+
+ def _decode(self, obj, context, path):
+ assert len(obj) == 2
+ return self._hextuple(obj)
+
+
+ProtocolVersion = Hex(Int16ub)
+
+Random = Hex(Bytes(32))
+
+CipherSuite = _CipherSuiteAdapter(Byte[2])
+
+ClientHello = Struct(
+ "legacy_version" / ProtocolVersion,
+ "random" / Random,
+ "legacy_session_id" / Prefixed(Byte, Hex(GreedyBytes)),
+ "cipher_suites" / _Vector(Int16ub, CipherSuite),
+ "legacy_compression_methods" / Prefixed(Byte, GreedyBytes),
+ "extensions" / _Vector(Int16ub, Extension),
+)
+
+# ServerHello
+
+ServerHello = Struct(
+ "legacy_version" / ProtocolVersion,
+ "random" / Random,
+ "legacy_session_id_echo" / Prefixed(Byte, Hex(GreedyBytes)),
+ "cipher_suite" / CipherSuite,
+ "legacy_compression_method" / Hex(Byte),
+ "extensions" / _Vector(Int16ub, Extension),
+)
+
+# Handshake
+
+HandshakeType = Enum(
+ Byte,
+ client_hello=1,
+ server_hello=2,
+ new_session_ticket=4,
+ end_of_early_data=5,
+ encrypted_extensions=8,
+ certificate=11,
+ certificate_request=13,
+ certificate_verify=15,
+ finished=20,
+ key_update=24,
+ message_hash=254,
+)
+
+Handshake = Struct(
+ "msg_type" / HandshakeType,
+ "length" / Int24ub,
+ "payload"
+ / Switch(
+ this.msg_type,
+ {
+ HandshakeType.client_hello: ClientHello,
+ HandshakeType.server_hello: ServerHello,
+ # HandshakeType.end_of_early_data: EndOfEarlyData,
+ # HandshakeType.encrypted_extensions: EncryptedExtensions,
+ # HandshakeType.certificate_request: CertificateRequest,
+ # HandshakeType.certificate: Certificate,
+ # HandshakeType.certificate_verify: CertificateVerify,
+ # HandshakeType.finished: Finished,
+ # HandshakeType.new_session_ticket: NewSessionTicket,
+ # HandshakeType.key_update: KeyUpdate,
+ },
+ default=FixedSized(this.length, GreedyBytes),
+ ),
+)
+
+# Records
+
+ContentType = Enum(
+ Byte,
+ invalid=0,
+ change_cipher_spec=20,
+ alert=21,
+ handshake=22,
+ application_data=23,
+)
+
+Plaintext = Struct(
+ "type" / ContentType,
+ "legacy_record_version" / ProtocolVersion,
+ "length" / Int16ub,
+ "fragment" / FixedSized(this.length, GreedyBytes),
+)
--
2.25.1