1#![deny(clippy::pedantic)]
2
3use std::{
4 borrow::Cow,
5 fmt::{Display, Formatter},
6 future::IntoFuture,
7 net::SocketAddr,
8 path::PathBuf,
9 str::FromStr,
10 sync::Arc,
11 time::Duration,
12};
13
14use anyhow::Context;
15use askama::Template;
16use axum::{
17 body::Body,
18 http,
19 http::{HeaderValue, StatusCode},
20 response::{IntoResponse, Response},
21 routing::get,
22 Extension, Router,
23};
24use bat::assets::HighlightingAssets;
25use clap::Parser;
26use database::schema::SCHEMA_VERSION;
27use nom::AsBytes;
28use once_cell::sync::{Lazy, OnceCell};
29use rocksdb::{Options, SliceTransform};
30use sha2::{digest::FixedOutput, Digest};
31use syntect::html::ClassStyle;
32use tokio::{
33 net::TcpListener,
34 signal::unix::{signal, SignalKind},
35 sync::mpsc,
36};
37use tower_http::cors::CorsLayer;
38use tower_layer::layer_fn;
39use tracing::{error, info, instrument, warn};
40use tracing_subscriber::EnvFilter;
41
42use crate::{
43 database::schema::prefixes::{
44 COMMIT_COUNT_FAMILY, COMMIT_FAMILY, REFERENCE_FAMILY, REPOSITORY_FAMILY, TAG_FAMILY,
45 },
46 git::Git,
47 layers::logger::LoggingMiddleware,
48};
49
50mod database;
51mod git;
52mod layers;
53mod methods;
54mod syntax_highlight;
55
56const CRATE_VERSION: &str = clap::crate_version!();
57
58static GLOBAL_CSS: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/statics/css/style.css",));
59static GLOBAL_CSS_HASH: Lazy<Box<str>> = Lazy::new(|| build_asset_hash(GLOBAL_CSS));
60
61static HIGHLIGHT_CSS_HASH: OnceCell<Box<str>> = OnceCell::new();
62static DARK_HIGHLIGHT_CSS_HASH: OnceCell<Box<str>> = OnceCell::new();
63
64#[derive(Parser, Debug)]
65#[clap(author, version, about)]
66pub struct Args {
67 /// Path to a directory in which the RocksDB database should be stored, will be created if it doesn't already exist
68 ///
69 /// The RocksDB database is very quick to generate, so this can be pointed to temporary storage
70 #[clap(short, long, value_parser)]
71 db_store: PathBuf,
72 /// The socket address to bind to (eg. 0.0.0.0:3333)
73 bind_address: SocketAddr,
74 /// The path in which your bare Git repositories reside (will be scanned recursively)
75 scan_path: PathBuf,
76 /// Configures the metadata refresh interval (eg. "never" or "60s")
77 #[clap(long, default_value_t = RefreshInterval::Duration(Duration::from_secs(300)))]
78 refresh_interval: RefreshInterval,
79}
80
81#[derive(Debug, Clone, Copy)]
82pub enum RefreshInterval {
83 Never,
84 Duration(Duration),
85}
86
87impl Display for RefreshInterval {
88 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
89 match self {
90 Self::Never => write!(f, "never"),
91 Self::Duration(s) => write!(f, "{}", humantime::format_duration(*s)),
92 }
93 }
94}
95
96impl FromStr for RefreshInterval {
97 type Err = &'static str;
98
99 fn from_str(s: &str) -> Result<Self, Self::Err> {
100 if s == "never" {
101 Ok(Self::Never)
102 } else if let Ok(v) = humantime::parse_duration(s) {
103 Ok(Self::Duration(v))
104 } else {
105 Err("must be seconds, a human readable duration (eg. '10m') or 'never'")
106 }
107 }
108}
109
110#[tokio::main]
111async fn main() -> Result<(), anyhow::Error> {
112 let args: Args = Args::parse();
113
114 if std::env::var_os("RUST_LOG").is_none() {
115 std::env::set_var("RUST_LOG", "info");
116 }
117
118 let subscriber = tracing_subscriber::fmt().with_env_filter(EnvFilter::from_default_env());
119 #[cfg(debug_assertions)]
120 let subscriber = subscriber.pretty();
121 subscriber.init();
122
123 let db = open_db(&args)?;
124
125 let indexer_wakeup_task =
126 run_indexer(db.clone(), args.scan_path.clone(), args.refresh_interval);
127
128 let bat_assets = HighlightingAssets::from_binary();
129 let syntax_set = bat_assets.get_syntax_set().unwrap().clone();
130
131 let theme = bat_assets.get_theme("GitHub");
132 let css = syntect::html::css_for_theme_with_class_style(theme, ClassStyle::Spaced).unwrap();
133 let css = Box::leak(
134 format!(r#"@media (prefers-color-scheme: light){{{css}}}"#)
135 .into_boxed_str()
136 .into_boxed_bytes(),
137 );
138 HIGHLIGHT_CSS_HASH.set(build_asset_hash(css)).unwrap();
139
140 let dark_theme = bat_assets.get_theme("TwoDark");
141 let dark_css =
142 syntect::html::css_for_theme_with_class_style(dark_theme, ClassStyle::Spaced).unwrap();
143 let dark_css = Box::leak(
144 format!(r#"@media (prefers-color-scheme: dark){{{dark_css}}}"#)
145 .into_boxed_str()
146 .into_boxed_bytes(),
147 );
148 DARK_HIGHLIGHT_CSS_HASH
149 .set(build_asset_hash(dark_css))
150 .unwrap();
151
152 let static_favicon = |content: &'static [u8]| {
153 move || async move {
154 let mut resp = Response::new(Body::from(content));
155 resp.headers_mut().insert(
156 http::header::CONTENT_TYPE,
157 HeaderValue::from_static("image/x-icon"),
158 );
159 resp
160 }
161 };
162
163 let static_css = |content: &'static [u8]| {
164 move || async move {
165 let mut resp = Response::new(Body::from(content));
166 resp.headers_mut().insert(
167 http::header::CONTENT_TYPE,
168 HeaderValue::from_static("text/css"),
169 );
170 resp
171 }
172 };
173
174 let app = Router::new()
175 .route("/", get(methods::index::handle))
176 .route(
177 &format!("/style-{}.css", *GLOBAL_CSS_HASH),
178 get(static_css(GLOBAL_CSS)),
179 )
180 .route(
181 &format!("/highlight-{}.css", HIGHLIGHT_CSS_HASH.get().unwrap()),
182 get(static_css(css)),
183 )
184 .route(
185 &format!(
186 "/highlight-dark-{}.css",
187 DARK_HIGHLIGHT_CSS_HASH.get().unwrap()
188 ),
189 get(static_css(dark_css)),
190 )
191 .route(
192 "/favicon.ico",
193 get(static_favicon(include_bytes!("../statics/favicon.ico"))),
194 )
195 .fallback(methods::repo::service)
196 .layer(layer_fn(LoggingMiddleware))
197 .layer(Extension(Arc::new(Git::new(syntax_set))))
198 .layer(Extension(db))
199 .layer(Extension(Arc::new(args.scan_path)))
200 .layer(CorsLayer::new());
201
202 let listener = TcpListener::bind(&args.bind_address).await?;
203 let server = axum::serve(
204 listener,
205 app.into_make_service_with_connect_info::<SocketAddr>(),
206 )
207 .into_future();
208
209 tokio::select! {
210 res = server => res.context("failed to run server"),
211 res = indexer_wakeup_task => res.context("failed to run indexer"),
212 _ = tokio::signal::ctrl_c() => {
213 info!("Received ctrl-c, shutting down");
214 Ok(())
215 }
216 }
217}
218
219fn open_db(args: &Args) -> Result<Arc<rocksdb::DB>, anyhow::Error> {
220 loop {
221 let mut db_options = Options::default();
222 db_options.create_missing_column_families(true);
223 db_options.create_if_missing(true);
224
225 let mut commit_family_options = Options::default();
226 commit_family_options.set_prefix_extractor(SliceTransform::create(
227 "commit_prefix",
228 |input| input.split(|&c| c == b'\0').next().unwrap_or(input),
229 None,
230 ));
231
232 let mut tag_family_options = Options::default();
233 tag_family_options.set_prefix_extractor(SliceTransform::create_fixed_prefix(
234 std::mem::size_of::<u64>(),
235 )); // repository id prefix
236
237 let db = rocksdb::DB::open_cf_with_opts(
238 &db_options,
239 &args.db_store,
240 vec![
241 (COMMIT_FAMILY, commit_family_options),
242 (REPOSITORY_FAMILY, Options::default()),
243 (TAG_FAMILY, tag_family_options),
244 (REFERENCE_FAMILY, Options::default()),
245 (COMMIT_COUNT_FAMILY, Options::default()),
246 ],
247 )?;
248
249 let needs_schema_regen = match db.get("schema_version")? {
250 Some(v) if v.as_bytes() != SCHEMA_VERSION.as_bytes() => Some(Some(v)),
251 Some(_) => None,
252 None => {
253 db.put("schema_version", SCHEMA_VERSION)?;
254 None
255 }
256 };
257
258 if let Some(version) = needs_schema_regen {
259 let old_version = version
260 .as_deref()
261 .map_or(Cow::Borrowed("unknown"), String::from_utf8_lossy);
262
263 warn!("Clearing outdated database ({old_version} != {SCHEMA_VERSION})");
264
265 drop(db);
266 rocksdb::DB::destroy(&Options::default(), &args.db_store)?;
267 } else {
268 break Ok(Arc::new(db));
269 }
270 }
271}
272
273async fn run_indexer(
274 db: Arc<rocksdb::DB>,
275 scan_path: PathBuf,
276 refresh_interval: RefreshInterval,
277) -> Result<(), tokio::task::JoinError> {
278 let (indexer_wakeup_send, mut indexer_wakeup_recv) = mpsc::channel(10);
279
280 std::thread::spawn(move || loop {
281 info!("Running periodic index");
282 crate::database::indexer::run(&scan_path, &db);
283 info!("Finished periodic index");
284
285 if indexer_wakeup_recv.blocking_recv().is_none() {
286 break;
287 }
288 });
289
290 tokio::spawn({
291 let mut sighup = signal(SignalKind::hangup()).expect("could not subscribe to sighup");
292 let build_sleeper = move || async move {
293 match refresh_interval {
294 RefreshInterval::Never => futures::future::pending().await,
295 RefreshInterval::Duration(v) => tokio::time::sleep(v).await,
296 };
297 };
298
299 async move {
300 loop {
301 tokio::select! {
302 _ = sighup.recv() => {},
303 () = build_sleeper() => {},
304 }
305
306 if indexer_wakeup_send.send(()).await.is_err() {
307 error!("Indexing thread has died and is no longer accepting wakeup messages");
308 }
309 }
310 }
311 })
312 .await
313}
314
315#[must_use]
316pub fn build_asset_hash(v: &[u8]) -> Box<str> {
317 let mut hasher = sha2::Sha256::default();
318 hasher.update(v);
319 let mut out = hex::encode(hasher.finalize_fixed());
320 out.truncate(10);
321 Box::from(out)
322}
323
324pub struct TemplateResponse<T> {
325 template: T,
326}
327
328impl<T: Template> IntoResponse for TemplateResponse<T> {
329 #[instrument(skip_all)]
330 fn into_response(self) -> Response {
331 match self.template.render() {
332 Ok(body) => {
333 let headers = [(
334 http::header::CONTENT_TYPE,
335 HeaderValue::from_static(T::MIME_TYPE),
336 )];
337
338 (headers, body).into_response()
339 }
340 Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
341 }
342 }
343}
344
345pub fn into_response<T: Template>(template: T) -> impl IntoResponse {
346 TemplateResponse { template }
347}
348
349pub enum ResponseEither<A, B> {
350 Left(A),
351 Right(B),
352}
353
354impl<A: IntoResponse, B: IntoResponse> IntoResponse for ResponseEither<A, B> {
355 fn into_response(self) -> Response {
356 match self {
357 Self::Left(a) => a.into_response(),
358 Self::Right(b) => b.into_response(),
359 }
360 }
361}