juiceback/
main.rs

1//! juiceback -- the brains behind Juicebox.
2//!
3//! This is where uploads come in (multipart or TUS), file records get saved
4//! to SQLite, admins log in, and everything gets pushed over to juicehost
5//! for storage. If you turn on the quic feature it also listens on a second
6//! port for HTTP/3, reusing the same router.
7
8use mimalloc::MiMalloc;
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::net::TcpListener;
12use tokio::signal;
13#[cfg(feature = "quic")]
14use tokio::sync::Notify;
15use tracing_subscriber::{
16    fmt::format::FmtSpan,
17    layer::SubscriberExt,
18    util::SubscriberInitExt,
19};
20
21#[global_allocator]
22static GLOBAL: MiMalloc = MiMalloc;
23
24mod auth;
25mod config;
26mod db;
27mod error;
28mod jobs;
29mod juicehost;
30#[cfg(feature = "quic")]
31mod quic;
32mod routes;
33mod state;
34mod tus;
35
36use crate::config::Config;
37use crate::state::AppState;
38
39async fn shutdown_signal() {
40    let ctrl_c = async {
41        signal::ctrl_c()
42            .await
43            .expect("Failed to install Ctrl+C handler");
44    };
45
46    #[cfg(unix)]
47    let terminate = async {
48        signal::unix::signal(signal::unix::SignalKind::terminate())
49            .expect("Failed to install signal handler")
50            .recv()
51            .await;
52    };
53
54    #[cfg(not(unix))]
55    let terminate = std::future::pending::<()>();
56
57    tokio::select! {
58        _ = ctrl_c => {
59            tracing::info!("juiceback shutting down...");
60        },
61        _ = terminate => {
62            tracing::info!("juiceback terminated...");
63        },
64    }
65}
66
67#[tokio::main]
68async fn main() -> Result<(), Box<dyn std::error::Error>> {
69    let args: Vec<String> = std::env::args().collect();
70    if args.len() > 1 && args[1] == "hash-password" {
71        let password = args.get(2).expect("usage: juiceback hash-password <password>");
72        let hash = auth::hash_password(password).expect("hashing failed");
73        println!("{}", hash);
74        return Ok(());
75    }
76
77    dotenvy::dotenv().ok();
78
79    let config = Config::load().expect("Failed to load configuration");
80
81    tracing_subscriber::registry()
82        .with(
83            tracing_subscriber::EnvFilter::try_from_default_env()
84                .unwrap_or_else(|_| config.log_level.clone().into()),
85        )
86        .with(tracing_subscriber::fmt::layer().with_span_events(FmtSpan::CLOSE))
87        .init();
88
89    tracing::info!("juiceback starting");
90    tracing::info!("files: {:?}", config.files_dir);
91    tracing::info!("db: {}", config.database_path);
92
93    tokio::fs::create_dir_all(&config.files_dir)
94        .await
95        .expect("did a great job with permissions congrats, i can't make the folder");
96
97    let manager = r2d2_sqlite::SqliteConnectionManager::file(&config.database_path);
98    let pool = r2d2::Pool::builder()
99        .max_size(config.db_pool_size)
100        .connection_timeout(Duration::from_secs(5))
101        .build(manager)
102        .expect("Failed to create database connection pool");
103
104    {
105        let conn = pool.get().expect("Failed to get connection for initialization");
106        conn.execute_batch(
107            "PRAGMA journal_mode=WAL;
108             PRAGMA synchronous=NORMAL;
109             PRAGMA busy_timeout=5000;"
110        ).expect("Failed to configure SQLite pragmas");
111        db::init_db(&conn).expect("Failed to initialize database schema");
112    }
113    tracing::info!("Database initialized");
114
115    let state = AppState::new(pool, config);
116
117    let state_cleanup = Arc::clone(&state);
118    tokio::spawn(async move {
119        jobs::cleanup::run_cleanup_loop(state_cleanup).await;
120    });
121
122    let app = routes::build_router(Arc::clone(&state));
123
124    let addr = format!("{}:{}", state.config.host, state.config.port);
125    let listener = TcpListener::bind(&addr)
126        .await
127        .expect("Failed to bind to address");
128
129    #[cfg(feature = "quic")]
130    {
131        let shutdown = Arc::new(Notify::new());
132        let quic_shutdown = Arc::clone(&shutdown);
133
134        let quic_listen: std::net::SocketAddr = format!("{}:{}", state.config.host, 6402)
135            .parse()
136            .expect("Invalid QUIC address");
137
138        tokio::select! {
139            _ = async {
140                quic::start_quic_server(Arc::clone(&state), quic_listen, quic_shutdown).await;
141            } => {}
142            _ = async {
143                tracing::info!("Listening on http://{}", addr);
144                axum::serve(
145                    listener,
146                    app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
147                )
148                .with_graceful_shutdown(shutdown_signal())
149                .await
150                .expect("Server failed");
151            } => {}
152            _ = shutdown_signal() => {
153                tracing::info!("juiceback shutting down...");
154                shutdown.notify_one();
155            }
156        }
157    }
158
159    #[cfg(not(feature = "quic"))]
160    {
161        tracing::info!("Listening on http://{}", addr);
162        axum::serve(
163            listener,
164            app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
165        )
166        .with_graceful_shutdown(shutdown_signal())
167        .await
168        .expect("Server failed");
169    }
170
171    Ok(())
172}
juiceback/main.rs

juiceback/
main.rs
