From 6fdd9c10d18270a5e30704f17e573ea14ee978ce Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki@neon.tech>
Date: Tue, 21 Mar 2023 15:24:49 +0200
Subject: [PATCH] Read storage auth token from spec file.

We read the pageserver connection string from the spec file, so let's
read the auth token from the same place.

We've been talking about pre-launching compute nodes that are not
associated with any particular tenant at startup, so that the spec
file is delivered to the compute node later. We cannot change the env
variables after the process has been launched.

We still pass the token to 'postgres' binary in the NEON_AUTH_TOKEN
env variable, but compute_ctl is now responsible for setting it.
---
 compute_tools/src/bin/compute_ctl.rs |  2 ++
 compute_tools/src/compute.rs         | 29 +++++++++++++++++-----------
 compute_tools/src/spec.rs            |  2 ++
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/compute_tools/src/bin/compute_ctl.rs b/compute_tools/src/bin/compute_ctl.rs
index a4e9262072..b96842e416 100644
--- a/compute_tools/src/bin/compute_ctl.rs
+++ b/compute_tools/src/bin/compute_ctl.rs
@@ -133,6 +133,7 @@ fn main() -> Result<()> {
         .settings
         .find("neon.pageserver_connstring")
         .expect("pageserver connstr should be provided");
+    let storage_auth_token = spec.storage_auth_token.clone();
     let tenant = spec
         .cluster
         .settings
@@ -153,6 +154,7 @@ fn main() -> Result<()> {
         tenant,
         timeline,
         pageserver_connstr,
+        storage_auth_token,
         metrics: ComputeMetrics::default(),
         state: RwLock::new(ComputeState::new()),
     };
diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index 09272262de..00d1e234ab 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -45,6 +45,7 @@ pub struct ComputeNode {
     pub tenant: String,
     pub timeline: String,
     pub pageserver_connstr: String,
+    pub storage_auth_token: Option<String>,
     pub metrics: ComputeMetrics,
     /// Volatile part of the `ComputeNode` so should be used under `RwLock`
     /// to allow HTTP API server to serve status requests, while configuration
@@ -129,18 +130,14 @@ impl ComputeNode {
 
         let mut config = postgres::Config::from_str(&self.pageserver_connstr)?;
 
-        // Like in the neon extension, if the $NEON_AUTH_TOKEN env variable is
-        // set, use it as the password when connecting to pageserver.
-        //
+        // Use the storage auth token from the config file, if given.
         // Note: this overrides any password set in the connection string.
-        match std::env::var("NEON_AUTH_TOKEN") {
-            Ok(val) => {
-                info!("Got pageserver auth token from NEON_AUTH_TOKEN env variable");
-                config.password(val);
-            }
-            Err(std::env::VarError::NotPresent) => info!("NEON_AUTH_TOKEN env variable not set"),
-            Err(e) => info!("could not parse NEON_AUTH_TOKEN env variable: {}", e),
-        };
+        if let Some(storage_auth_token) = &self.storage_auth_token {
+            info!("Got storage auth token from spec file");
+            config.password(storage_auth_token);
+        } else {
+            info!("Storage auth token not set");
+        }
 
         let mut client = config.connect(NoTls)?;
         let basebackup_cmd = match lsn {
@@ -179,6 +176,11 @@ impl ComputeNode {
         let sync_handle = Command::new(&self.pgbin)
             .args(["--sync-safekeepers"])
             .env("PGDATA", &self.pgdata) // we cannot use -D in this mode
+            .envs(if let Some(storage_auth_token) = &self.storage_auth_token {
+                vec![("NEON_AUTH_TOKEN", storage_auth_token)]
+            } else {
+                vec![]
+            })
             .stdout(Stdio::piped())
             .spawn()
             .expect("postgres --sync-safekeepers failed to start");
@@ -256,6 +258,11 @@ impl ComputeNode {
         // Run postgres as a child process.
         let mut pg = Command::new(&self.pgbin)
             .args(["-D", &self.pgdata])
+            .envs(if let Some(storage_auth_token) = &self.storage_auth_token {
+                vec![("NEON_AUTH_TOKEN", storage_auth_token)]
+            } else {
+                vec![]
+            })
             .spawn()
             .expect("cannot start postgres process");
 
diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs
index 47f1d69cff..9694ba9a88 100644
--- a/compute_tools/src/spec.rs
+++ b/compute_tools/src/spec.rs
@@ -24,6 +24,8 @@ pub struct ComputeSpec {
     pub cluster: Cluster,
     pub delta_operations: Option<Vec<DeltaOp>>,
 
+    pub storage_auth_token: Option<String>,
+
     pub startup_tracing_context: Option<HashMap<String, String>>,
 }