fix dbname

## Problem

The `WalAcceptor` main loop currently uses two nested loops to consume
inbound messages. This makes it hard to slot in periodic events like
metrics collection. It also duplicates the event processing code, and assumes
all messages in steady state are AppendRequests (other messages types may
be dropped if following an AppendRequest).

## Summary of changes

Refactor the `WalAcceptor` loop to be event driven.

.github/workflows/build_and_test.yml

@@ -839,6 +839,7 @@ jobs:
       - name: Build vm image
         run: |
           ./vm-builder \
+            -size=2G \
             -spec=compute/vm-image-spec-${{ matrix.version.debian }}.yaml \
             -src=neondatabase/compute-node-${{ matrix.version.pg }}:${{ needs.tag.outputs.build-tag }} \
             -dst=neondatabase/vm-compute-node-${{ matrix.version.pg }}:${{ needs.tag.outputs.build-tag }}
@@ -1118,7 +1119,6 @@ jobs:
               -f deployPgSniRouter=true \
               -f deployProxyLink=true \
               -f deployPrivatelinkProxy=true \
-              -f deployLegacyProxyScram=true \
               -f deployProxyScram=true \
               -f deployProxyAuthBroker=true \
               -f branch=main \

compute/compute-node.Dockerfile

@@ -666,7 +666,7 @@ RUN apt-get update && \
 #
 # Use new version only for v17
 # because Release_2024_09_1 has some backward incompatible changes
-# https://github.com/rdkit/rdkit/releases/tag/Release_2024_09_1 
+# https://github.com/rdkit/rdkit/releases/tag/Release_2024_09_1
 ENV PATH="/usr/local/pgsql/bin/:/usr/local/pgsql/:$PATH"
 RUN case "${PG_VERSION}" in \
     "v17") \
@@ -860,13 +860,14 @@ ENV PATH="/home/nonroot/.cargo/bin:/usr/local/pgsql/bin/:$PATH"
 USER nonroot
 WORKDIR /home/nonroot
 
-RUN case "${PG_VERSION}" in "v17") \
+RUN curl -sSO https://static.rust-lang.org/rustup/dist/$(uname -m)-unknown-linux-gnu/rustup-init && \
-    echo "v17 is not supported yet by pgrx. Quit" && exit 0;; \
-    esac && \
-    curl -sSO https://static.rust-lang.org/rustup/dist/$(uname -m)-unknown-linux-gnu/rustup-init && \
     chmod +x rustup-init && \
     ./rustup-init -y --no-modify-path --profile minimal --default-toolchain stable && \
     rm rustup-init && \
+    case "${PG_VERSION}" in \
+        'v17') \
+            echo 'v17 is not supported yet by pgrx. Quit' && exit 0;; \
+    esac && \
     cargo install --locked --version 0.11.3 cargo-pgrx && \
     /bin/bash -c 'cargo pgrx init --pg${PG_VERSION:1}=/usr/local/pgsql/bin/pg_config'
 
@@ -1041,6 +1042,31 @@ RUN wget https://github.com/pgpartman/pg_partman/archive/refs/tags/v5.1.0.tar.gz
     make -j $(getconf _NPROCESSORS_ONLN) install && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pg_partman.control
 
+#########################################################################################
+#
+# Layer "pg_mooncake"
+# compile pg_mooncake extension
+#
+#########################################################################################
+FROM rust-extensions-build AS pg-mooncake-build
+ARG PG_VERSION
+COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
+
+ENV PG_MOONCAKE_VERSION=0a7de4c0b5c7b1a5e2175e1c5f4625b97b7346f1
+ENV PATH="/usr/local/pgsql/bin/:$PATH"
+
+RUN case "${PG_VERSION}" in \
+        'v14') \
+            echo "pg_mooncake is not supported on Postgres ${PG_VERSION}" && exit 0;; \
+    esac && \
+    git clone --depth 1 --branch neon https://github.com/Mooncake-Labs/pg_mooncake.git pg_mooncake-src && \
+    cd pg_mooncake-src && \
+    git checkout "${PG_MOONCAKE_VERSION}" && \
+    git submodule update --init --depth 1 --recursive && \
+    make BUILD_TYPE=release -j $(getconf _NPROCESSORS_ONLN) && \
+    make BUILD_TYPE=release -j $(getconf _NPROCESSORS_ONLN) install && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pg_mooncake.control
+
 #########################################################################################
 #
 # Layer "neon-pg-ext-build"
@@ -1084,6 +1110,7 @@ COPY --from=wal2json-pg-build /usr/local/pgsql /usr/local/pgsql
 COPY --from=pg-anon-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-ivm-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-partman-build /usr/local/pgsql/ /usr/local/pgsql/
+COPY --from=pg-mooncake-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY pgxn/ pgxn/
 
 RUN make -j $(getconf _NPROCESSORS_ONLN) \

poetry.lock

@@ -1521,21 +1521,6 @@ files = [
 [package.dependencies]
 six = "*"
 
-[[package]]
-name = "jwcrypto"
-version = "1.5.6"
-description = "Implementation of JOSE Web standards"
-optional = false
-python-versions = ">= 3.8"
-files = [
-    {file = "jwcrypto-1.5.6-py3-none-any.whl", hash = "sha256:150d2b0ebbdb8f40b77f543fb44ffd2baeff48788be71f67f03566692fd55789"},
-    {file = "jwcrypto-1.5.6.tar.gz", hash = "sha256:771a87762a0c081ae6166958a954f80848820b2ab066937dc8b8379d65b1b039"},
-]
-
-[package.dependencies]
-cryptography = ">=3.4"
-typing-extensions = ">=4.5.0"
-
 [[package]]
 name = "kafka-python"
 version = "2.0.2"
@@ -2126,6 +2111,7 @@ files = [
     {file = "psycopg2_binary-2.9.9-cp311-cp311-win32.whl", hash = "sha256:dc4926288b2a3e9fd7b50dc6a1909a13bbdadfc67d93f3374d984e56f885579d"},
     {file = "psycopg2_binary-2.9.9-cp311-cp311-win_amd64.whl", hash = "sha256:b76bedd166805480ab069612119ea636f5ab8f8771e640ae103e05a4aae3e417"},
     {file = "psycopg2_binary-2.9.9-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:8532fd6e6e2dc57bcb3bc90b079c60de896d2128c5d9d6f24a63875a95a088cf"},
+    {file = "psycopg2_binary-2.9.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b0605eaed3eb239e87df0d5e3c6489daae3f7388d455d0c0b4df899519c6a38d"},
     {file = "psycopg2_binary-2.9.9-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8f8544b092a29a6ddd72f3556a9fcf249ec412e10ad28be6a0c0d948924f2212"},
     {file = "psycopg2_binary-2.9.9-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2d423c8d8a3c82d08fe8af900ad5b613ce3632a1249fd6a223941d0735fce493"},
     {file = "psycopg2_binary-2.9.9-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2e5afae772c00980525f6d6ecf7cbca55676296b580c0e6abb407f15f3706996"},
@@ -2134,6 +2120,8 @@ files = [
     {file = "psycopg2_binary-2.9.9-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:cb16c65dcb648d0a43a2521f2f0a2300f40639f6f8c1ecbc662141e4e3e1ee07"},
     {file = "psycopg2_binary-2.9.9-cp312-cp312-musllinux_1_1_ppc64le.whl", hash = "sha256:911dda9c487075abd54e644ccdf5e5c16773470a6a5d3826fda76699410066fb"},
     {file = "psycopg2_binary-2.9.9-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:57fede879f08d23c85140a360c6a77709113efd1c993923c59fde17aa27599fe"},
+    {file = "psycopg2_binary-2.9.9-cp312-cp312-win32.whl", hash = "sha256:64cf30263844fa208851ebb13b0732ce674d8ec6a0c86a4e160495d299ba3c93"},
+    {file = "psycopg2_binary-2.9.9-cp312-cp312-win_amd64.whl", hash = "sha256:81ff62668af011f9a48787564ab7eded4e9fb17a4a6a74af5ffa6a457400d2ab"},
     {file = "psycopg2_binary-2.9.9-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:2293b001e319ab0d869d660a704942c9e2cce19745262a8aba2115ef41a0a42a"},
     {file = "psycopg2_binary-2.9.9-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:03ef7df18daf2c4c07e2695e8cfd5ee7f748a1d54d802330985a78d2a5a6dca9"},
     {file = "psycopg2_binary-2.9.9-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0a602ea5aff39bb9fac6308e9c9d82b9a35c2bf288e184a816002c9fae930b77"},
@@ -2615,6 +2603,7 @@ files = [
     {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
     {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
     {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
+    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
     {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
     {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
     {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
@@ -2923,20 +2912,6 @@ files = [
     {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"},
 ]
 
-[[package]]
-name = "types-jwcrypto"
-version = "1.5.0.20240925"
-description = "Typing stubs for jwcrypto"
-optional = false
-python-versions = ">=3.8"
-files = [
-    {file = "types-jwcrypto-1.5.0.20240925.tar.gz", hash = "sha256:50e17b790378c96239344476c7bd13b52d0c7eeb6d16c2d53723e48cc6bbf4fe"},
-    {file = "types_jwcrypto-1.5.0.20240925-py3-none-any.whl", hash = "sha256:2d12a2d528240d326075e896aafec7056b9136bf3207fa6ccf3fcb8fbf9e11a1"},
-]
-
-[package.dependencies]
-cryptography = "*"
-
 [[package]]
 name = "types-psutil"
 version = "5.9.5.12"
@@ -3184,6 +3159,16 @@ files = [
     {file = "wrapt-1.14.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:8ad85f7f4e20964db4daadcab70b47ab05c7c1cf2a7c1e51087bfaa83831854c"},
     {file = "wrapt-1.14.1-cp310-cp310-win32.whl", hash = "sha256:a9a52172be0b5aae932bef82a79ec0a0ce87288c7d132946d645eba03f0ad8a8"},
     {file = "wrapt-1.14.1-cp310-cp310-win_amd64.whl", hash = "sha256:6d323e1554b3d22cfc03cd3243b5bb815a51f5249fdcbb86fda4bf62bab9e164"},
+    {file = "wrapt-1.14.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:ecee4132c6cd2ce5308e21672015ddfed1ff975ad0ac8d27168ea82e71413f55"},
+    {file = "wrapt-1.14.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2020f391008ef874c6d9e208b24f28e31bcb85ccff4f335f15a3251d222b92d9"},
+    {file = "wrapt-1.14.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2feecf86e1f7a86517cab34ae6c2f081fd2d0dac860cb0c0ded96d799d20b335"},
+    {file = "wrapt-1.14.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:240b1686f38ae665d1b15475966fe0472f78e71b1b4903c143a842659c8e4cb9"},
+    {file = "wrapt-1.14.1-cp311-cp311-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9008dad07d71f68487c91e96579c8567c98ca4c3881b9b113bc7b33e9fd78b8"},
+    {file = "wrapt-1.14.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:6447e9f3ba72f8e2b985a1da758767698efa72723d5b59accefd716e9e8272bf"},
+    {file = "wrapt-1.14.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:acae32e13a4153809db37405f5eba5bac5fbe2e2ba61ab227926a22901051c0a"},
+    {file = "wrapt-1.14.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:49ef582b7a1152ae2766557f0550a9fcbf7bbd76f43fbdc94dd3bf07cc7168be"},
+    {file = "wrapt-1.14.1-cp311-cp311-win32.whl", hash = "sha256:358fe87cc899c6bb0ddc185bf3dbfa4ba646f05b1b0b9b5a27c2cb92c2cea204"},
+    {file = "wrapt-1.14.1-cp311-cp311-win_amd64.whl", hash = "sha256:26046cd03936ae745a502abf44dac702a5e6880b2b01c29aea8ddf3353b68224"},
     {file = "wrapt-1.14.1-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:43ca3bbbe97af00f49efb06e352eae40434ca9d915906f77def219b88e85d907"},
     {file = "wrapt-1.14.1-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:6b1a564e6cb69922c7fe3a678b9f9a3c54e72b469875aa8018f18b4d1dd1adf3"},
     {file = "wrapt-1.14.1-cp35-cp35m-manylinux2010_i686.whl", hash = "sha256:00b6d4ea20a906c0ca56d84f93065b398ab74b927a7a3dbd470f6fc503f95dc3"},
@@ -3421,4 +3406,4 @@ cffi = ["cffi (>=1.11)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9"
-content-hash = "ad5c9ee7723359af22bbd7fa41538dcf78913c02e947a13a8f9a87eb3a59039e"
+content-hash = "0f4804119f417edf8e1fbd6d715d2e8d70ad731334fa9570304a2203f83339cf"

proxy/src/auth/backend/hacks.rs

@@ -1,5 +1,5 @@
 use tokio::io::{AsyncRead, AsyncWrite};
-use tracing::{info, warn};
+use tracing::{debug, info};
 
 use super::{ComputeCredentials, ComputeUserInfo, ComputeUserInfoNoEndpoint};
 use crate::auth::{self, AuthFlow};
@@ -21,7 +21,7 @@ pub(crate) async fn authenticate_cleartext(
     secret: AuthSecret,
     config: &'static AuthenticationConfig,
 ) -> auth::Result<ComputeCredentials> {
-    warn!("cleartext auth flow override is enabled, proceeding");
+    debug!("cleartext auth flow override is enabled, proceeding");
     ctx.set_auth_method(crate::context::AuthMethod::Cleartext);
 
     // pause the timer while we communicate with the client
@@ -61,7 +61,7 @@ pub(crate) async fn password_hack_no_authentication(
     info: ComputeUserInfoNoEndpoint,
     client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
 ) -> auth::Result<(ComputeUserInfo, Vec<u8>)> {
-    warn!("project not specified, resorting to the password hack auth flow");
+    debug!("project not specified, resorting to the password hack auth flow");
     ctx.set_auth_method(crate::context::AuthMethod::Cleartext);
 
     // pause the timer while we communicate with the client

proxy/src/bin/proxy.rs

@@ -137,9 +137,6 @@ struct ProxyCliArgs {
     /// size of the threadpool for password hashing
     #[clap(long, default_value_t = 4)]
     scram_thread_pool_size: u8,
-    /// Disable dynamic rate limiter and store the metrics to ensure its production behaviour.
-    #[clap(long, default_value_t = true, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
-    disable_dynamic_rate_limiter: bool,
     /// Endpoint rate limiter max number of requests per second.
     ///
     /// Provided in the form `<Requests Per Second>@<Bucket Duration Size>`.
@@ -615,9 +612,6 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
              and metric-collection-interval must be specified"
         ),
     };
-    if !args.disable_dynamic_rate_limiter {
-        bail!("dynamic rate limiter should be disabled");
-    }
 
     let config::ConcurrencyLockOptions {
         shards,

pyproject.toml

@@ -42,9 +42,6 @@ pytest-repeat = "^0.9.3"
 websockets = "^12.0"
 clickhouse-connect = "^0.7.16"
 kafka-python = "^2.0.2"
-jwcrypto = "^1.5.6"
-h2 = "^4.1.0"
-types-jwcrypto = "^1.5.0.20240925"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "==1.3.0"

safekeeper/src/receive_wal.rs

@@ -21,18 +21,15 @@ use postgres_backend::QueryError;
 use pq_proto::BeMessage;
 use serde::Deserialize;
 use serde::Serialize;
+use std::future;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
-use tokio::sync::mpsc::channel;
+use tokio::sync::mpsc::{channel, Receiver, Sender};
-use tokio::sync::mpsc::error::TryRecvError;
-use tokio::sync::mpsc::Receiver;
-use tokio::sync::mpsc::Sender;
 use tokio::task;
 use tokio::task::JoinHandle;
-use tokio::time::Duration;
+use tokio::time::{Duration, MissedTickBehavior};
-use tokio::time::Instant;
 use tracing::*;
 use utils::id::TenantTimelineId;
 use utils::lsn::Lsn;
@@ -444,9 +441,9 @@ async fn network_write<IO: AsyncRead + AsyncWrite + Unpin>(
     }
 }
 
-// Send keepalive messages to walproposer, to make sure it receives updates
+/// The WAL flush interval. This ensures we periodically flush the WAL and send AppendResponses to
-// even when it writes a steady stream of messages.
+/// walproposer, even when it's writing a steady stream of messages.
-const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(1);
+const FLUSH_INTERVAL: Duration = Duration::from_secs(1);
 
 /// Encapsulates a task which takes messages from msg_rx, processes and pushes
 /// replies to reply_tx.
@@ -494,67 +491,76 @@ impl WalAcceptor {
     async fn run(&mut self) -> anyhow::Result<()> {
         let walreceiver_guard = self.tli.get_walreceivers().register(self.conn_id);
 
-        // After this timestamp we will stop processing AppendRequests and send a response
+        // Periodically flush the WAL.
-        // to the walproposer. walproposer sends at least one AppendRequest per second,
+        let mut flush_ticker = tokio::time::interval(FLUSH_INTERVAL);
-        // we will send keepalives by replying to these requests once per second.
+        flush_ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
-        let mut next_keepalive = Instant::now();
+        flush_ticker.tick().await; // skip the initial, immediate tick
 
-        while let Some(mut next_msg) = self.msg_rx.recv().await {
+        // Tracks unflushed appends.
-            // Update walreceiver state in shmem for reporting.
+        let mut dirty = false;
-            if let ProposerAcceptorMessage::Elected(_) = &next_msg {
-                walreceiver_guard.get().status = WalReceiverStatus::Streaming;
-            }
 
-            let reply_msg = if matches!(next_msg, ProposerAcceptorMessage::AppendRequest(_)) {
+        loop {
-                // Loop through AppendRequests while available to write as many WAL records as
+            let reply = tokio::select! {
-                // possible without fsyncing.
+                // Process inbound message.
-                //
+                msg = self.msg_rx.recv() => {
-                // Make sure the WAL is flushed before returning, see:
+                    // If disconnected, break to flush WAL and return.
-                // https://github.com/neondatabase/neon/issues/9259
+                    let Some(mut msg) = msg else {
-                //
-                // Note: this will need to be rewritten if we want to read non-AppendRequest messages here.
-                // Otherwise, we might end up in a situation where we read a message, but don't
-                // process it.
-                while let ProposerAcceptorMessage::AppendRequest(append_request) = next_msg {
-                    let noflush_msg = ProposerAcceptorMessage::NoFlushAppendRequest(append_request);
-
-                    if let Some(reply) = self.tli.process_msg(&noflush_msg).await? {
-                        if self.reply_tx.send(reply).await.is_err() {
-                            break; // disconnected, flush WAL and return on next send/recv
-                        }
-                    }
-
-                    // get out of this loop if keepalive time is reached
-                    if Instant::now() >= next_keepalive {
                         break;
+                    };
+
+                    // Update walreceiver state in shmem for reporting.
+                    if let ProposerAcceptorMessage::Elected(_) = &msg {
+                        walreceiver_guard.get().status = WalReceiverStatus::Streaming;
                     }
 
-                    // continue pulling AppendRequests if available
+                    // Don't flush the WAL on every append, only periodically via flush_ticker.
-                    match self.msg_rx.try_recv() {
+                    // This batches multiple appends per fsync. If the channel is empty after
-                        Ok(msg) => next_msg = msg,
+                    // sending the reply, we'll schedule an immediate flush.
-                        Err(TryRecvError::Empty) => break,
+                    if let ProposerAcceptorMessage::AppendRequest(append_request) = msg {
-                        // on disconnect, flush WAL and return on next send/recv
+                        msg = ProposerAcceptorMessage::NoFlushAppendRequest(append_request);
-                        Err(TryRecvError::Disconnected) => break,
+                        dirty = true;
-                    };
+                    }
+
+                    self.tli.process_msg(&msg).await?
                 }
 
-                // flush all written WAL to the disk
+                // While receiving AppendRequests, flush the WAL periodically and respond with an
-                self.tli
+                // AppendResponse to let walproposer know we're still alive.
-                    .process_msg(&ProposerAcceptorMessage::FlushWAL)
+                _ = flush_ticker.tick(), if dirty => {
-                    .await?
+                    dirty = false;
-            } else {
+                    self.tli
-                // process message other than AppendRequest
+                        .process_msg(&ProposerAcceptorMessage::FlushWAL)
-                self.tli.process_msg(&next_msg).await?
+                        .await?
+                }
+
+                // If there are no pending messages, flush the WAL immediately.
+                //
+                // TODO: this should be done via flush_ticker.reset_immediately(), but that's always
+                // delayed by 1ms due to this bug: https://github.com/tokio-rs/tokio/issues/6866.
+                _ = future::ready(()), if dirty && self.msg_rx.is_empty() => {
+                    dirty = false;
+                    flush_ticker.reset();
+                    self.tli
+                        .process_msg(&ProposerAcceptorMessage::FlushWAL)
+                        .await?
+                }
             };
 
-            if let Some(reply) = reply_msg {
+            // Send reply, if any.
+            if let Some(reply) = reply {
                 if self.reply_tx.send(reply).await.is_err() {
-                    return Ok(()); // chan closed, streaming terminated
+                    break; // disconnected, break to flush WAL and return
                 }
-                // reset keepalive time
-                next_keepalive = Instant::now() + KEEPALIVE_INTERVAL;
             }
         }
+
+        // Flush WAL on disconnect, see https://github.com/neondatabase/neon/issues/9259.
+        if dirty {
+            self.tli
+                .process_msg(&ProposerAcceptorMessage::FlushWAL)
+                .await?;
+        }
+
         Ok(())
     }
 }

test_runner/conftest.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 pytest_plugins = (
     "fixtures.pg_version",
     "fixtures.parametrize",
-    "fixtures.h2server",
     "fixtures.httpserver",
     "fixtures.compute_reconfigure",
     "fixtures.storage_controller_proxy",

test_runner/fixtures/h2server.py

@@ -1,198 +0,0 @@
-"""
-https://python-hyper.org/projects/hyper-h2/en/stable/asyncio-example.html
-
-auth-broker -> local-proxy needs a h2 connection, so we need a h2 server :)
-"""
-
-import asyncio
-import collections
-import io
-import json
-from collections.abc import AsyncIterable
-
-import pytest_asyncio
-from h2.config import H2Configuration
-from h2.connection import H2Connection
-from h2.errors import ErrorCodes
-from h2.events import (
-    ConnectionTerminated,
-    DataReceived,
-    RemoteSettingsChanged,
-    RequestReceived,
-    StreamEnded,
-    StreamReset,
-    WindowUpdated,
-)
-from h2.exceptions import ProtocolError, StreamClosedError
-from h2.settings import SettingCodes
-
-RequestData = collections.namedtuple("RequestData", ["headers", "data"])
-
-
-class H2Server:
-    def __init__(self, host, port) -> None:
-        self.host = host
-        self.port = port
-
-
-class H2Protocol(asyncio.Protocol):
-    def __init__(self):
-        config = H2Configuration(client_side=False, header_encoding="utf-8")
-        self.conn = H2Connection(config=config)
-        self.transport = None
-        self.stream_data = {}
-        self.flow_control_futures = {}
-
-    def connection_made(self, transport: asyncio.Transport):  # type: ignore[override]
-        self.transport = transport
-        self.conn.initiate_connection()
-        self.transport.write(self.conn.data_to_send())
-
-    def connection_lost(self, _exc):
-        for future in self.flow_control_futures.values():
-            future.cancel()
-        self.flow_control_futures = {}
-
-    def data_received(self, data: bytes):
-        assert self.transport is not None
-        try:
-            events = self.conn.receive_data(data)
-        except ProtocolError:
-            self.transport.write(self.conn.data_to_send())
-            self.transport.close()
-        else:
-            self.transport.write(self.conn.data_to_send())
-            for event in events:
-                if isinstance(event, RequestReceived):
-                    self.request_received(event.headers, event.stream_id)
-                elif isinstance(event, DataReceived):
-                    self.receive_data(event.data, event.stream_id)
-                elif isinstance(event, StreamEnded):
-                    self.stream_complete(event.stream_id)
-                elif isinstance(event, ConnectionTerminated):
-                    self.transport.close()
-                elif isinstance(event, StreamReset):
-                    self.stream_reset(event.stream_id)
-                elif isinstance(event, WindowUpdated):
-                    self.window_updated(event.stream_id, event.delta)
-                elif isinstance(event, RemoteSettingsChanged):
-                    if SettingCodes.INITIAL_WINDOW_SIZE in event.changed_settings:
-                        self.window_updated(None, 0)
-
-                self.transport.write(self.conn.data_to_send())
-
-    def request_received(self, headers: list[tuple[str, str]], stream_id: int):
-        headers_map = collections.OrderedDict(headers)
-
-        # Store off the request data.
-        request_data = RequestData(headers_map, io.BytesIO())
-        self.stream_data[stream_id] = request_data
-
-    def stream_complete(self, stream_id: int):
-        """
-        When a stream is complete, we can send our response.
-        """
-        try:
-            request_data = self.stream_data[stream_id]
-        except KeyError:
-            # Just return, we probably 405'd this already
-            return
-
-        headers = request_data.headers
-        body = request_data.data.getvalue().decode("utf-8")
-
-        data = json.dumps({"headers": headers, "body": body}, indent=4).encode("utf8")
-
-        response_headers = (
-            (":status", "200"),
-            ("content-type", "application/json"),
-            ("content-length", str(len(data))),
-        )
-        self.conn.send_headers(stream_id, response_headers)
-        asyncio.ensure_future(self.send_data(data, stream_id))
-
-    def receive_data(self, data: bytes, stream_id: int):
-        """
-        We've received some data on a stream. If that stream is one we're
-        expecting data on, save it off. Otherwise, reset the stream.
-        """
-        try:
-            stream_data = self.stream_data[stream_id]
-        except KeyError:
-            self.conn.reset_stream(stream_id, error_code=ErrorCodes.PROTOCOL_ERROR)
-        else:
-            stream_data.data.write(data)
-
-    def stream_reset(self, stream_id):
-        """
-        A stream reset was sent. Stop sending data.
-        """
-        if stream_id in self.flow_control_futures:
-            future = self.flow_control_futures.pop(stream_id)
-            future.cancel()
-
-    async def send_data(self, data, stream_id):
-        """
-        Send data according to the flow control rules.
-        """
-        while data:
-            while self.conn.local_flow_control_window(stream_id) < 1:
-                try:
-                    await self.wait_for_flow_control(stream_id)
-                except asyncio.CancelledError:
-                    return
-
-            chunk_size = min(
-                self.conn.local_flow_control_window(stream_id),
-                len(data),
-                self.conn.max_outbound_frame_size,
-            )
-
-            try:
-                self.conn.send_data(
-                    stream_id, data[:chunk_size], end_stream=(chunk_size == len(data))
-                )
-            except (StreamClosedError, ProtocolError):
-                # The stream got closed and we didn't get told. We're done
-                # here.
-                break
-
-            assert self.transport is not None
-            self.transport.write(self.conn.data_to_send())
-            data = data[chunk_size:]
-
-    async def wait_for_flow_control(self, stream_id):
-        """
-        Waits for a Future that fires when the flow control window is opened.
-        """
-        f: asyncio.Future[None] = asyncio.Future()
-        self.flow_control_futures[stream_id] = f
-        await f
-
-    def window_updated(self, stream_id, delta):
-        """
-        A window update frame was received. Unblock some number of flow control
-        Futures.
-        """
-        if stream_id and stream_id in self.flow_control_futures:
-            f = self.flow_control_futures.pop(stream_id)
-            f.set_result(delta)
-        elif not stream_id:
-            for f in self.flow_control_futures.values():
-                f.set_result(delta)
-
-            self.flow_control_futures = {}
-
-
-@pytest_asyncio.fixture(scope="function")
-async def http2_echoserver() -> AsyncIterable[H2Server]:
-    loop = asyncio.get_event_loop()
-    serve = await loop.create_server(H2Protocol, "127.0.0.1", 0)
-    (host, port) = serve.sockets[0].getsockname()
-
-    asyncio.create_task(serve.wait_closed())
-
-    server = H2Server(host, port)
-    yield server
-
-    serve.close()

test_runner/fixtures/neon_fixtures.py

@@ -35,7 +35,6 @@ import toml
 from _pytest.config import Config
 from _pytest.config.argparsing import Parser
 from _pytest.fixtures import FixtureRequest
-from jwcrypto import jwk
 
 # Type-related stuff
 from psycopg2.extensions import connection as PgConnection
@@ -55,7 +54,6 @@ from fixtures.common_types import (
     TimelineId,
 )
 from fixtures.endpoint.http import EndpointHttpClient
-from fixtures.h2server import H2Server
 from fixtures.log_helper import log
 from fixtures.metrics import Metrics, MetricsGetter, parse_metrics
 from fixtures.neon_cli import NeonLocalCli, Pagectl
@@ -3095,41 +3093,12 @@ class PSQL:
         )
 
 
-def generate_proxy_tls_certs(common_name: str, key_path: Path, crt_path: Path):
-    if not key_path.exists():
-        r = subprocess.run(
-            [
-                "openssl",
-                "req",
-                "-new",
-                "-x509",
-                "-days",
-                "365",
-                "-nodes",
-                "-text",
-                "-out",
-                str(crt_path),
-                "-keyout",
-                str(key_path),
-                "-subj",
-                f"/CN={common_name}",
-                "-addext",
-                f"subjectAltName = DNS:{common_name}",
-            ]
-        )
-        assert r.returncode == 0
-
-
 class NeonProxy(PgProtocol):
     link_auth_uri: str = "http://dummy-uri"
 
     class AuthBackend(abc.ABC):
         """All auth backends must inherit from this class"""
 
-        @property
-        def default_conn_url(self) -> Optional[str]:
-            return None
-
         @abc.abstractmethod
         def extra_args(self) -> list[str]:
             pass
@@ -3143,7 +3112,7 @@ class NeonProxy(PgProtocol):
                 *["--allow-self-signed-compute", "true"],
             ]
 
-    class Console(AuthBackend):
+    class ControlPlane(AuthBackend):
         def __init__(self, endpoint: str, fixed_rate_limit: Optional[int] = None):
             self.endpoint = endpoint
             self.fixed_rate_limit = fixed_rate_limit
@@ -3167,21 +3136,6 @@ class NeonProxy(PgProtocol):
                 ]
             return args
 
-    @dataclass(frozen=True)
-    class Postgres(AuthBackend):
-        pg_conn_url: str
-
-        @property
-        def default_conn_url(self) -> Optional[str]:
-            return self.pg_conn_url
-
-        def extra_args(self) -> list[str]:
-            return [
-                # Postgres auth backend params
-                *["--auth-backend", "postgres"],
-                *["--auth-endpoint", self.pg_conn_url],
-            ]
-
     def __init__(
         self,
         neon_binpath: Path,
@@ -3196,7 +3150,7 @@ class NeonProxy(PgProtocol):
     ):
         host = "127.0.0.1"
         domain = "proxy.localtest.me"  # resolves to 127.0.0.1
-        super().__init__(dsn=auth_backend.default_conn_url, host=domain, port=proxy_port)
+        super().__init__(host=domain, port=proxy_port)
 
         self.domain = domain
         self.host = host
@@ -3218,7 +3172,29 @@ class NeonProxy(PgProtocol):
         # generate key of it doesn't exist
         crt_path = self.test_output_dir / "proxy.crt"
         key_path = self.test_output_dir / "proxy.key"
-        generate_proxy_tls_certs("*.localtest.me", key_path, crt_path)
+
+        if not key_path.exists():
+            r = subprocess.run(
+                [
+                    "openssl",
+                    "req",
+                    "-new",
+                    "-x509",
+                    "-days",
+                    "365",
+                    "-nodes",
+                    "-text",
+                    "-out",
+                    str(crt_path),
+                    "-keyout",
+                    str(key_path),
+                    "-subj",
+                    "/CN=*.localtest.me",
+                    "-addext",
+                    "subjectAltName = DNS:*.localtest.me",
+                ]
+            )
+            assert r.returncode == 0
 
         args = [
             str(self.neon_binpath / "proxy"),
@@ -3398,125 +3374,6 @@ class NeonProxy(PgProtocol):
         assert out == "ok"
 
 
-class NeonAuthBroker:
-    class ControlPlane:
-        def __init__(self, endpoint: str):
-            self.endpoint = endpoint
-
-        def extra_args(self) -> list[str]:
-            args = [
-                *["--auth-backend", "console"],
-                *["--auth-endpoint", self.endpoint],
-            ]
-            return args
-
-    def __init__(
-        self,
-        neon_binpath: Path,
-        test_output_dir: Path,
-        http_port: int,
-        mgmt_port: int,
-        external_http_port: int,
-        auth_backend: NeonAuthBroker.ControlPlane,
-    ):
-        self.domain = "apiauth.localtest.me"  # resolves to 127.0.0.1
-        self.host = "127.0.0.1"
-        self.http_port = http_port
-        self.external_http_port = external_http_port
-        self.neon_binpath = neon_binpath
-        self.test_output_dir = test_output_dir
-        self.mgmt_port = mgmt_port
-        self.auth_backend = auth_backend
-        self.http_timeout_seconds = 15
-        self._popen: Optional[subprocess.Popen[bytes]] = None
-
-    def start(self) -> NeonAuthBroker:
-        assert self._popen is None
-
-        # generate key of it doesn't exist
-        crt_path = self.test_output_dir / "proxy.crt"
-        key_path = self.test_output_dir / "proxy.key"
-        generate_proxy_tls_certs("apiauth.localtest.me", key_path, crt_path)
-
-        args = [
-            str(self.neon_binpath / "proxy"),
-            *["--http", f"{self.host}:{self.http_port}"],
-            *["--mgmt", f"{self.host}:{self.mgmt_port}"],
-            *["--wss", f"{self.host}:{self.external_http_port}"],
-            *["-c", str(crt_path)],
-            *["-k", str(key_path)],
-            *["--sql-over-http-pool-opt-in", "false"],
-            *["--is-auth-broker", "true"],
-            *self.auth_backend.extra_args(),
-        ]
-
-        logfile = open(self.test_output_dir / "proxy.log", "w")
-        self._popen = subprocess.Popen(args, stdout=logfile, stderr=logfile)
-        self._wait_until_ready()
-        return self
-
-    # Sends SIGTERM to the proxy if it has been started
-    def terminate(self):
-        if self._popen:
-            self._popen.terminate()
-
-    # Waits for proxy to exit if it has been opened with a default timeout of
-    # two seconds. Raises subprocess.TimeoutExpired if the proxy does not exit in time.
-    def wait_for_exit(self, timeout=2):
-        if self._popen:
-            self._popen.wait(timeout=timeout)
-
-    @backoff.on_exception(backoff.expo, requests.exceptions.RequestException, max_time=10)
-    def _wait_until_ready(self):
-        assert (
-            self._popen and self._popen.poll() is None
-        ), "Proxy exited unexpectedly. Check test log."
-        requests.get(f"http://{self.host}:{self.http_port}/v1/status")
-
-    async def query(self, query, args, **kwargs):
-        user = kwargs["user"]
-        token = kwargs["token"]
-        expected_code = kwargs.get("expected_code")
-
-        log.info(f"Executing http query: {query}")
-
-        connstr = f"postgresql://{user}@{self.domain}/postgres"
-        async with httpx.AsyncClient(verify=str(self.test_output_dir / "proxy.crt")) as client:
-            response = await client.post(
-                f"https://{self.domain}:{self.external_http_port}/sql",
-                json={"query": query, "params": args},
-                headers={
-                    "Neon-Connection-String": connstr,
-                    "Authorization": f"Bearer {token}",
-                },
-            )
-
-            if expected_code is not None:
-                assert response.status_code == expected_code, f"response: {response.json()}"
-            return response.json()
-
-    def get_metrics(self) -> str:
-        request_result = requests.get(f"http://{self.host}:{self.http_port}/metrics")
-        return request_result.text
-
-    def __enter__(self) -> NeonAuthBroker:
-        return self
-
-    def __exit__(
-        self,
-        _exc_type: Optional[type[BaseException]],
-        _exc_value: Optional[BaseException],
-        _traceback: Optional[TracebackType],
-    ):
-        if self._popen is not None:
-            self._popen.terminate()
-            try:
-                self._popen.wait(timeout=5)
-            except subprocess.TimeoutExpired:
-                log.warning("failed to gracefully terminate proxy; killing")
-                self._popen.kill()
-
-
 @pytest.fixture(scope="function")
 def link_proxy(
     port_distributor: PortDistributor, neon_binpath: Path, test_output_dir: Path
@@ -3547,20 +3404,39 @@ def static_proxy(
     port_distributor: PortDistributor,
     neon_binpath: Path,
     test_output_dir: Path,
+    httpserver: HTTPServer,
 ) -> Iterator[NeonProxy]:
-    """Neon proxy that routes directly to vanilla postgres."""
+    """Neon proxy that routes directly to vanilla postgres and a mocked cplane HTTP API."""
 
     port = vanilla_pg.default_options["port"]
     host = vanilla_pg.default_options["host"]
     dbname = vanilla_pg.default_options["dbname"]
-    auth_endpoint = f"postgres://proxy:password@{host}:{port}/{dbname}"
 
-    # For simplicity, we use the same user for both `--auth-endpoint` and `safe_psql`
     vanilla_pg.start()
     vanilla_pg.safe_psql("create user proxy with login superuser password 'password'")
-    vanilla_pg.safe_psql("CREATE SCHEMA IF NOT EXISTS neon_control_plane")
+    [(rolpassword,)] = vanilla_pg.safe_psql(
-    vanilla_pg.safe_psql(
+        "select rolpassword from pg_catalog.pg_authid where rolname = 'proxy'"
-        "CREATE TABLE neon_control_plane.endpoints (endpoint_id VARCHAR(255) PRIMARY KEY, allowed_ips VARCHAR(255))"
+    )
+
+    # return local postgres addr on ProxyWakeCompute.
+    httpserver.expect_request("/cplane/proxy_wake_compute").respond_with_json(
+        {
+            "address": f"{host}:{port}",
+            "aux": {
+                "endpoint_id": "ep-foo-bar-1234",
+                "branch_id": "br-foo-bar",
+                "project_id": "foo-bar",
+            },
+        }
+    )
+
+    # return local postgres addr on ProxyWakeCompute.
+    httpserver.expect_request("/cplane/proxy_get_role_secret").respond_with_json(
+        {
+            "role_secret": rolpassword,
+            "allowed_ips": None,
+            "project_id": "foo-bar",
+        }
     )
 
     proxy_port = port_distributor.get_port()
@@ -3575,92 +3451,12 @@ def static_proxy(
         http_port=http_port,
         mgmt_port=mgmt_port,
         external_http_port=external_http_port,
-        auth_backend=NeonProxy.Postgres(auth_endpoint),
+        auth_backend=NeonProxy.ControlPlane(httpserver.url_for("/cplane")),
     ) as proxy:
-        proxy.start()
+        proxy.default_options["user"] = "proxy"
-        yield proxy
+        proxy.default_options["password"] = "password"
+        proxy.default_options["dbname"] = dbname
 
-
-@pytest.fixture(scope="function")
-def neon_authorize_jwk() -> jwk.JWK:
-    kid = str(uuid.uuid4())
-    key = jwk.JWK.generate(kty="RSA", size=2048, alg="RS256", use="sig", kid=kid)
-    assert isinstance(key, jwk.JWK)
-    return key
-
-
-@pytest.fixture(scope="function")
-def cplane_endpoint_jwks(
-    httpserver: HTTPServer,
-    neon_authorize_jwk: jwk.JWK,
-    role_names: list[str],
-    audience: str | None,
-) -> jwk.JWK:
-    # return static fixture jwks.
-    jwk = neon_authorize_jwk.export_public(as_dict=True)
-    httpserver.expect_request("/authorize/jwks.json").respond_with_json({"keys": [jwk]})
-
-    id = str(uuid.uuid4())
-
-    # return jwks mock addr on GetEndpointJwks
-    httpserver.expect_request(re.compile("^/cplane/endpoints/.+/jwks$")).respond_with_json(
-        {
-            "jwks": [
-                {
-                    "id": id,
-                    "jwks_url": httpserver.url_for("/authorize/jwks.json"),
-                    "provider_name": "test",
-                    "jwt_audience": audience,
-                    "role_names": role_names,
-                }
-            ]
-        }
-    )
-
-    return neon_authorize_jwk
-
-
-@pytest.fixture(scope="function")
-def cplane_wake_compute_local_proxy(
-    httpserver: HTTPServer,
-    http2_echoserver: H2Server,
-) -> None:
-    local_proxy_addr = f"{http2_echoserver.host}:{http2_echoserver.port}"
-
-    # return local_proxy addr on ProxyWakeCompute.
-    httpserver.expect_request("/cplane/proxy_wake_compute").respond_with_json(
-        {
-            "address": local_proxy_addr,
-            "aux": {
-                "endpoint_id": "ep-foo-bar-1234",
-                "branch_id": "br-foo-bar",
-                "project_id": "foo-bar",
-            },
-        }
-    )
-
-
-@pytest.fixture(scope="function")
-def static_auth_broker(
-    port_distributor: PortDistributor,
-    neon_binpath: Path,
-    test_output_dir: Path,
-    httpserver: HTTPServer,
-) -> Iterable[NeonAuthBroker]:
-    """Neon Auth Broker that routes to a mocked local_proxy and a mocked cplane HTTP API."""
-
-    mgmt_port = port_distributor.get_port()
-    http_port = port_distributor.get_port()
-    external_http_port = port_distributor.get_port()
-
-    with NeonAuthBroker(
-        neon_binpath=neon_binpath,
-        test_output_dir=test_output_dir,
-        http_port=http_port,
-        mgmt_port=mgmt_port,
-        external_http_port=external_http_port,
-        auth_backend=NeonAuthBroker.ControlPlane(httpserver.url_for("/cplane")),
-    ) as proxy:
         proxy.start()
         yield proxy
 

test_runner/regress/test_auth_broker.py

@@ -1,114 +0,0 @@
-import json
-
-import pytest
-from fixtures.neon_fixtures import NeonAuthBroker
-from jwcrypto import jwk, jwt
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]])
-@pytest.mark.parametrize("audience", [None])
-async def test_auth_broker_happy(
-    static_auth_broker: NeonAuthBroker,
-    cplane_wake_compute_local_proxy: None,
-    cplane_endpoint_jwks: jwk.JWK,
-):
-    """
-    Signs a JWT and uses it to authorize a query to local_proxy.
-    """
-
-    token = jwt.JWT(
-        header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"}, claims={"sub": "user1"}
-    )
-    token.make_signed_token(cplane_endpoint_jwks)
-    res = await static_auth_broker.query(
-        "foo", ["arg1"], user="anonymous", token=token.serialize(), expected_code=200
-    )
-
-    # local proxy mock just echos back the request
-    # check that we forward the correct data
-
-    assert (
-        res["headers"]["authorization"] == f"Bearer {token.serialize()}"
-    ), "JWT should be forwarded"
-
-    assert (
-        "anonymous" in res["headers"]["neon-connection-string"]
-    ), "conn string should be forwarded"
-
-    assert json.loads(res["body"]) == {
-        "query": "foo",
-        "params": ["arg1"],
-    }, "Query body should be forwarded"
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]])
-@pytest.mark.parametrize("audience", [None])
-async def test_auth_broker_incorrect_role(
-    static_auth_broker: NeonAuthBroker,
-    cplane_wake_compute_local_proxy: None,
-    cplane_endpoint_jwks: jwk.JWK,
-):
-    """
-    Connects to auth broker with the wrong role associated with the JWKs
-    """
-
-    token = jwt.JWT(
-        header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"}, claims={"sub": "user1"}
-    )
-    token.make_signed_token(cplane_endpoint_jwks)
-    res = await static_auth_broker.query(
-        "foo", ["arg1"], user="wrong_role", token=token.serialize(), expected_code=400
-    )
-
-    # if the user is wrong, we announce that the jwk was not found.
-    assert "jwk not found" in res["message"]
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]])
-@pytest.mark.parametrize("audience", ["neon"])
-async def test_auth_broker_incorrect_aud(
-    static_auth_broker: NeonAuthBroker,
-    cplane_wake_compute_local_proxy: None,
-    cplane_endpoint_jwks: jwk.JWK,
-):
-    """
-    Connects to auth broker with the wrong audience associated with the JWKs
-    """
-
-    token = jwt.JWT(
-        header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"},
-        claims={"sub": "user1", "aud": "wrong_aud"},
-    )
-    token.make_signed_token(cplane_endpoint_jwks)
-    res = await static_auth_broker.query(
-        "foo", ["arg1"], user="anonymous", token=token.serialize(), expected_code=400
-    )
-
-    assert "invalid JWT token audience" in res["message"]
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("role_names", [["anonymous", "authenticated"]])
-@pytest.mark.parametrize("audience", ["neon"])
-async def test_auth_broker_missing_aud(
-    static_auth_broker: NeonAuthBroker,
-    cplane_wake_compute_local_proxy: None,
-    cplane_endpoint_jwks: jwk.JWK,
-):
-    """
-    Connects to auth broker with no audience
-    """
-
-    token = jwt.JWT(
-        header={"kid": cplane_endpoint_jwks.key_id, "alg": "RS256"},
-        claims={"sub": "user1"},
-    )
-    token.make_signed_token(cplane_endpoint_jwks)
-    res = await static_auth_broker.query(
-        "foo", ["arg1"], user="anonymous", token=token.serialize(), expected_code=400
-    )
-
-    assert "invalid JWT token audience" in res["message"]

test_runner/regress/test_proxy_allowed_ips.py

@@ -6,20 +6,27 @@ from fixtures.neon_fixtures import (
     NeonProxy,
     VanillaPostgres,
 )
+from pytest_httpserver import HTTPServer
 
 TABLE_NAME = "neon_control_plane.endpoints"
 
 
-# Proxy uses the same logic for psql and websockets.
+def test_proxy_psql_not_allowed_ips(
-@pytest.mark.asyncio
+    static_proxy: NeonProxy,
-async def test_proxy_psql_allowed_ips(static_proxy: NeonProxy, vanilla_pg: VanillaPostgres):
+    vanilla_pg: VanillaPostgres,
-    # Shouldn't be able to connect to this project
+    httpserver: HTTPServer,
-    vanilla_pg.safe_psql(
+):
-        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('private-project', '8.8.8.8')"
+    [(rolpassword,)] = vanilla_pg.safe_psql(
+        "select rolpassword from pg_catalog.pg_authid where rolname = 'proxy'"
     )
-    # Should be able to connect to this project
+
-    vanilla_pg.safe_psql(
+    # Shouldn't be able to connect to this project
-        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('generic-project', '::1,127.0.0.1')"
+    httpserver.expect_request("/cplane/proxy_get_role_secret").respond_with_json(
+        {
+            "role_secret": rolpassword,
+            "allowed_ips": ["8.8.8.8"],
+            "project_id": "foo-bar",
+        }
     )
 
     def check_cannot_connect(**kwargs):
@@ -37,6 +44,25 @@ async def test_proxy_psql_allowed_ips(static_proxy: NeonProxy, vanilla_pg: Vanil
     # with SNI
     check_cannot_connect(query="select 1", host="private-project.localtest.me")
 
+
+def test_proxy_psql_allowed_ips(
+    static_proxy: NeonProxy,
+    vanilla_pg: VanillaPostgres,
+    httpserver: HTTPServer,
+):
+    [(rolpassword,)] = vanilla_pg.safe_psql(
+        "select rolpassword from pg_catalog.pg_authid where rolname = 'proxy'"
+    )
+
+    # Should be able to connect to this project
+    httpserver.expect_request("/cplane/proxy_get_role_secret").respond_with_json(
+        {
+            "role_secret": rolpassword,
+            "allowed_ips": ["::1", "127.0.0.1"],
+            "project_id": "foo-bar",
+        }
+    )
+
     # no SNI, deprecated `options=project` syntax (before we had several endpoint in project)
     out = static_proxy.safe_psql(query="select 1", sslsni=0, options="project=generic-project")
     assert out[0][0] == 1
@@ -50,27 +76,61 @@ async def test_proxy_psql_allowed_ips(static_proxy: NeonProxy, vanilla_pg: Vanil
     assert out[0][0] == 1
 
 
-@pytest.mark.asyncio
+def test_proxy_http_not_allowed_ips(
-async def test_proxy_http_allowed_ips(static_proxy: NeonProxy, vanilla_pg: VanillaPostgres):
+    static_proxy: NeonProxy,
-    static_proxy.safe_psql("create user http_auth with password 'http' superuser")
+    vanilla_pg: VanillaPostgres,
+    httpserver: HTTPServer,
+):
+    vanilla_pg.safe_psql("create user http_auth with password 'http' superuser")
 
-    # Shouldn't be able to connect to this project
+    [(rolpassword,)] = vanilla_pg.safe_psql(
-    vanilla_pg.safe_psql(
+        "select rolpassword from pg_catalog.pg_authid where rolname = 'http_auth'"
-        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('proxy', '8.8.8.8')"
     )
 
-    def query(status: int, query: str, *args):
+    httpserver.expect_oneshot_request("/cplane/proxy_get_role_secret").respond_with_json(
+        {
+            "role_secret": rolpassword,
+            "allowed_ips": ["8.8.8.8"],
+            "project_id": "foo-bar",
+        }
+    )
+
+    with httpserver.wait() as waiting:
         static_proxy.http_query(
-            query,
+            "select 1;",
-            args,
+            [],
             user="http_auth",
             password="http",
-            expected_code=status,
+            expected_code=400,
         )
+    assert waiting.result
 
-    query(400, "select 1;")  # ip address is not allowed
+
-    # Should be able to connect to this project
+def test_proxy_http_allowed_ips(
-    vanilla_pg.safe_psql(
+    static_proxy: NeonProxy,
-        f"UPDATE {TABLE_NAME} SET allowed_ips = '8.8.8.8,127.0.0.1' WHERE endpoint_id = 'proxy'"
+    vanilla_pg: VanillaPostgres,
+    httpserver: HTTPServer,
+):
+    vanilla_pg.safe_psql("create user http_auth with password 'http' superuser")
+
+    [(rolpassword,)] = vanilla_pg.safe_psql(
+        "select rolpassword from pg_catalog.pg_authid where rolname = 'http_auth'"
     )
-    query(200, "select 1;")  # should work now
+
+    httpserver.expect_oneshot_request("/cplane/proxy_get_role_secret").respond_with_json(
+        {
+            "role_secret": rolpassword,
+            "allowed_ips": ["8.8.8.8", "127.0.0.1"],
+            "project_id": "foo-bar",
+        }
+    )
+
+    with httpserver.wait() as waiting:
+        static_proxy.http_query(
+            "select 1;",
+            [],
+            user="http_auth",
+            password="http",
+            expected_code=200,
+        )
+    assert waiting.result

test_runner/stubs/h2/README.md

@@ -1 +0,0 @@
-generated via `poetry run stubgen -p h2 -o test_runner/stubs`

test_runner/stubs/h2/config.pyi

@@ -1,42 +0,0 @@
-from _typeshed import Incomplete
-
-class _BooleanConfigOption:
-    name: Incomplete
-    attr_name: Incomplete
-    def __init__(self, name) -> None: ...
-    def __get__(self, instance, owner): ...
-    def __set__(self, instance, value) -> None: ...
-
-class DummyLogger:
-    def __init__(self, *vargs) -> None: ...
-    def debug(self, *vargs, **kwargs) -> None: ...
-    def trace(self, *vargs, **kwargs) -> None: ...
-
-class OutputLogger:
-    file: Incomplete
-    trace_level: Incomplete
-    def __init__(self, file: Incomplete | None = ..., trace_level: bool = ...) -> None: ...
-    def debug(self, fmtstr, *args) -> None: ...
-    def trace(self, fmtstr, *args) -> None: ...
-
-class H2Configuration:
-    client_side: Incomplete
-    validate_outbound_headers: Incomplete
-    normalize_outbound_headers: Incomplete
-    validate_inbound_headers: Incomplete
-    normalize_inbound_headers: Incomplete
-    logger: Incomplete
-    def __init__(
-        self,
-        client_side: bool = ...,
-        header_encoding: Incomplete | None = ...,
-        validate_outbound_headers: bool = ...,
-        normalize_outbound_headers: bool = ...,
-        validate_inbound_headers: bool = ...,
-        normalize_inbound_headers: bool = ...,
-        logger: Incomplete | None = ...,
-    ) -> None: ...
-    @property
-    def header_encoding(self): ...
-    @header_encoding.setter
-    def header_encoding(self, value) -> None: ...

test_runner/stubs/h2/connection.pyi

@@ -1,142 +0,0 @@
-from enum import Enum, IntEnum
-
-from _typeshed import Incomplete
-
-from .config import H2Configuration as H2Configuration
-from .errors import ErrorCodes as ErrorCodes
-from .events import AlternativeServiceAvailable as AlternativeServiceAvailable
-from .events import ConnectionTerminated as ConnectionTerminated
-from .events import PingAckReceived as PingAckReceived
-from .events import PingReceived as PingReceived
-from .events import PriorityUpdated as PriorityUpdated
-from .events import RemoteSettingsChanged as RemoteSettingsChanged
-from .events import SettingsAcknowledged as SettingsAcknowledged
-from .events import UnknownFrameReceived as UnknownFrameReceived
-from .events import WindowUpdated as WindowUpdated
-from .exceptions import DenialOfServiceError as DenialOfServiceError
-from .exceptions import FlowControlError as FlowControlError
-from .exceptions import FrameTooLargeError as FrameTooLargeError
-from .exceptions import NoAvailableStreamIDError as NoAvailableStreamIDError
-from .exceptions import NoSuchStreamError as NoSuchStreamError
-from .exceptions import ProtocolError as ProtocolError
-from .exceptions import RFC1122Error as RFC1122Error
-from .exceptions import StreamClosedError as StreamClosedError
-from .exceptions import StreamIDTooLowError as StreamIDTooLowError
-from .exceptions import TooManyStreamsError as TooManyStreamsError
-from .frame_buffer import FrameBuffer as FrameBuffer
-from .settings import SettingCodes as SettingCodes
-from .settings import Settings as Settings
-from .stream import H2Stream as H2Stream
-from .stream import StreamClosedBy as StreamClosedBy
-from .utilities import guard_increment_window as guard_increment_window
-from .windows import WindowManager as WindowManager
-
-class ConnectionState(Enum):
-    IDLE: int
-    CLIENT_OPEN: int
-    SERVER_OPEN: int
-    CLOSED: int
-
-class ConnectionInputs(Enum):
-    SEND_HEADERS: int
-    SEND_PUSH_PROMISE: int
-    SEND_DATA: int
-    SEND_GOAWAY: int
-    SEND_WINDOW_UPDATE: int
-    SEND_PING: int
-    SEND_SETTINGS: int
-    SEND_RST_STREAM: int
-    SEND_PRIORITY: int
-    RECV_HEADERS: int
-    RECV_PUSH_PROMISE: int
-    RECV_DATA: int
-    RECV_GOAWAY: int
-    RECV_WINDOW_UPDATE: int
-    RECV_PING: int
-    RECV_SETTINGS: int
-    RECV_RST_STREAM: int
-    RECV_PRIORITY: int
-    SEND_ALTERNATIVE_SERVICE: int
-    RECV_ALTERNATIVE_SERVICE: int
-
-class AllowedStreamIDs(IntEnum):
-    EVEN: int
-    ODD: int
-
-class H2ConnectionStateMachine:
-    state: Incomplete
-    def __init__(self) -> None: ...
-    def process_input(self, input_): ...
-
-class H2Connection:
-    DEFAULT_MAX_OUTBOUND_FRAME_SIZE: int
-    DEFAULT_MAX_INBOUND_FRAME_SIZE: Incomplete
-    HIGHEST_ALLOWED_STREAM_ID: Incomplete
-    MAX_WINDOW_INCREMENT: Incomplete
-    DEFAULT_MAX_HEADER_LIST_SIZE: Incomplete
-    MAX_CLOSED_STREAMS: Incomplete
-    state_machine: Incomplete
-    streams: Incomplete
-    highest_inbound_stream_id: int
-    highest_outbound_stream_id: int
-    encoder: Incomplete
-    decoder: Incomplete
-    config: Incomplete
-    local_settings: Incomplete
-    remote_settings: Incomplete
-    outbound_flow_control_window: Incomplete
-    max_outbound_frame_size: Incomplete
-    max_inbound_frame_size: Incomplete
-    incoming_buffer: Incomplete
-    def __init__(self, config: Incomplete | None = ...) -> None: ...
-    @property
-    def open_outbound_streams(self): ...
-    @property
-    def open_inbound_streams(self): ...
-    @property
-    def inbound_flow_control_window(self): ...
-    def initiate_connection(self) -> None: ...
-    def initiate_upgrade_connection(self, settings_header: Incomplete | None = ...): ...
-    def get_next_available_stream_id(self): ...
-    def send_headers(
-        self,
-        stream_id,
-        headers,
-        end_stream: bool = ...,
-        priority_weight: Incomplete | None = ...,
-        priority_depends_on: Incomplete | None = ...,
-        priority_exclusive: Incomplete | None = ...,
-    ) -> None: ...
-    def send_data(
-        self, stream_id, data, end_stream: bool = ..., pad_length: Incomplete | None = ...
-    ) -> None: ...
-    def end_stream(self, stream_id) -> None: ...
-    def increment_flow_control_window(
-        self, increment, stream_id: Incomplete | None = ...
-    ) -> None: ...
-    def push_stream(self, stream_id, promised_stream_id, request_headers) -> None: ...
-    def ping(self, opaque_data) -> None: ...
-    def reset_stream(self, stream_id, error_code: int = ...) -> None: ...
-    def close_connection(
-        self,
-        error_code: int = ...,
-        additional_data: Incomplete | None = ...,
-        last_stream_id: Incomplete | None = ...,
-    ) -> None: ...
-    def update_settings(self, new_settings) -> None: ...
-    def advertise_alternative_service(
-        self, field_value, origin: Incomplete | None = ..., stream_id: Incomplete | None = ...
-    ) -> None: ...
-    def prioritize(
-        self,
-        stream_id,
-        weight: Incomplete | None = ...,
-        depends_on: Incomplete | None = ...,
-        exclusive: Incomplete | None = ...,
-    ) -> None: ...
-    def local_flow_control_window(self, stream_id): ...
-    def remote_flow_control_window(self, stream_id): ...
-    def acknowledge_received_data(self, acknowledged_size, stream_id) -> None: ...
-    def data_to_send(self, amount: Incomplete | None = ...): ...
-    def clear_outbound_data_buffer(self) -> None: ...
-    def receive_data(self, data): ...

test_runner/stubs/h2/errors.pyi

@@ -1,17 +0,0 @@
-import enum
-
-class ErrorCodes(enum.IntEnum):
-    NO_ERROR: int
-    PROTOCOL_ERROR: int
-    INTERNAL_ERROR: int
-    FLOW_CONTROL_ERROR: int
-    SETTINGS_TIMEOUT: int
-    STREAM_CLOSED: int
-    FRAME_SIZE_ERROR: int
-    REFUSED_STREAM: int
-    CANCEL: int
-    COMPRESSION_ERROR: int
-    CONNECT_ERROR: int
-    ENHANCE_YOUR_CALM: int
-    INADEQUATE_SECURITY: int
-    HTTP_1_1_REQUIRED: int

test_runner/stubs/h2/events.pyi

@@ -1,106 +0,0 @@
-from _typeshed import Incomplete
-
-from .settings import ChangedSetting as ChangedSetting
-
-class Event: ...
-
-class RequestReceived(Event):
-    stream_id: Incomplete
-    headers: Incomplete
-    stream_ended: Incomplete
-    priority_updated: Incomplete
-    def __init__(self) -> None: ...
-
-class ResponseReceived(Event):
-    stream_id: Incomplete
-    headers: Incomplete
-    stream_ended: Incomplete
-    priority_updated: Incomplete
-    def __init__(self) -> None: ...
-
-class TrailersReceived(Event):
-    stream_id: Incomplete
-    headers: Incomplete
-    stream_ended: Incomplete
-    priority_updated: Incomplete
-    def __init__(self) -> None: ...
-
-class _HeadersSent(Event): ...
-class _ResponseSent(_HeadersSent): ...
-class _RequestSent(_HeadersSent): ...
-class _TrailersSent(_HeadersSent): ...
-class _PushedRequestSent(_HeadersSent): ...
-
-class InformationalResponseReceived(Event):
-    stream_id: Incomplete
-    headers: Incomplete
-    priority_updated: Incomplete
-    def __init__(self) -> None: ...
-
-class DataReceived(Event):
-    stream_id: Incomplete
-    data: Incomplete
-    flow_controlled_length: Incomplete
-    stream_ended: Incomplete
-    def __init__(self) -> None: ...
-
-class WindowUpdated(Event):
-    stream_id: Incomplete
-    delta: Incomplete
-    def __init__(self) -> None: ...
-
-class RemoteSettingsChanged(Event):
-    changed_settings: Incomplete
-    def __init__(self) -> None: ...
-    @classmethod
-    def from_settings(cls, old_settings, new_settings): ...
-
-class PingReceived(Event):
-    ping_data: Incomplete
-    def __init__(self) -> None: ...
-
-class PingAckReceived(Event):
-    ping_data: Incomplete
-    def __init__(self) -> None: ...
-
-class StreamEnded(Event):
-    stream_id: Incomplete
-    def __init__(self) -> None: ...
-
-class StreamReset(Event):
-    stream_id: Incomplete
-    error_code: Incomplete
-    remote_reset: bool
-    def __init__(self) -> None: ...
-
-class PushedStreamReceived(Event):
-    pushed_stream_id: Incomplete
-    parent_stream_id: Incomplete
-    headers: Incomplete
-    def __init__(self) -> None: ...
-
-class SettingsAcknowledged(Event):
-    changed_settings: Incomplete
-    def __init__(self) -> None: ...
-
-class PriorityUpdated(Event):
-    stream_id: Incomplete
-    weight: Incomplete
-    depends_on: Incomplete
-    exclusive: Incomplete
-    def __init__(self) -> None: ...
-
-class ConnectionTerminated(Event):
-    error_code: Incomplete
-    last_stream_id: Incomplete
-    additional_data: Incomplete
-    def __init__(self) -> None: ...
-
-class AlternativeServiceAvailable(Event):
-    origin: Incomplete
-    field_value: Incomplete
-    def __init__(self) -> None: ...
-
-class UnknownFrameReceived(Event):
-    frame: Incomplete
-    def __init__(self) -> None: ...

test_runner/stubs/h2/exceptions.pyi

@@ -1,48 +0,0 @@
-from _typeshed import Incomplete
-
-class H2Error(Exception): ...
-
-class ProtocolError(H2Error):
-    error_code: Incomplete
-
-class FrameTooLargeError(ProtocolError):
-    error_code: Incomplete
-
-class FrameDataMissingError(ProtocolError):
-    error_code: Incomplete
-
-class TooManyStreamsError(ProtocolError): ...
-
-class FlowControlError(ProtocolError):
-    error_code: Incomplete
-
-class StreamIDTooLowError(ProtocolError):
-    stream_id: Incomplete
-    max_stream_id: Incomplete
-    def __init__(self, stream_id, max_stream_id) -> None: ...
-
-class NoAvailableStreamIDError(ProtocolError): ...
-
-class NoSuchStreamError(ProtocolError):
-    stream_id: Incomplete
-    def __init__(self, stream_id) -> None: ...
-
-class StreamClosedError(NoSuchStreamError):
-    stream_id: Incomplete
-    error_code: Incomplete
-    def __init__(self, stream_id) -> None: ...
-
-class InvalidSettingsValueError(ProtocolError, ValueError):
-    error_code: Incomplete
-    def __init__(self, msg, error_code) -> None: ...
-
-class InvalidBodyLengthError(ProtocolError):
-    expected_length: Incomplete
-    actual_length: Incomplete
-    def __init__(self, expected, actual) -> None: ...
-
-class UnsupportedFrameError(ProtocolError): ...
-class RFC1122Error(H2Error): ...
-
-class DenialOfServiceError(ProtocolError):
-    error_code: Incomplete

test_runner/stubs/h2/frame_buffer.pyi

@@ -1,19 +0,0 @@
-from .exceptions import (
-    FrameDataMissingError as FrameDataMissingError,
-)
-from .exceptions import (
-    FrameTooLargeError as FrameTooLargeError,
-)
-from .exceptions import (
-    ProtocolError as ProtocolError,
-)
-
-CONTINUATION_BACKLOG: int
-
-class FrameBuffer:
-    data: bytes
-    max_frame_size: int
-    def __init__(self, server: bool = ...) -> None: ...
-    def add_data(self, data) -> None: ...
-    def __iter__(self): ...
-    def __next__(self): ...

test_runner/stubs/h2/settings.pyi

@@ -1,61 +0,0 @@
-import enum
-from collections.abc import MutableMapping
-from typing import Any
-
-from _typeshed import Incomplete
-from h2.errors import ErrorCodes as ErrorCodes
-from h2.exceptions import InvalidSettingsValueError as InvalidSettingsValueError
-
-class SettingCodes(enum.IntEnum):
-    HEADER_TABLE_SIZE: Incomplete
-    ENABLE_PUSH: Incomplete
-    MAX_CONCURRENT_STREAMS: Incomplete
-    INITIAL_WINDOW_SIZE: Incomplete
-    MAX_FRAME_SIZE: Incomplete
-    MAX_HEADER_LIST_SIZE: Incomplete
-    ENABLE_CONNECT_PROTOCOL: Incomplete
-
-class ChangedSetting:
-    setting: Incomplete
-    original_value: Incomplete
-    new_value: Incomplete
-    def __init__(self, setting, original_value, new_value) -> None: ...
-
-class Settings(MutableMapping[str, Any]):
-    def __init__(self, client: bool = ..., initial_values: Incomplete | None = ...) -> None: ...
-    def acknowledge(self): ...
-    @property
-    def header_table_size(self): ...
-    @header_table_size.setter
-    def header_table_size(self, value) -> None: ...
-    @property
-    def enable_push(self): ...
-    @enable_push.setter
-    def enable_push(self, value) -> None: ...
-    @property
-    def initial_window_size(self): ...
-    @initial_window_size.setter
-    def initial_window_size(self, value) -> None: ...
-    @property
-    def max_frame_size(self): ...
-    @max_frame_size.setter
-    def max_frame_size(self, value) -> None: ...
-    @property
-    def max_concurrent_streams(self): ...
-    @max_concurrent_streams.setter
-    def max_concurrent_streams(self, value) -> None: ...
-    @property
-    def max_header_list_size(self): ...
-    @max_header_list_size.setter
-    def max_header_list_size(self, value) -> None: ...
-    @property
-    def enable_connect_protocol(self): ...
-    @enable_connect_protocol.setter
-    def enable_connect_protocol(self, value) -> None: ...
-    def __getitem__(self, key): ...
-    def __setitem__(self, key, value) -> None: ...
-    def __delitem__(self, key) -> None: ...
-    def __iter__(self): ...
-    def __len__(self) -> int: ...
-    def __eq__(self, other): ...
-    def __ne__(self, other): ...

test_runner/stubs/h2/stream.pyi

@@ -1,184 +0,0 @@
-from enum import Enum, IntEnum
-
-from _typeshed import Incomplete
-
-from .errors import ErrorCodes as ErrorCodes
-from .events import (
-    AlternativeServiceAvailable as AlternativeServiceAvailable,
-)
-from .events import (
-    DataReceived as DataReceived,
-)
-from .events import (
-    InformationalResponseReceived as InformationalResponseReceived,
-)
-from .events import (
-    PushedStreamReceived as PushedStreamReceived,
-)
-from .events import (
-    RequestReceived as RequestReceived,
-)
-from .events import (
-    ResponseReceived as ResponseReceived,
-)
-from .events import (
-    StreamEnded as StreamEnded,
-)
-from .events import (
-    StreamReset as StreamReset,
-)
-from .events import (
-    TrailersReceived as TrailersReceived,
-)
-from .events import (
-    WindowUpdated as WindowUpdated,
-)
-from .exceptions import (
-    FlowControlError as FlowControlError,
-)
-from .exceptions import (
-    InvalidBodyLengthError as InvalidBodyLengthError,
-)
-from .exceptions import (
-    ProtocolError as ProtocolError,
-)
-from .exceptions import (
-    StreamClosedError as StreamClosedError,
-)
-from .utilities import (
-    HeaderValidationFlags as HeaderValidationFlags,
-)
-from .utilities import (
-    authority_from_headers as authority_from_headers,
-)
-from .utilities import (
-    extract_method_header as extract_method_header,
-)
-from .utilities import (
-    guard_increment_window as guard_increment_window,
-)
-from .utilities import (
-    is_informational_response as is_informational_response,
-)
-from .utilities import (
-    normalize_inbound_headers as normalize_inbound_headers,
-)
-from .utilities import (
-    normalize_outbound_headers as normalize_outbound_headers,
-)
-from .utilities import (
-    validate_headers as validate_headers,
-)
-from .utilities import (
-    validate_outbound_headers as validate_outbound_headers,
-)
-from .windows import WindowManager as WindowManager
-
-class StreamState(IntEnum):
-    IDLE: int
-    RESERVED_REMOTE: int
-    RESERVED_LOCAL: int
-    OPEN: int
-    HALF_CLOSED_REMOTE: int
-    HALF_CLOSED_LOCAL: int
-    CLOSED: int
-
-class StreamInputs(Enum):
-    SEND_HEADERS: int
-    SEND_PUSH_PROMISE: int
-    SEND_RST_STREAM: int
-    SEND_DATA: int
-    SEND_WINDOW_UPDATE: int
-    SEND_END_STREAM: int
-    RECV_HEADERS: int
-    RECV_PUSH_PROMISE: int
-    RECV_RST_STREAM: int
-    RECV_DATA: int
-    RECV_WINDOW_UPDATE: int
-    RECV_END_STREAM: int
-    RECV_CONTINUATION: int
-    SEND_INFORMATIONAL_HEADERS: int
-    RECV_INFORMATIONAL_HEADERS: int
-    SEND_ALTERNATIVE_SERVICE: int
-    RECV_ALTERNATIVE_SERVICE: int
-    UPGRADE_CLIENT: int
-    UPGRADE_SERVER: int
-
-class StreamClosedBy(Enum):
-    SEND_END_STREAM: int
-    RECV_END_STREAM: int
-    SEND_RST_STREAM: int
-    RECV_RST_STREAM: int
-
-STREAM_OPEN: Incomplete
-
-class H2StreamStateMachine:
-    state: Incomplete
-    stream_id: Incomplete
-    client: Incomplete
-    headers_sent: Incomplete
-    trailers_sent: Incomplete
-    headers_received: Incomplete
-    trailers_received: Incomplete
-    stream_closed_by: Incomplete
-    def __init__(self, stream_id) -> None: ...
-    def process_input(self, input_): ...
-    def request_sent(self, previous_state): ...
-    def response_sent(self, previous_state): ...
-    def request_received(self, previous_state): ...
-    def response_received(self, previous_state): ...
-    def data_received(self, previous_state): ...
-    def window_updated(self, previous_state): ...
-    def stream_half_closed(self, previous_state): ...
-    def stream_ended(self, previous_state): ...
-    def stream_reset(self, previous_state): ...
-    def send_new_pushed_stream(self, previous_state): ...
-    def recv_new_pushed_stream(self, previous_state): ...
-    def send_push_promise(self, previous_state): ...
-    def recv_push_promise(self, previous_state): ...
-    def send_end_stream(self, previous_state) -> None: ...
-    def send_reset_stream(self, previous_state) -> None: ...
-    def reset_stream_on_error(self, previous_state) -> None: ...
-    def recv_on_closed_stream(self, previous_state) -> None: ...
-    def send_on_closed_stream(self, previous_state) -> None: ...
-    def recv_push_on_closed_stream(self, previous_state) -> None: ...
-    def send_push_on_closed_stream(self, previous_state) -> None: ...
-    def send_informational_response(self, previous_state): ...
-    def recv_informational_response(self, previous_state): ...
-    def recv_alt_svc(self, previous_state): ...
-    def send_alt_svc(self, previous_state) -> None: ...
-
-class H2Stream:
-    state_machine: Incomplete
-    stream_id: Incomplete
-    max_outbound_frame_size: Incomplete
-    request_method: Incomplete
-    outbound_flow_control_window: Incomplete
-    config: Incomplete
-    def __init__(self, stream_id, config, inbound_window_size, outbound_window_size) -> None: ...
-    @property
-    def inbound_flow_control_window(self): ...
-    @property
-    def open(self): ...
-    @property
-    def closed(self): ...
-    @property
-    def closed_by(self): ...
-    def upgrade(self, client_side) -> None: ...
-    def send_headers(self, headers, encoder, end_stream: bool = ...): ...
-    def push_stream_in_band(self, related_stream_id, headers, encoder): ...
-    def locally_pushed(self): ...
-    def send_data(self, data, end_stream: bool = ..., pad_length: Incomplete | None = ...): ...
-    def end_stream(self): ...
-    def advertise_alternative_service(self, field_value): ...
-    def increase_flow_control_window(self, increment): ...
-    def receive_push_promise_in_band(self, promised_stream_id, headers, header_encoding): ...
-    def remotely_pushed(self, pushed_headers): ...
-    def receive_headers(self, headers, end_stream, header_encoding): ...
-    def receive_data(self, data, end_stream, flow_control_len): ...
-    def receive_window_update(self, increment): ...
-    def receive_continuation(self) -> None: ...
-    def receive_alt_svc(self, frame): ...
-    def reset_stream(self, error_code: int = ...): ...
-    def stream_reset(self, frame): ...
-    def acknowledge_received_data(self, acknowledged_size): ...

test_runner/stubs/h2/utilities.pyi

@@ -1,25 +0,0 @@
-from typing import NamedTuple
-
-from _typeshed import Incomplete
-
-from .exceptions import FlowControlError as FlowControlError
-from .exceptions import ProtocolError as ProtocolError
-
-UPPER_RE: Incomplete
-CONNECTION_HEADERS: Incomplete
-
-def extract_method_header(headers): ...
-def is_informational_response(headers): ...
-def guard_increment_window(current, increment): ...
-def authority_from_headers(headers): ...
-
-class HeaderValidationFlags(NamedTuple):
-    is_client: Incomplete
-    is_trailer: Incomplete
-    is_response_header: Incomplete
-    is_push_promise: Incomplete
-
-def validate_headers(headers, hdr_validation_flags): ...
-def normalize_outbound_headers(headers, hdr_validation_flags): ...
-def normalize_inbound_headers(headers, hdr_validation_flags): ...
-def validate_outbound_headers(headers, hdr_validation_flags): ...

test_runner/stubs/h2/windows.pyi

@@ -1,13 +0,0 @@
-from _typeshed import Incomplete
-
-from .exceptions import FlowControlError as FlowControlError
-
-LARGEST_FLOW_CONTROL_WINDOW: Incomplete
-
-class WindowManager:
-    max_window_size: Incomplete
-    current_window_size: Incomplete
-    def __init__(self, max_window_size) -> None: ...
-    def window_consumed(self, size) -> None: ...
-    def window_opened(self, size) -> None: ...
-    def process_bytes(self, size): ...