forked from naim94a/udpt
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.rs
169 lines (148 loc) · 5.34 KB
/
main.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
use clap::ValueHint;
use log::{error, info, trace, warn};
mod config;
mod server;
mod stackvec;
mod tracker;
mod webserver;
use config::Configuration;
use std::process::exit;
fn setup_logging(cfg: &Configuration) {
let log_level = match cfg.get_log_level() {
None => log::LevelFilter::Info,
Some(level) => {
match level.as_str() {
"off" => log::LevelFilter::Off,
"trace" => log::LevelFilter::Trace,
"debug" => log::LevelFilter::Debug,
"info" => log::LevelFilter::Info,
"warn" => log::LevelFilter::Warn,
"error" => log::LevelFilter::Error,
_ => {
eprintln!("udpt: unknown log level encountered '{}'", level.as_str());
exit(-1);
}
}
}
};
if let Err(err) = fern::Dispatch::new()
.format(|out, message, record| {
out.finish(format_args!(
"{} [{}][{}] {}",
chrono::Local::now().format("%+"),
record.target(),
record.level(),
message
))
})
.level(log_level)
.chain(std::io::stdout())
.apply()
{
eprintln!("udpt: failed to initialize logging. {}", err);
std::process::exit(-1);
}
info!("logging initialized.");
}
#[tokio::main]
async fn main() {
let parser = clap::Command::new(env!("CARGO_PKG_NAME"))
.about(env!("CARGO_PKG_DESCRIPTION"))
.author(env!("CARGO_PKG_AUTHORS"))
.version(env!("CARGO_PKG_VERSION"))
.arg(
clap::Arg::new("config")
.short('c')
.value_hint(ValueHint::FilePath)
.help("Configuration file to load.")
.required(true),
);
let matches = parser.get_matches();
let cfg_path = matches.get_one::<String>("config").unwrap();
let cfg = match Configuration::load_file(cfg_path) {
Ok(v) => std::sync::Arc::new(v),
Err(e) => {
eprintln!("udpt: failed to open configuration: {}", e);
return;
}
};
setup_logging(&cfg);
let tracker_obj = match cfg.get_db_path() {
Some(path) => {
let file_path = std::path::Path::new(path);
if !file_path.exists() {
warn!("database file \"{}\" doesn't exist.", path);
tracker::TorrentTracker::new(cfg.get_mode().clone())
} else {
let mut input_file = match tokio::fs::File::open(file_path).await {
Ok(v) => v,
Err(err) => {
error!("failed to open \"{}\". error: {}", path.as_str(), err);
panic!("error opening file. check logs.");
}
};
match tracker::TorrentTracker::load_database(cfg.get_mode().clone(), &mut input_file).await {
Ok(v) => {
info!("database loaded.");
v
}
Err(err) => {
error!("failed to load database. error: {}", err);
panic!("failed to load database. check logs.");
}
}
}
}
None => tracker::TorrentTracker::new(cfg.get_mode().clone()),
};
let tracker = std::sync::Arc::new(tracker_obj);
if cfg.get_http_config().is_some() {
let https_tracker = tracker.clone();
let http_cfg = cfg.clone();
info!("Starting http server");
tokio::spawn(async move {
let http_cfg = http_cfg.get_http_config().unwrap();
let bind_addr = http_cfg.get_address();
let tokens = http_cfg.get_access_tokens();
let server = webserver::build_server(https_tracker, tokens.clone());
server.bind(bind_addr.parse::<std::net::SocketAddr>().unwrap()).await;
});
}
let udp_server = server::UDPTracker::new(cfg.clone(), tracker.clone())
.await
.expect("failed to bind udp socket");
trace!("Waiting for UDP packets");
let udp_server = tokio::spawn(async move {
if let Err(err) = udp_server.accept_packets().await {
eprintln!("error: {}", err);
}
});
let weak_tracker = std::sync::Arc::downgrade(&tracker);
if let Some(db_path) = cfg.get_db_path() {
let db_path = db_path.clone();
let interval = cfg.get_cleanup_interval().unwrap_or(600);
tokio::spawn(async move {
let interval = std::time::Duration::from_secs(interval);
let mut interval = tokio::time::interval(interval);
interval.tick().await; // first tick is immediate...
loop {
interval.tick().await;
if let Some(tracker) = weak_tracker.upgrade() {
tracker.periodic_task(&db_path).await;
} else {
break;
}
}
});
}
let ctrl_c = tokio::signal::ctrl_c();
tokio::select! {
_ = udp_server => { warn!("udp server exited.") },
_ = ctrl_c => { info!("CTRL-C, exiting...") },
}
if let Some(path) = cfg.get_db_path() {
info!("saving database...");
tracker.periodic_task(path).await;
}
info!("goodbye.");
}