dev and prod with all the oauth joy

Cargo.lock

···
 checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
 dependencies = [
  "const-oid",
+ "pem-rfc7468",
  "zeroize",
 ]
 
···
  "elliptic-curve",
  "rfc6979",
  "signature",
+ "spki",
 ]
 
 [[package]]
···
  "ff",
  "generic-array",
  "group",
+ "pem-rfc7468",
+ "pkcs8",
  "rand_core 0.6.4",
  "sec1",
  "subtle",
···
  "jose-b64",
  "jose-jwa",
  "p256",
+ "p384",
+ "rsa",
  "serde",
  "zeroize",
 ]
···
 version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+dependencies = [
+ "spin",
+]
 
 [[package]]
 name = "lazycell"
···
 ]
 
 [[package]]
+name = "num-bigint-dig"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc84195820f291c7697304f3cbdadd1cb7199c0efc917ff5eafd71225c136151"
+dependencies = [
+ "byteorder",
+ "lazy_static",
+ "libm",
+ "num-integer",
+ "num-iter",
+ "num-traits",
+ "rand 0.8.5",
+ "smallvec",
+ "zeroize",
+]
+
+[[package]]
 name = "num-conv"
 version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
···
 ]
 
 [[package]]
+name = "num-iter"
+version = "0.1.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf"
+dependencies = [
+ "autocfg",
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
 name = "num-modular"
 version = "0.6.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
···
 checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
 dependencies = [
  "autocfg",
+ "libm",
 ]
 
 [[package]]
···
 ]
 
 [[package]]
+name = "p384"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe42f1670a52a47d448f14b6a5c61dd78fce51856e68edaa38f7ae3a46b8d6b6"
+dependencies = [
+ "elliptic-curve",
+ "primeorder",
+]
+
+[[package]]
 name = "parking"
 version = "2.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
···
 dependencies = [
  "base64 0.22.1",
  "serde",
+]
+
+[[package]]
+name = "pem-rfc7468"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412"
+dependencies = [
+ "base64ct",
 ]
 
 [[package]]
···
 checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
 
 [[package]]
+name = "pkcs1"
+version = "0.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f"
+dependencies = [
+ "der",
+ "pkcs8",
+ "spki",
+]
+
+[[package]]
+name = "pkcs8"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
+dependencies = [
+ "der",
+ "spki",
+]
+
+[[package]]
 name = "pkg-config"
 version = "0.3.32"
 source = "registry+https://github.com/rust-lang/crates.io-index"
···
 ]
 
 [[package]]
+name = "rsa"
+version = "0.9.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78928ac1ed176a5ca1d17e578a1825f3d81ca54cf41053a592584b020cfd691b"
+dependencies = [
+ "const-oid",
+ "digest",
+ "num-bigint-dig",
+ "num-integer",
+ "num-traits",
+ "pkcs1",
+ "pkcs8",
+ "rand_core 0.6.4",
+ "signature",
+ "spki",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
 name = "rustc-demangle"
 version = "0.1.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
···
  "base16ct",
  "der",
  "generic-array",
+ "pkcs8",
  "subtle",
  "zeroize",
 ]
···
 checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
 dependencies = [
  "lock_api",
+]
+
+[[package]]
+name = "spki"
+version = "0.7.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d"
+dependencies = [
+ "base64ct",
+ "der",
 ]
 
 [[package]]
···
  "clap",
  "ctrlc",
  "dashmap",
+ "elliptic-curve",
  "handlebars",
  "hickory-resolver",
+ "jose-jwk",
  "jsonwebtoken",
  "metrics",
  "metrics-exporter-prometheus 0.17.2",
+ "p256",
+ "pkcs8",
  "rand 0.9.1",
  "reqwest",
  "serde",

who-am-i/Cargo.toml

···
 clap = { version = "4.5.40", features = ["derive", "env"] }
 ctrlc = "3.4.7"
 dashmap = "6.1.0"
+elliptic-curve = "0.13.8"
 handlebars = { version = "6.3.2", features = ["dir_source"] }
 hickory-resolver = "0.25.2"
+jose-jwk = "0.1.2"
 jsonwebtoken = "9.3.1"
 metrics = "0.24.2"
+p256 = "0.13.2"
+pkcs8 = "0.10.2"
 rand = "0.9.1"
 reqwest = { version = "0.12.22", features = ["native-tls-vendored"] }
 serde = { version = "1.0.219", features = ["derive"] }

who-am-i/src/main.rs

···
     /// eg: `cat /dev/urandom | head -c 64 | base64`
     #[arg(long, env)]
     app_secret: String,
+    /// path to at-oauth private key (PEM pk8 format)
+    ///
+    /// generate with:
+    ///
+    ///     openssl ecparam -genkey -noout -name prime256v1 \
+    ///         | openssl pkcs8 -topk8 -nocrypt -out <PATH-TO-PRIV-KEY>.pem
+    #[arg(long, env)]
+    oauth_private_key: Option<PathBuf>,
     /// path to jwt private key (PEM pk8 format)
     ///
     /// generate with:
···
     /// wrap the jwk in an array, then in an object under "keys":
     ///
     ///     { "keys": [<JWK obj>] }
+    ///
+    /// TODO: remove this, serve automatically
     #[arg(long)]
     jwks: PathBuf,
+    /// this server's client-reachable base url, for oauth redirect + jwt check
+    ///
+    /// required unless running in localhost mode with --dev
+    #[arg(long, env)]
+    base_url: Option<String>,
+    /// host:port to bind to on startup
+    #[arg(long, env, default_value = "127.0.0.1:9997")]
+    bind: String,
     /// Enable dev mode
     ///
-    /// enables automatic template reloading
+    /// enables automatic template reloading, uses localhost oauth config, etc
     #[arg(long, action)]
     dev: bool,
     /// Hosts who are allowed to one-click auth
···
 
     let args = Args::parse();
 
+    // let bind = args.bind.to_socket_addrs().expect("--bind must be ToSocketAddrs");
+
+    let base = args.base_url.unwrap_or_else(|| {
+        if args.dev {
+            format!("http://{}", args.bind)
+        } else {
+            panic!("not in --dev mode so --base-url is required")
+        }
+    });
+
+    if !args.dev && args.oauth_private_key.is_none() {
+        panic!("--at-oauth-key is required except in --dev");
+    } else if args.dev && args.oauth_private_key.is_some() {
+        eprintln!("warn: --at-oauth-key is ignored in dev (localhost config)");
+    }
+
     if args.allowed_hosts.is_empty() {
         panic!("at least one --allowed-host host must be set");
     }
···
     serve(
         shutdown,
         args.app_secret,
+        args.oauth_private_key,
         tokens,
+        base,
+        args.bind,
         args.allowed_hosts,
         args.dev,
     )

who-am-i/src/oauth.rs

···
+use jose_jwk::Class;
+use jose_jwk::Jwk;
+use jose_jwk::Key;
+use jose_jwk::Parameters;
+use std::fs;
+use std::path::PathBuf;
+// use p256::SecretKey;
 use atrium_api::{agent::SessionManager, types::string::Did};
 use atrium_common::resolver::Resolver;
 use atrium_identity::{
···
     handle::{AtprotoHandleResolver, AtprotoHandleResolverConfig, DnsTxtResolver},
 };
 use atrium_oauth::{
-    AtprotoLocalhostClientMetadata, AuthorizeOptions, CallbackParams, DefaultHttpClient,
-    KnownScope, OAuthClient, OAuthClientConfig, OAuthResolverConfig, Scope,
+    AtprotoClientMetadata, AtprotoLocalhostClientMetadata, AuthMethod, AuthorizeOptions,
+    CallbackParams, DefaultHttpClient, GrantType, KnownScope, OAuthClient, OAuthClientConfig,
+    OAuthClientMetadata, OAuthResolverConfig, Scope,
     store::{session::MemorySessionStore, state::MemoryStateStore},
 };
+use elliptic_curve::SecretKey;
 use hickory_resolver::{ResolveError, TokioResolver};
+use jose_jwk::JwkSet;
+use pkcs8::DecodePrivateKey;
 use serde::Deserialize;
 use std::sync::Arc;
 use thiserror::Error;
···
 }
 
 impl OAuth {
-    pub fn new() -> Result<Self, AuthSetupError> {
+    pub fn new(oauth_private_key: Option<PathBuf>, base: String) -> Result<Self, AuthSetupError> {
         let http_client = Arc::new(DefaultHttpClient::default());
         let did_resolver = || {
             CommonDidResolver::new(CommonDidResolverConfig {
···
         };
         let dns_txt_resolver =
             HickoryDnsTxtResolver::new().map_err(AuthSetupError::HickoryResolverError)?;
-        let client_config = OAuthClientConfig {
-            client_metadata: AtprotoLocalhostClientMetadata {
-                redirect_uris: Some(vec![String::from("http://127.0.0.1:9997/authorized")]),
-                scopes: Some(READONLY_SCOPE.to_vec()),
-            },
-            keys: None,
-            resolver: OAuthResolverConfig {
-                did_resolver: did_resolver(),
-                handle_resolver: AtprotoHandleResolver::new(AtprotoHandleResolverConfig {
-                    dns_txt_resolver,
-                    http_client: Arc::clone(&http_client),
-                }),
-                authorization_server_metadata: Default::default(),
-                protected_resource_metadata: Default::default(),
-            },
-            state_store: MemoryStateStore::default(),
-            session_store: MemorySessionStore::default(),
+
+        let resolver = OAuthResolverConfig {
+            did_resolver: did_resolver(),
+            handle_resolver: AtprotoHandleResolver::new(AtprotoHandleResolverConfig {
+                dns_txt_resolver,
+                http_client: Arc::clone(&http_client),
+            }),
+            authorization_server_metadata: Default::default(),
+            protected_resource_metadata: Default::default(),
         };
 
-        let client = OAuthClient::new(client_config).map_err(AuthSetupError::AtriumClientError)?;
+        let state_store = MemoryStateStore::default();
+        let session_store = MemorySessionStore::default();
+
+        let client = if let Some(path) = oauth_private_key {
+            let key_contents: Vec<u8> = fs::read(path).unwrap();
+            let key_string = String::from_utf8(key_contents).unwrap();
+            let key = SecretKey::<p256::NistP256>::from_pkcs8_pem(&key_string)
+                .map(|secret_key| Jwk {
+                    key: Key::from(&secret_key.into()),
+                    prm: Parameters {
+                        kid: Some("at-oauth-00".to_string()),
+                        cls: Some(Class::Signing),
+                        ..Default::default()
+                    },
+                })
+                .expect("to get private key");
+            OAuthClient::new(OAuthClientConfig {
+                client_metadata: AtprotoClientMetadata {
+                    client_id: format!("{base}/client-metadata.json"),
+                    client_uri: Some(base.clone()),
+                    redirect_uris: vec![format!("{base}/authorized")],
+                    token_endpoint_auth_method: AuthMethod::PrivateKeyJwt,
+                    grant_types: vec![GrantType::AuthorizationCode, GrantType::RefreshToken],
+                    scopes: READONLY_SCOPE.to_vec(),
+                    jwks_uri: Some(format!("{base}/.well-known/at-jwks.json")),
+                    token_endpoint_auth_signing_alg: Some(String::from("ES256")),
+                },
+                keys: Some(vec![key]),
+                resolver,
+                state_store,
+                session_store,
+            })
+            .map_err(AuthSetupError::AtriumClientError)?
+        } else {
+            OAuthClient::new(OAuthClientConfig {
+                client_metadata: AtprotoLocalhostClientMetadata {
+                    redirect_uris: Some(vec![String::from("http://127.0.0.1:9997/authorized")]),
+                    scopes: Some(READONLY_SCOPE.to_vec()),
+                },
+                keys: None,
+                resolver,
+                state_store,
+                session_store,
+            })
+            .map_err(AuthSetupError::AtriumClientError)?
+        };
 
         Ok(Self {
             client: Arc::new(client),
             did_resolver: Arc::new(did_resolver()),
         })
+    }
+
+    pub fn client_metadata(&self) -> OAuthClientMetadata {
+        self.client.client_metadata.clone()
+    }
+
+    pub fn jwks(&self) -> JwkSet {
+        self.client.jwks()
     }
 
     pub async fn begin(&self, handle: &str) -> Result<String, atrium_oauth::Error> {

who-am-i/src/server.rs

···
 use atrium_api::types::string::Did;
+use atrium_oauth::OAuthClientMetadata;
 use axum::{
     Router,
     extract::{FromRef, Json as ExtractJson, Query, State},
···
 use axum_extra::extract::cookie::{Cookie, Key, SameSite, SignedCookieJar};
 use axum_template::{RenderHtml, engine::Engine};
 use handlebars::{Handlebars, handlebars_helper};
+use jose_jwk::JwkSet;
+use std::path::PathBuf;
 
 use serde::Deserialize;
 use serde_json::{Value, json};
···
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 pub async fn serve(
     shutdown: CancellationToken,
     app_secret: String,
+    oauth_private_key: Option<PathBuf>,
     tokens: Tokens,
+    base: String,
+    bind: String,
     allowed_hosts: Vec<String>,
     dev: bool,
 ) {
···
     // clients have to pick up their identity-resolving tasks within this period
     let task_pickup_expiration = Duration::from_secs(15);
 
-    let oauth = OAuth::new().unwrap();
+    let oauth = OAuth::new(oauth_private_key, base).unwrap();
 
     let state = AppState {
         engine: Engine::new(hbs),
···
         .route("/style.css", get(css))
         .route("/prompt", get(prompt))
         .route("/user-info", post(user_info))
+        .route("/client-metadata.json", get(client_metadata))
         .route("/auth", get(start_oauth))
         .route("/authorized", get(complete_oauth))
         .route("/disconnect", post(disconnect))
+        .route("/.well-known/at-jwks.json", get(at_jwks)) // todo combine jwks eps (key id is enough?)
         .route("/.well-known/jwks.json", get(jwks))
         .with_state(state);
 
-    let listener = TcpListener::bind("0.0.0.0:9997")
+    eprintln!("starting server at http://{bind}");
+    let listener = TcpListener::bind(bind)
         .await
         .expect("listener binding to work");
 
···
             Json(json!({ "handle": handle })).into_response()
         }
     }
+}
+
+async fn client_metadata(
+    State(AppState { oauth, .. }): State<AppState>,
+) -> Json<OAuthClientMetadata> {
+    Json(oauth.client_metadata())
+}
+
+async fn at_jwks(State(AppState { oauth, .. }): State<AppState>) -> Json<JwkSet> {
+    Json(oauth.jwks())
 }
 
 #[derive(Debug, Deserialize)]