From e109d5aac02186828b1e0f126f2a05f815ecf51a Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Fri, 4 Oct 2024 09:16:21 +0100
Subject: [PATCH] add local_proxy to computespec

---
 Cargo.lock                       |  1 +
 compute_tools/Cargo.toml         |  1 +
 compute_tools/src/lib.rs         |  1 +
 compute_tools/src/local_proxy.rs | 56 ++++++++++++++++++++++++++++++++
 control_plane/src/endpoint.rs    |  1 +
 libs/compute_api/src/spec.rs     |  4 +++
 6 files changed, 64 insertions(+)
 create mode 100644 compute_tools/src/local_proxy.rs

diff --git a/Cargo.lock b/Cargo.lock
index 419276f543..6108549660 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1265,6 +1265,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "bytes",
+ "camino",
  "cfg-if",
  "chrono",
  "clap",
diff --git a/compute_tools/Cargo.toml b/compute_tools/Cargo.toml
index 6bb3e211b6..91e0b9d5b8 100644
--- a/compute_tools/Cargo.toml
+++ b/compute_tools/Cargo.toml
@@ -11,6 +11,7 @@ testing = []
 
 [dependencies]
 anyhow.workspace = true
+camino.workspace = true
 chrono.workspace = true
 cfg-if.workspace = true
 clap.workspace = true
diff --git a/compute_tools/src/lib.rs b/compute_tools/src/lib.rs
index 0795eb6171..477f423aa2 100644
--- a/compute_tools/src/lib.rs
+++ b/compute_tools/src/lib.rs
@@ -15,6 +15,7 @@ pub mod catalog;
 pub mod compute;
 pub mod disk_quota;
 pub mod extension_server;
+pub mod local_proxy;
 pub mod lsn_lease;
 mod migration;
 pub mod monitor;
diff --git a/compute_tools/src/local_proxy.rs b/compute_tools/src/local_proxy.rs
new file mode 100644
index 0000000000..3de3c58786
--- /dev/null
+++ b/compute_tools/src/local_proxy.rs
@@ -0,0 +1,56 @@
+//! Local Proxy is a feature of our BaaS Neon Authorize project.
+//!
+//! Local Proxy validates JWTs and manages the pg_session_jwt extension.
+//! It also maintains a connection pool to postgres.
+
+use anyhow::{Context, Result};
+use camino::Utf8Path;
+use compute_api::spec::LocalProxySpec;
+use nix::sys::signal::Signal;
+use utils::pid_file::{self, PidFileRead};
+
+pub fn configure(local_proxy: &LocalProxySpec) -> Result<()> {
+    write_local_proxy_conf("/etc/local_proxy/config.json".as_ref(), local_proxy)?;
+    notify_local_proxy("/etc/local_proxy/pid".as_ref())?;
+
+    Ok(())
+}
+
+/// Create or completely rewrite configuration file specified by `path`
+fn write_local_proxy_conf(path: &Utf8Path, local_proxy: &LocalProxySpec) -> Result<()> {
+    let config =
+        serde_json::to_string_pretty(local_proxy).context("serializing LocalProxySpec to json")?;
+    std::fs::write(path, config).with_context(|| format!("writing {path}"))?;
+
+    Ok(())
+}
+
+/// Notify local proxy about a new config file.
+fn notify_local_proxy(path: &Utf8Path) -> Result<()> {
+    match pid_file::read(path)? {
+        // if the file doesn't exist, or isn't locked, local_proxy isn't running
+        // and will naturally pick up our config later
+        PidFileRead::NotExist | PidFileRead::NotHeldByAnyProcess(_) => {}
+        PidFileRead::LockedByOtherProcess(pid) => {
+            // From the pid_file docs:
+            //
+            // > 1. The other process might exit at any time, turning the given PID stale.
+            // > 2. There is a small window in which `claim_for_current_process` has already
+            // >    locked the file but not yet updates its contents. [`read`] will return
+            // >    this variant here, but with the old file contents, i.e., a stale PID.
+            // >
+            // > The kernel is free to recycle PID once it has been `wait(2)`ed upon by
+            // > its creator. Thus, acting upon a stale PID, e.g., by issuing a `kill`
+            // > system call on it, bears the risk of killing an unrelated process.
+            // > This is an inherent limitation of using pidfiles.
+            // > The only race-free solution is to have a supervisor-process with a lifetime
+            // > that exceeds that of all of its child-processes (e.g., `runit`, `supervisord`).
+            //
+            // This is an ok risk as we only send a SIGHUP which likely won't actually
+            // kill the process, only reload config.
+            nix::sys::signal::kill(pid, Signal::SIGHUP).context("sending signal to local_proxy")?;
+        }
+    }
+
+    Ok(())
+}
diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs
index 18f396b886..7cdf621737 100644
--- a/control_plane/src/endpoint.rs
+++ b/control_plane/src/endpoint.rs
@@ -599,6 +599,7 @@ impl Endpoint {
             remote_extensions,
             pgbouncer_settings: None,
             shard_stripe_size: Some(shard_stripe_size),
+            local_proxy_config: None,
         };
         let spec_path = self.endpoint_path().join("spec.json");
         std::fs::write(spec_path, serde_json::to_string_pretty(&spec)?)?;
diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs
index 83515a00a0..f6e4c3b944 100644
--- a/libs/compute_api/src/spec.rs
+++ b/libs/compute_api/src/spec.rs
@@ -106,6 +106,10 @@ pub struct ComputeSpec {
     // Stripe size for pageserver sharding, in pages
     #[serde(default)]
     pub shard_stripe_size: Option<usize>,
+
+    /// Local Proxy configuration used for JWT authentication
+    #[serde(default)]
+    pub local_proxy_config: Option<LocalProxySpec>,
 }
 
 /// Feature flag to signal `compute_ctl` to enable certain experimental functionality.