Skip to content

Commit a5737d5

Browse files
therealalephtherealaleph
authored
Merge pull request #4 from v4g4b0nd-0x76/fix/gracefull-shutdown-ui
Thanks @v4g4b0nd-0x76 — proper listener teardown on Stop is exactly what was needed. The 2-second grace window + force-abort fallback is a clean pattern.
2 parents 84503ec + 2290738 commit a5737d5

3 files changed

Lines changed: 227 additions & 132 deletions

File tree

‎src/bin/ui.rs‎

Lines changed: 154 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -95,10 +95,16 @@ enum Cmd {
9595
PollStats,
9696
/// Probe a single SNI against the given google_ip. Result is written
9797
/// into UiState::sni_probe keyed by the SNI string.
98-
TestSni { google_ip: String, sni: String },
98+
TestSni {
99+
google_ip: String,
100+
sni: String,
101+
},
99102
/// Probe a batch of SNI names. Results appear in UiState::sni_probe one
100103
/// by one as each probe finishes.
101-
TestAllSni { google_ip: String, snis: Vec<String> },
104+
TestAllSni {
105+
google_ip: String,
106+
snis: Vec<String>,
107+
},
102108
}
103109

104110
struct App {
@@ -213,7 +219,10 @@ fn sni_pool_for_form(user: Option<&[String]>, front_domain: &str) -> Vec<SniRow>
213219
if !user_clean.is_empty() {
214220
return user_clean
215221
.into_iter()
216-
.map(|name| SniRow { name, enabled: true })
222+
.map(|name| SniRow {
223+
name,
224+
enabled: true,
225+
})
217226
.collect();
218227
}
219228
// Default: primary + the other Google-edge subdomains, primary first,
@@ -223,11 +232,17 @@ fn sni_pool_for_form(user: Option<&[String]>, front_domain: &str) -> Vec<SniRow>
223232
let mut out = Vec::new();
224233
if !primary.is_empty() {
225234
seen.insert(primary.clone());
226-
out.push(SniRow { name: primary, enabled: true });
235+
out.push(SniRow {
236+
name: primary,
237+
enabled: true,
238+
});
227239
}
228240
for s in DEFAULT_GOOGLE_SNI_POOL {
229241
if seen.insert(s.to_string()) {
230-
out.push(SniRow { name: (*s).to_string(), enabled: true });
242+
out.push(SniRow {
243+
name: (*s).to_string(),
244+
enabled: true,
245+
});
231246
}
232247
}
233248
out
@@ -281,7 +296,11 @@ impl FormState {
281296
enable_batching: false,
282297
upstream_socks5: {
283298
let v = self.upstream_socks5.trim();
284-
if v.is_empty() { None } else { Some(v.to_string()) }
299+
if v.is_empty() {
300+
None
301+
} else {
302+
Some(v.to_string())
303+
}
285304
},
286305
parallel_relay: self.parallel_relay,
287306
sni_hosts: {
@@ -296,7 +315,11 @@ impl FormState {
296315
// If the user's pool is empty/all-off we still save as None so
297316
// the backend falls back to sensible defaults instead of dying
298317
// on an empty pool.
299-
if active.is_empty() { None } else { Some(active) }
318+
if active.is_empty() {
319+
None
320+
} else {
321+
Some(active)
322+
}
300323
},
301324
})
302325
}
@@ -307,8 +330,7 @@ fn save_config(cfg: &Config) -> Result<PathBuf, String> {
307330
if let Some(parent) = path.parent() {
308331
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
309332
}
310-
let json = serde_json::to_string_pretty(&ConfigWire::from(cfg))
311-
.map_err(|e| e.to_string())?;
333+
let json = serde_json::to_string_pretty(&ConfigWire::from(cfg)).map_err(|e| e.to_string())?;
312334
std::fs::write(&path, json).map_err(|e| e.to_string())?;
313335
Ok(path)
314336
}
@@ -368,7 +390,10 @@ impl<'a> From<&'a Config> for ConfigWire<'a> {
368390
hosts: &c.hosts,
369391
upstream_socks5: c.upstream_socks5.as_deref(),
370392
parallel_relay: c.parallel_relay,
371-
sni_hosts: c.sni_hosts.as_ref().map(|v| v.iter().map(String::as_str).collect()),
393+
sni_hosts: c
394+
.sni_hosts
395+
.as_ref()
396+
.map(|v| v.iter().map(String::as_str).collect()),
372397
}
373398
}
374399
}
@@ -766,14 +791,13 @@ impl App {
766791
});
767792
}
768793
}
769-
if ui.button("Keep working only").on_hover_text(
770-
"Uncheck every SNI that didn't pass the last probe."
771-
).clicked() {
794+
if ui
795+
.button("Keep working only")
796+
.on_hover_text("Uncheck every SNI that didn't pass the last probe.")
797+
.clicked()
798+
{
772799
for row in &mut self.form.sni_pool {
773-
let ok = matches!(
774-
probe_map.get(&row.name),
775-
Some(SniProbeState::Ok(_))
776-
);
800+
let ok = matches!(probe_map.get(&row.name), Some(SniProbeState::Ok(_)));
777801
row.enabled = ok;
778802
}
779803
}
@@ -785,13 +809,20 @@ impl App {
785809
if ui.button("Clear status").clicked() {
786810
self.shared.state.lock().unwrap().sni_probe.clear();
787811
}
788-
if ui.button("Reset to defaults").on_hover_text(
789-
"Replace the list with the built-in Google SNI pool. Custom entries \
790-
are dropped."
791-
).clicked() {
812+
if ui
813+
.button("Reset to defaults")
814+
.on_hover_text(
815+
"Replace the list with the built-in Google SNI pool. Custom entries \
816+
are dropped.",
817+
)
818+
.clicked()
819+
{
792820
self.form.sni_pool = DEFAULT_GOOGLE_SNI_POOL
793821
.iter()
794-
.map(|s| SniRow { name: (*s).to_string(), enabled: true })
822+
.map(|s| SniRow {
823+
name: (*s).to_string(),
824+
enabled: true,
825+
})
795826
.collect();
796827
self.shared.state.lock().unwrap().sni_probe.clear();
797828
}
@@ -804,51 +835,55 @@ impl App {
804835
let mut test_name: Option<String> = None;
805836
const STATUS_W: f32 = 150.0;
806837
const NAME_W: f32 = 230.0;
807-
egui::ScrollArea::vertical().max_height(280.0).show(ui, |ui| {
808-
for (i, row) in self.form.sni_pool.iter_mut().enumerate() {
809-
ui.horizontal(|ui| {
810-
ui.checkbox(&mut row.enabled, "");
811-
ui.add(
812-
egui::TextEdit::singleline(&mut row.name)
813-
.desired_width(NAME_W)
814-
.font(egui::TextStyle::Monospace),
815-
);
816-
let status_txt = match probe_map.get(&row.name) {
817-
Some(SniProbeState::Ok(ms)) => {
818-
egui::RichText::new(format!("ok {} ms", ms))
819-
.color(egui::Color32::from_rgb(80, 180, 100))
820-
.monospace()
821-
}
822-
Some(SniProbeState::Failed(e)) => {
823-
let short = if e.len() > 22 { &e[..22] } else { e };
824-
egui::RichText::new(format!("fail {}", short))
825-
.color(egui::Color32::from_rgb(220, 110, 110))
826-
.monospace()
827-
}
828-
Some(SniProbeState::InFlight) => {
829-
egui::RichText::new("testing…")
838+
egui::ScrollArea::vertical()
839+
.max_height(280.0)
840+
.show(ui, |ui| {
841+
for (i, row) in self.form.sni_pool.iter_mut().enumerate() {
842+
ui.horizontal(|ui| {
843+
ui.checkbox(&mut row.enabled, "");
844+
ui.add(
845+
egui::TextEdit::singleline(&mut row.name)
846+
.desired_width(NAME_W)
847+
.font(egui::TextStyle::Monospace),
848+
);
849+
let status_txt = match probe_map.get(&row.name) {
850+
Some(SniProbeState::Ok(ms)) => {
851+
egui::RichText::new(format!("ok {} ms", ms))
852+
.color(egui::Color32::from_rgb(80, 180, 100))
853+
.monospace()
854+
}
855+
Some(SniProbeState::Failed(e)) => {
856+
let short = if e.len() > 22 { &e[..22] } else { e };
857+
egui::RichText::new(format!("fail {}", short))
858+
.color(egui::Color32::from_rgb(220, 110, 110))
859+
.monospace()
860+
}
861+
Some(SniProbeState::InFlight) => {
862+
egui::RichText::new("testing…")
863+
.color(egui::Color32::GRAY)
864+
.monospace()
865+
}
866+
None => egui::RichText::new("untested")
830867
.color(egui::Color32::GRAY)
831-
.monospace()
868+
.monospace(),
869+
};
870+
ui.add_sized(
871+
[STATUS_W, 18.0],
872+
egui::Label::new(status_txt).truncate(),
873+
);
874+
if ui.small_button("Test").clicked() {
875+
test_name = Some(row.name.clone());
832876
}
833-
None => {
834-
egui::RichText::new("untested")
835-
.color(egui::Color32::GRAY)
836-
.monospace()
877+
if ui
878+
.small_button("remove")
879+
.on_hover_text("Remove this row")
880+
.clicked()
881+
{
882+
to_remove = Some(i);
837883
}
838-
};
839-
ui.add_sized([STATUS_W, 18.0], egui::Label::new(status_txt).truncate());
840-
if ui.small_button("Test").clicked() {
841-
test_name = Some(row.name.clone());
842-
}
843-
if ui.small_button("remove")
844-
.on_hover_text("Remove this row")
845-
.clicked()
846-
{
847-
to_remove = Some(i);
848-
}
849-
});
850-
}
851-
});
884+
});
885+
}
886+
});
852887

853888
if let Some(name) = test_name {
854889
let name = name.trim().to_string();
@@ -930,12 +965,16 @@ fn fmt_bytes(b: u64) -> String {
930965
fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
931966
let rt = Runtime::new().expect("failed to create tokio runtime");
932967

933-
let mut active: Option<(JoinHandle<()>, Arc<AsyncMutex<Option<Arc<DomainFronter>>>>)> = None;
968+
let mut active: Option<(
969+
JoinHandle<()>,
970+
Arc<AsyncMutex<Option<Arc<DomainFronter>>>>,
971+
tokio::sync::oneshot::Sender<()>,
972+
)> = None;
934973

935974
loop {
936975
match rx.recv_timeout(Duration::from_millis(250)) {
937976
Ok(Cmd::PollStats) => {
938-
if let Some((_, fronter_slot)) = &active {
977+
if let Some((_, fronter_slot, _)) = &active {
939978
let slot = fronter_slot.clone();
940979
let shared = shared.clone();
941980
rt.spawn(async move {
@@ -950,6 +989,7 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
950989
});
951990
}
952991
}
992+
// In background_thread function, modify the Cmd::Start handler:
953993
Ok(Cmd::Start(cfg)) => {
954994
if active.is_some() {
955995
push_log(&shared, "[ui] already running");
@@ -961,6 +1001,8 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
9611001
Arc::new(AsyncMutex::new(None));
9621002
let fronter_slot2 = fronter_slot.clone();
9631003

1004+
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
1005+
9641006
let handle = rt.spawn(async move {
9651007
let base = data_dir::data_dir();
9661008
let mitm = match MitmCertManager::new_in(&base) {
@@ -986,27 +1028,49 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
9861028
s.running = true;
9871029
s.started_at = Some(Instant::now());
9881030
}
989-
push_log(&shared2, &format!(
990-
"[ui] listening HTTP {}:{} SOCKS5 {}:{}",
991-
cfg.listen_host, cfg.listen_port,
992-
cfg.listen_host, cfg.socks5_port.unwrap_or(cfg.listen_port + 1)
993-
));
994-
let _ = server.run().await;
1031+
push_log(
1032+
&shared2,
1033+
&format!(
1034+
"[ui] listening HTTP {}:{} SOCKS5 {}:{}",
1035+
cfg.listen_host,
1036+
cfg.listen_port,
1037+
cfg.listen_host,
1038+
cfg.socks5_port.unwrap_or(cfg.listen_port + 1)
1039+
),
1040+
);
1041+
1042+
let _ = server.run(shutdown_rx).await;
1043+
9951044
shared2.state.lock().unwrap().running = false;
9961045
push_log(&shared2, "[ui] proxy stopped");
9971046
});
9981047

999-
active = Some((handle, fronter_slot));
1048+
active = Some((handle, fronter_slot, shutdown_tx));
10001049
}
1050+
10011051
Ok(Cmd::Stop) => {
1002-
if let Some((handle, _)) = active.take() {
1003-
handle.abort();
1052+
if let Some((handle, _, shutdown_tx)) = active.take() {
1053+
push_log(&shared, "[ui] stop requested");
1054+
let _ = shutdown_tx.send(());
1055+
1056+
// Give the proxy 2 seconds to shut down gracefully
1057+
rt.block_on(async {
1058+
tokio::select! {
1059+
_ = handle => {
1060+
push_log(&shared, "[ui] proxy stopped gracefully");
1061+
}
1062+
_ = tokio::time::sleep(tokio::time::Duration::from_secs(2)) => {
1063+
push_log(&shared, "[ui] shutdown timeout, forcing abort");
1064+
}
1065+
}
1066+
});
1067+
10041068
shared.state.lock().unwrap().running = false;
10051069
shared.state.lock().unwrap().started_at = None;
10061070
shared.state.lock().unwrap().last_stats = None;
1007-
push_log(&shared, "[ui] stop requested");
10081071
}
10091072
}
1073+
10101074
Ok(Cmd::Test(cfg)) => {
10111075
let shared2 = shared.clone();
10121076
push_log(&shared, "[ui] running test...");
@@ -1018,7 +1082,10 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
10181082
} else {
10191083
"Test failed — see terminal for details.".into()
10201084
};
1021-
push_log(&shared2, &format!("[ui] test result: {}", if ok { "pass" } else { "fail" }));
1085+
push_log(
1086+
&shared2,
1087+
&format!("[ui] test result: {}", if ok { "pass" } else { "fail" }),
1088+
);
10221089
// Also run ip scan on demand (cheap).
10231090
let _ = scan_ips::run(&cfg).await;
10241091
});
@@ -1055,7 +1122,9 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
10551122
let result = scan_sni::probe_one(&google_ip, &sni).await;
10561123
let state = match result.latency_ms {
10571124
Some(ms) => SniProbeState::Ok(ms),
1058-
None => SniProbeState::Failed(result.error.unwrap_or_else(|| "failed".into())),
1125+
None => {
1126+
SniProbeState::Failed(result.error.unwrap_or_else(|| "failed".into()))
1127+
}
10591128
};
10601129
shared2.state.lock().unwrap().sni_probe.insert(sni, state);
10611130
});
@@ -1074,7 +1143,9 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
10741143
for (sni, r) in results {
10751144
let state = match r.latency_ms {
10761145
Some(ms) => SniProbeState::Ok(ms),
1077-
None => SniProbeState::Failed(r.error.unwrap_or_else(|| "failed".into())),
1146+
None => {
1147+
SniProbeState::Failed(r.error.unwrap_or_else(|| "failed".into()))
1148+
}
10781149
};
10791150
st.sni_probe.insert(sni, state);
10801151
}
@@ -1093,7 +1164,7 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
10931164
}
10941165

10951166
// Clean up finished task.
1096-
if let Some((handle, _)) = &active {
1167+
if let Some((handle, _, _)) = &active {
10971168
if handle.is_finished() {
10981169
active = None;
10991170
shared.state.lock().unwrap().running = false;
@@ -1105,7 +1176,9 @@ fn background_thread(shared: Arc<Shared>, rx: Receiver<Cmd>) {
11051176
fn push_log(shared: &Shared, msg: &str) {
11061177
let line = format!(
11071178
"{} {}",
1108-
time::OffsetDateTime::now_utc().format(&time::format_description::well_known::Iso8601::DEFAULT).unwrap_or_default(),
1179+
time::OffsetDateTime::now_utc()
1180+
.format(&time::format_description::well_known::Iso8601::DEFAULT)
1181+
.unwrap_or_default(),
11091182
msg
11101183
);
11111184
let mut s = shared.state.lock().unwrap();

0 commit comments

Comments
 (0)