···
1
1
-
use clap::Parser;
1
1
+
use clap::{ArgAction, Parser};
2
2
use tokio_util::sync::CancellationToken;
3
3
use who_am_i::serve;
4
4
···
18
18
/// enables automatic template reloading
19
19
#[arg(long, action)]
20
20
dev: bool,
21
21
+
/// Hosts who are allowed to one-click auth
22
22
+
///
23
23
+
/// Pass this argument multiple times to allow multiple hosts
24
24
+
#[arg(long, short = 'o', action = ArgAction::Append)]
25
25
+
one_click: Vec<String>,
21
26
}
22
27
23
28
#[tokio::main]
···
29
34
30
35
let args = Args::parse();
31
36
32
32
-
serve(shutdown, args.app_secret, args.dev).await;
37
37
+
if args.one_click.is_empty() {
38
38
+
panic!("at least one --one-click host must be set");
39
39
+
}
40
40
+
41
41
+
println!("starting with allowed hosts:");
42
42
+
for host in &args.one_click {
43
43
+
println!(" - {host}");
44
44
+
}
45
45
+
46
46
+
serve(shutdown, args.app_secret, args.one_click, args.dev).await;
33
47
}
···
13
13
14
14
use serde::Deserialize;
15
15
use serde_json::json;
16
16
+
use std::collections::HashSet;
16
17
use std::sync::Arc;
17
18
use std::time::Duration;
18
19
use tokio::net::TcpListener;
···
31
32
#[derive(Clone)]
32
33
struct AppState {
33
34
pub key: Key,
35
35
+
pub one_clicks: Arc<HashSet<String>>,
34
36
pub engine: AppEngine,
35
37
pub client: Arc<Client>,
36
38
pub resolving: ExpiringTaskMap<Option<String>>,
···
43
45
}
44
46
}
45
47
46
46
-
pub async fn serve(shutdown: CancellationToken, app_secret: String, dev: bool) {
48
48
+
pub async fn serve(
49
49
+
shutdown: CancellationToken,
50
50
+
app_secret: String,
51
51
+
one_click: Vec<String>,
52
52
+
dev: bool,
53
53
+
) {
47
54
let mut hbs = Handlebars::new();
48
55
hbs.set_dev_mode(dev);
49
56
hbs.register_templates_directory("templates", Default::default())
···
58
65
let state = AppState {
59
66
engine: Engine::new(hbs),
60
67
key: Key::from(app_secret.as_bytes()), // TODO: via config
68
68
+
one_clicks: Arc::new(HashSet::from_iter(one_click)),
61
69
client: Arc::new(client()),
62
70
resolving: ExpiringTaskMap::new(task_pickup_expiration),
63
71
shutdown: shutdown.clone(),
···
85
93
async fn prompt(
86
94
State(AppState {
87
95
engine,
96
96
+
one_clicks,
88
97
resolving,
89
98
shutdown,
90
99
..
···
104
113
let Some(parent_host) = url.host_str() else {
105
114
return "could nto get host from url".into_response();
106
115
};
116
116
+
if !one_clicks.contains(parent_host) {
117
117
+
return format!("host {parent_host:?} not in one_clicks, disallowing for now")
118
118
+
.into_response();
119
119
+
}
107
120
if let Some(did) = jar.get(DID_COOKIE_KEY) {
108
121
let did = did.value_trimmed().to_string();
109
122