From ba9dc3fa0b4f5d6d5bde76c88a2089d9909c91fd Mon Sep 17 00:00:00 2001
From: Jake Howard <git@theorangeone.net>
Date: Fri, 1 Sep 2023 17:11:15 +0100
Subject: [PATCH] Locate dokku root

---
 Cargo.lock  | 20 +++++++++++++++++++-
 Cargo.toml  |  1 +
 src/main.rs | 54 +++++++++++++++++++++++++++++++++++++++++++++++++----
 3 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 78c1087..2952524 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -36,7 +36,7 @@ checksum = "3b829e4e32b91e643de6eafe82b1d90675f5874230191a4ffbc1b336dec4d6bf"
 dependencies = [
  "async-trait",
  "axum-core",
- "bitflags",
+ "bitflags 1.3.2",
  "bytes",
  "futures-util",
  "http",
@@ -98,6 +98,12 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
+[[package]]
+name = "bitflags"
+version = "2.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4682ae6287fcf752ecaabbfcc7b6f9b72aa33933dc23a554d853aea8eea8635"
+
 [[package]]
 name = "bytes"
 version = "1.4.0"
@@ -124,6 +130,7 @@ name = "dokku-hosts-api"
 version = "0.1.0"
 dependencies = [
  "axum",
+ "nix",
  "tokio",
 ]
 
@@ -303,6 +310,17 @@ dependencies = [
  "windows-sys",
 ]
 
+[[package]]
+name = "nix"
+version = "0.27.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2eb04e9c688eff1c89d72b407f168cf79bb9e867a9d3323ed6c01519eb9cc053"
+dependencies = [
+ "bitflags 2.4.0",
+ "cfg-if",
+ "libc",
+]
+
 [[package]]
 name = "num_cpus"
 version = "1.15.0"
diff --git a/Cargo.toml b/Cargo.toml
index f2b92f7..4d315e4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,4 +5,5 @@ edition = "2021"
 
 [dependencies]
 axum = "0.6.20"
+nix = { version="0.27.1", features=["user"]}
 tokio = { version="1.32.0", features=["macros", "rt-multi-thread"]}
diff --git a/src/main.rs b/src/main.rs
index 5e3c881..6ea5cdf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,17 +1,63 @@
+use axum::extract::State;
 use axum::{routing::get, Router};
 use std::net::{IpAddr, Ipv6Addr, SocketAddr};
+use std::path::PathBuf;
 
-async fn root() -> &'static str {
-    "Hello, World!"
+#[derive(Clone)]
+struct AppState {
+    dokku_root: PathBuf,
+}
+
+fn get_dokku_root() -> Option<PathBuf> {
+    // First, check if we have a custom directory
+    if let Some(dokku_root) = std::env::var_os("DOKKU_ROOT") {
+        return Some(PathBuf::from(dokku_root));
+    }
+
+    // Next, check for a dokku user and its home directory
+    match nix::unistd::User::from_name("dokku") {
+        Ok(Some(user)) => return Some(user.dir),
+
+        // If the user doesn't exist, do nothing
+        Ok(None) => return None,
+
+        // If there was an error, give up
+        _ => {}
+    }
+
+    // Finally, try a hard-coded path
+    let dokku_home_guess = PathBuf::from("/home/dokku");
+    if dokku_home_guess.is_dir() {
+        return Some(dokku_home_guess);
+    }
+
+    // If there's nothing else, give up
+    None
+}
+
+async fn hosts(State(state): State<AppState>) -> String {
+    format!("{}", state.dokku_root.display())
 }
 
 #[tokio::main]
 async fn main() {
-    let app = Router::new().route("/hosts", get(root));
+    let dokku_root = match get_dokku_root() {
+        Some(path) => path,
+        None => {
+            println!("Failed to find dokku root");
+            std::process::exit(1);
+        }
+    };
+
+    println!("Using dokku root: {}", dokku_root.display());
+
+    let state = AppState { dokku_root };
+
+    let app = Router::new().route("/hosts", get(hosts)).with_state(state);
 
     let addr = &SocketAddr::new(IpAddr::from(Ipv6Addr::UNSPECIFIED), 3000);
     println!("Listening on {addr}");
-    axum::Server::bind(&addr)
+    axum::Server::bind(addr)
         .serve(app.into_make_service())
         .await
         .unwrap();