use std::{collections::HashMap, net::SocketAddr}; use async_stream::stream; use axum::{ body::Body, extract::{ConnectInfo, Query}, http::{header, HeaderMap, StatusCode}, response::IntoResponse, }; use bytes::Bytes; use html_escape::{encode_text, encode_unquoted_attribute}; use crate::engines::{ self, Engine, EngineProgressUpdate, ProgressUpdate, ProgressUpdateData, Response, SearchQuery, }; fn render_beginning_of_html(query: &str) -> String { format!( r#" {} - metasearch
"#, encode_text(query), encode_unquoted_attribute(query) ) } fn render_end_of_html() -> String { r#"
"#.to_string() } fn render_engine_list(engines: &[engines::Engine]) -> String { let mut html = String::new(); for engine in engines { html.push_str(&format!( r#"{engine}"#, engine = encode_text(&engine.id()) )); } format!(r#"
{html}
"#) } fn render_search_result(result: &engines::SearchResult) -> String { format!( r#"
{url}

{title}

{desc}

{engines_html}
"#, url_attr = encode_unquoted_attribute(&result.url), url = encode_text(&result.url), title = encode_text(&result.title), desc = encode_text(&result.description), engines_html = render_engine_list(&result.engines.iter().copied().collect::>()) ) } fn render_featured_snippet(featured_snippet: &engines::FeaturedSnippet) -> String { format!( r#"

{desc}

{url}

{title}

 {engines_html}
"#, desc = encode_text(&featured_snippet.description), url_attr = encode_unquoted_attribute(&featured_snippet.url), url = encode_text(&featured_snippet.url), title = encode_text(&featured_snippet.title), engines_html = render_engine_list(&[featured_snippet.engine]) ) } fn render_results(response: Response) -> String { let mut html = String::new(); if let Some(infobox) = response.infobox { html.push_str(&format!( r#"
{infobox_html}{engines_html}
"#, infobox_html = &infobox.html, engines_html = render_engine_list(&[infobox.engine]) )); } if let Some(answer) = response.answer { html.push_str(&format!( r#"
{answer_html}{engines_html}
"#, answer_html = &answer.html, engines_html = render_engine_list(&[answer.engine]) )); } if let Some(featured_snippet) = response.featured_snippet { html.push_str(&render_featured_snippet(&featured_snippet)); } for result in &response.search_results { html.push_str(&render_search_result(result)); } html } fn render_engine_progress_update( engine: Engine, progress_update: &EngineProgressUpdate, time_ms: u64, ) -> String { let message = match progress_update { EngineProgressUpdate::Requesting => "requesting", EngineProgressUpdate::Downloading => "downloading", EngineProgressUpdate::Parsing => "parsing", EngineProgressUpdate::Done => "done", }; format!(r#"{time_ms:>4}ms {engine} {message}"#) } pub async fn route( Query(params): Query>, headers: HeaderMap, ConnectInfo(addr): ConnectInfo, ) -> impl IntoResponse { let query = params .get("q") .cloned() .unwrap_or_default() .trim() .replace('\n', " "); if query.is_empty() { // redirect to index return ( StatusCode::FOUND, [ (header::LOCATION, "/"), (header::CONTENT_TYPE, "text/html; charset=utf-8"), ], Body::from("No query provided, click here to go back to index"), ); } let query = SearchQuery { query, request_headers: headers .clone() .into_iter() .map(|(k, v)| { ( k.map(|k| k.to_string()).unwrap_or_default(), v.to_str().unwrap_or_default().to_string(), ) }) .collect(), ip: headers // this could be exploited under some setups, but the ip is only used for the // "what is my ip" answer so it doesn't really matter .get("x-forwarded-for") .map(|ip| ip.to_str().unwrap_or_default().to_string()) .unwrap_or_else(|| addr.ip().to_string()), }; let s = stream! { type R = Result; // the html is sent in three chunks (technically more if you count progress updates): // 1) the beginning of the html, including the search bar // 1.5) the progress updates // 2) the results // 3) the post-search infobox (usually not sent) + the end of the html let first_part = render_beginning_of_html(&query); // second part is in the loop let mut third_part = String::new(); yield R::Ok(Bytes::from(first_part)); let (progress_tx, mut progress_rx) = tokio::sync::mpsc::unbounded_channel(); let search_future = tokio::spawn(async move { engines::search(query, progress_tx).await }); while let Some(progress_update) = progress_rx.recv().await { match progress_update.data { ProgressUpdateData::Engine { engine, update } => { let progress_html = format!( r#"

{}

"#, render_engine_progress_update(engine, &update, progress_update.time_ms) ); yield R::Ok(Bytes::from(progress_html)); }, ProgressUpdateData::Response(results) => { let mut second_part = String::new(); second_part.push_str(""); // close progress-updates second_part.push_str(""); second_part.push_str(&render_results(results)); yield Ok(Bytes::from(second_part)); }, ProgressUpdateData::PostSearchInfobox(infobox) => { third_part.push_str(&format!( r#"
{infobox_html}{engines_html}
"#, infobox_html = &infobox.html, engines_html = render_engine_list(&[infobox.engine]) )); } } } if let Err(e) = search_future.await? { let error_html = format!( r#"

Error: {}

"#, encode_text(&e.to_string()) ); yield R::Ok(Bytes::from(error_html)); return; }; third_part.push_str(&render_end_of_html()); yield Ok(Bytes::from(third_part)); }; let stream = Body::from_stream(s); ( StatusCode::OK, [ (header::CONTENT_TYPE, "text/html; charset=utf-8"), (header::TRANSFER_ENCODING, "chunked"), ], stream, ) }