geolite 0.0.18

use std::net::SocketAddr;
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use tiny_http::{Header, Method, Response, Server, StatusCode};

use crate::admin_level::{level, level_error};
use crate::admin_level_hierarchy::tantivy_index;
use crate::database::store::base_files;

mod etag;
mod openapi;
mod status;

const INDEX_HTML: &str = include_str!("static/index.html");

static SHUTDOWN: AtomicBool = AtomicBool::new(false);

// SIGINT (ctrl-c) and SIGTERM (`docker stop`); these posix values are stable across linux/macos.
const SIGINT: std::os::raw::c_int = 2;
const SIGTERM: std::os::raw::c_int = 15;

// libc's signal(2), declared directly so no extra crate is pulled in for a one-flag handler.
unsafe extern "C" {
  fn signal(signum: std::os::raw::c_int, handler: extern "C" fn(std::os::raw::c_int)) -> usize;
}

extern "C" fn on_shutdown_signal(_signum: std::os::raw::c_int) {
  SHUTDOWN.store(true, Ordering::Relaxed);
}

pub struct bound_server {
  server: Arc<Server>,
  addr: SocketAddr,
}

impl bound_server {
  pub fn addr(&self) -> SocketAddr {
    self.addr
  }
}

pub fn bind(addr: &str) -> bound_server {
  let server = Server::http(addr).expect("failed to start http server");
  let addr = server
    .server_addr()
    .to_ip()
    .expect("the http server listens on a tcp socket");
  bound_server {
    server: Arc::new(server),
    addr,
  }
}

pub struct served_base {
  files: base_files,
  index: Option<Arc<tantivy_index>>,
}

struct worker_base<'a> {
  files: &'a base_files,
  label: String,
  conn: rusqlite::Connection,
  index: Option<&'a tantivy_index>,
}

pub fn serve(bound: bound_server, bases: Vec<base_files>, threads: u8) {
  let addr = bound.addr();
  let server = bound.server;

  // degraded boot: a base without its index answers coordinates only (reported via /status)
  // instead of refusing to start
  let boosts = crate::presets::DEFAULT.index_addresses.boosts;
  let served: Vec<served_base> = bases
    .into_iter()
    .map(|files| {
      let index = crate::admin_level_hierarchy::search_index::load(Path::new(&files.index), boosts)
        .map(Arc::new);
      if index.is_none() {
        eprintln!(
          "\x1b[1;33mwarn\x1b[0m: tantivy index not found at {} — text_to_address disabled for {}; run `geolite exec index-addresses {}` to enable",
          files.index,
          files.label(),
          files.iso
        );
      }
      served_base { files, index }
    })
    .collect();
  let served = Arc::new(served);
  crate::admin_level::country_index::country_index::shared();

  // SAFETY: on_shutdown_signal only performs an atomic store, which is async-signal-safe.
  unsafe {
    signal(SIGINT, on_shutdown_signal);
    signal(SIGTERM, on_shutdown_signal);
  }

  // the e2e harness reads this line to learn the port a `--port 0` server took: keep its shape
  println!(
    "\x1b[1;32mlistening\x1b[0m on {addr}  bases: {}",
    served
      .iter()
      .map(|base| base.files.sqlite.as_str())
      .collect::<Vec<_>>()
      .join(", ")
  );

  let handles: Vec<_> = (0..threads)
    .map(|_| {
      let server = server.clone();
      let served = served.clone();
      std::thread::spawn(move || {
        let mut bases: Vec<worker_base> = served
          .iter()
          .map(|base| worker_base {
            files: &base.files,
            label: base.files.label(),
            conn: crate::database::store::open_base_immutable(&base.files.sqlite),
            index: base.index.as_deref(),
          })
          .collect();
        // recv_timeout instead of the blocking incoming_requests() so each worker rechecks the
        // shutdown flag every poll interval (tiny_http's unblock() only wakes one thread).
        while !SHUTDOWN.load(Ordering::Relaxed) {
          match server.recv_timeout(Duration::from_millis(250)) {
            Ok(Some(request)) => handle(request, &mut bases),
            Ok(None) => {}
            Err(_) => break,
          }
        }
      })
    })
    .collect();

  for h in handles {
    h.join().ok();
  }
  println!("\x1b[1;33mshutdown\x1b[0m: http server stopped");
}

fn parse_countries(raw: &str) -> Result<Vec<String>, String> {
  let isos: Vec<String> = raw
    .split(',')
    .map(str::trim)
    .filter(|iso| !iso.is_empty())
    .map(str::to_string)
    .collect();
  if isos.is_empty() {
    return Err("countries must be a comma-separated list of iso 3166-1 alpha-2 codes".to_string());
  }
  Ok(isos)
}

// the bases named, in the order named, a country named twice asked once: a base is held alone
// while a question runs
fn select<'a, 'b>(
  bases: &'b mut [worker_base<'a>],
  countries: Option<&[String]>,
) -> Result<Vec<&'b mut worker_base<'a>>, String> {
  let Some(countries) = countries else {
    return Ok(bases.iter_mut().collect());
  };
  let mut wanted: Vec<usize> = Vec::with_capacity(countries.len());
  for iso in countries {
    let found = bases
      .iter()
      .position(|base| base.files.iso.eq_ignore_ascii_case(iso))
      .ok_or_else(|| format!("countries: no base for '{iso}'"))?;
    if !wanted.contains(&found) {
      wanted.push(found);
    }
  }
  let mut picked: Vec<(usize, &'b mut worker_base<'a>)> = bases
    .iter_mut()
    .enumerate()
    .filter_map(|(i, base)| wanted.iter().position(|&w| w == i).map(|rank| (rank, base)))
    .collect();
  picked.sort_by_key(|(rank, _)| *rank);
  Ok(picked.into_iter().map(|(_, base)| base).collect())
}

fn handle(request: tiny_http::Request, bases: &mut [worker_base]) {
  let start = Instant::now();
  let url = request.url().to_string();
  let method = request.method().clone();

  let (path, query_string) = match url.split_once('?') {
    Some((p, q)) => (p.to_string(), q.to_string()),
    None => (url.clone(), String::new()),
  };

  let known_route = matches!(
    path.as_str(),
    "/" | "/geocode" | "/address" | "/status" | "/openapi.json" | "/docs"
  );

  if known_route && method == Method::Options {
    let response = force_content_length(
      Response::from_string("")
        .with_status_code(StatusCode(204))
        .with_header(cors_origin())
        .with_header(Header::from_bytes("Access-Control-Allow-Methods", "GET, OPTIONS").unwrap())
        .with_header(cors_allowed_headers()),
    );
    request.respond(response).ok();
    log_access(&method, &url, 204, start.elapsed().as_millis());
    return;
  }

  if known_route && method != Method::Get {
    let status = respond_json(
      request,
      StatusCode(405),
      r#"{"error":"method not allowed"}"#,
    );
    log_access(&method, &url, status, start.elapsed().as_millis());
    return;
  }

  let status = match path.as_str() {
    "/" => respond_html(request, StatusCode(200), INDEX_HTML),
    "/status" => status(request, bases),
    "/openapi.json" => respond_ok(request, openapi::spec_json()),
    "/address" => address(request, bases, &query_string),
    "/docs" => respond_html(request, StatusCode(200), openapi::SWAGGER_HTML),
    "/geocode" => geocode(request, bases, &query_string),
    _ => respond_json(request, StatusCode(404), r#"{"error":"not found"}"#),
  };

  log_access(&method, &url, status, start.elapsed().as_millis());
}

fn status(request: tiny_http::Request, bases: &[worker_base]) -> u16 {
  let report = status::build(
    bases
      .iter()
      .map(|base| (base.label.clone(), &base.conn, base.index)),
  );
  match serde_json::to_string(&report) {
    Ok(body) if report.is_ok => respond_ok(request, &body),
    Ok(body) => respond_json(request, StatusCode(503), &body),
    Err(_) => respond_serialization_error(request),
  }
}

fn address(request: tiny_http::Request, bases: &[worker_base], query_string: &str) -> u16 {
  let Some(raw) = query_param(query_string, "id") else {
    return respond_json(
      request,
      StatusCode(400),
      r#"{"error":"missing query param: id"}"#,
    );
  };
  let Some(id) = crate::address::address_id::parse(&raw) else {
    return respond_json(
      request,
      StatusCode(400),
      r#"{"error":"id must be a country code and the osm elements of a path, e.g. br.r59470.r2931354.w255710390"}"#,
    );
  };
  let located = bases
    .iter()
    .find(|base| base.files.iso.eq_ignore_ascii_case(&id.iso))
    .and_then(|base| crate::address::record(&base.conn, &id).map(|found| (base, found)));
  let Some((base, mut found)) = located else {
    return respond_json(
      request,
      StatusCode(404),
      r#"{"error":"address not found"}"#,
    );
  };
  let etag = etag::of_address(&id);
  if etag::none_match(&request, &etag) {
    return respond_not_modified(request, &etag);
  }
  crate::address::with_shapes(&base.conn, &mut found);
  respond_serialized_with_etag(request, &found, &etag)
}

struct geocode_options {
  friendly_name_format: Option<String>,
  min_quality: Option<f64>,
  bounding: Option<crate::admin_level::geometry::bounding_geometry>,
  last_admin_levels: Option<Vec<level>>,
  countries: Option<Vec<String>>,
}

fn parse_geocode_options(query_string: &str) -> Result<geocode_options, String> {
  Ok(geocode_options {
    friendly_name_format: query_param(query_string, "friendly_name_format")
      .map(|s| crate::address::validate_friendly_name_format(&s))
      .transpose()?,
    min_quality: query_param(query_string, "quality")
      .map(|s| crate::address::parse_min_quality(&s))
      .transpose()?,
    bounding: query_param(query_string, "bounding_wkt")
      .map(|s| crate::admin_level::geometry::parse_bounding_wkt(&s))
      .transpose()?,
    last_admin_levels: query_param(query_string, "last_admin_levels")
      .map(|s| parse_last_admin_levels(&s))
      .transpose()?,
    countries: query_param(query_string, "countries")
      .map(|s| parse_countries(&s))
      .transpose()?,
  })
}

fn geocode(request: tiny_http::Request, bases: &mut [worker_base], query_string: &str) -> u16 {
  let Some(raw) = query_param(query_string, "query") else {
    return respond_json(
      request,
      StatusCode(400),
      r#"{"error":"missing query param: query"}"#,
    );
  };
  let options = match parse_geocode_options(query_string) {
    Ok(options) => options,
    Err(msg) => return respond_bad_request(request, &msg),
  };
  let selected = match select(bases, options.countries.as_deref()) {
    Ok(selected) => selected,
    Err(msg) => return respond_bad_request(request, &msg),
  };
  let input = crate::address::query_input::parse(&raw);
  let no_index = selected.iter().all(|base| base.index.is_none());
  if no_index && input == crate::address::query_input::text {
    // degraded boot: text search needs the tantivy index; coordinate queries still work.
    return respond_json(
      request,
      StatusCode(503),
      r#"{"error":"text_to_address service unavailable"}"#,
    );
  }
  let opts = crate::address::query_opts {
    friendly_name_format: options.friendly_name_format.as_deref(),
    min_quality: options.min_quality,
    bounding: options.bounding,
    last_admin_levels: options.last_admin_levels,
  };
  let etag = etag::of_question(&raw, &opts, options.countries.as_deref());
  if etag::none_match(&request, &etag) {
    return respond_not_modified(request, &etag);
  }
  let mut asked: Vec<crate::address::base> = selected
    .into_iter()
    .map(|base| crate::address::base {
      file: &base.label,
      iso: &base.files.iso,
      conn: &mut base.conn,
      index: base.index,
    })
    .collect();
  let result = match input {
    crate::address::query_input::coordinates {
      latitude,
      longitude,
    } => crate::address::query_by_coordinates_across(&mut asked, latitude, longitude, &opts),
    crate::address::query_input::text => {
      crate::address::query_by_text_across(&mut asked, &raw, &opts)
    }
  };
  respond_serialized_with_etag(request, &result, &etag)
}

fn respond_bad_request(request: tiny_http::Request, msg: &str) -> u16 {
  let body = serde_json::json!({ "error": msg }).to_string();
  respond_json(request, StatusCode(400), &body)
}

fn respond_serialization_error(request: tiny_http::Request) -> u16 {
  respond_json(
    request,
    StatusCode(500),
    r#"{"error":"serialization error"}"#,
  )
}

fn respond_serialized_with_etag(
  request: tiny_http::Request,
  answer: &impl serde::Serialize,
  etag: &str,
) -> u16 {
  match serde_json::to_string(answer) {
    Ok(body) => respond_ok_with_etag(request, &body, etag),
    Err(_) => respond_serialization_error(request),
  }
}

fn cors_origin() -> Header {
  Header::from_bytes("Access-Control-Allow-Origin", "*").unwrap()
}

fn cors_allowed_headers() -> Header {
  Header::from_bytes("Access-Control-Allow-Headers", "Content-Type, If-None-Match").unwrap()
}

fn cors_exposed_headers() -> Header {
  Header::from_bytes("Access-Control-Expose-Headers", "ETag").unwrap()
}

fn json_content_type() -> Header {
  Header::from_bytes("Content-Type", "application/json").unwrap()
}

fn force_content_length<R: std::io::Read>(response: Response<R>) -> Response<R> {
  // tiny_http switches to chunked transfer-encoding (dropping content-length) for
  // bodies >= 32 kib; raising the threshold keeps every response on identity
  // encoding so content-length is always present for http/1.1 keep-alive clients.
  response.with_chunked_threshold(usize::MAX)
}

fn log_access(method: &Method, url: &str, status: u16, ms: u128) {
  let color = match status {
    200..=299 => "1;32",
    300..=399 => "1;36",
    400..=499 => "1;33",
    _ => "1;31",
  };
  println!("\x1b[{color}m{status}\x1b[0m {method} {url} in {ms}ms");
}

fn query_param(query_string: &str, key: &str) -> Option<String> {
  for pair in query_string.split('&') {
    if let Some((k, v)) = pair.split_once('=')
      && url_decode(k) == key
    {
      return Some(url_decode(v));
    }
  }
  None
}

fn hex_digit(b: u8) -> Option<u8> {
  match b {
    b'0'..=b'9' => Some(b - b'0'),
    b'a'..=b'f' => Some(b - b'a' + 10),
    b'A'..=b'F' => Some(b - b'A' + 10),
    _ => None,
  }
}

fn url_decode(s: &str) -> String {
  let bytes = s.as_bytes();
  let mut out: Vec<u8> = Vec::with_capacity(bytes.len());
  let mut i = 0;
  while i < bytes.len() {
    match bytes[i] {
      b'%' if i + 2 < bytes.len() => match (hex_digit(bytes[i + 1]), hex_digit(bytes[i + 2])) {
        (Some(h), Some(l)) => {
          out.push((h << 4) | l);
          i += 3;
        }
        _ => {
          out.push(b'%');
          i += 1;
        }
      },
      b'+' => {
        out.push(b' ');
        i += 1;
      }
      b => {
        out.push(b);
        i += 1;
      }
    }
  }
  String::from_utf8_lossy(&out).into_owned()
}

fn parse_last_admin_levels(s: &str) -> Result<Vec<level>, String> {
  let parts: Vec<&str> = s.split(',').map(str::trim).collect();
  if parts.iter().all(|p| p.is_empty()) {
    return Err("last_admin_levels must be a comma-separated list of levels".to_string());
  }
  parts
    .iter()
    .map(|p| {
      level::parse(p).map_err(|error| match error {
        level_error::not_a_level_number(_) => format!("last_admin_levels: invalid level '{p}'"),
        level_error::outside_the_scale(value) => {
          format!("last_admin_levels: level {value} is not supported")
        }
      })
    })
    .collect()
}

fn respond_json(request: tiny_http::Request, status: StatusCode, body: &str) -> u16 {
  let code = status.0;
  let response = force_content_length(
    Response::from_string(body)
      .with_status_code(status)
      .with_header(json_content_type())
      .with_header(cors_origin()),
  );
  request.respond(response).ok();
  code
}

fn respond_ok(request: tiny_http::Request, body: &str) -> u16 {
  let etag = etag::of_body(body);
  if etag::none_match(&request, &etag) {
    respond_not_modified(request, &etag)
  } else {
    respond_ok_with_etag(request, body, &etag)
  }
}

fn respond_ok_with_etag(request: tiny_http::Request, body: &str, etag: &str) -> u16 {
  let response = force_content_length(
    Response::from_string(body)
      .with_status_code(StatusCode(200))
      .with_header(json_content_type())
      .with_header(etag::header(etag))
      .with_header(cors_origin())
      .with_header(cors_exposed_headers()),
  );
  request.respond(response).ok();
  200
}

fn respond_not_modified(request: tiny_http::Request, etag: &str) -> u16 {
  // from_data, not from_string(""): the latter adds a text/plain content type a 304 must not carry
  let response = force_content_length(
    Response::from_data(Vec::new())
      .with_status_code(StatusCode(304))
      .with_header(etag::header(etag))
      .with_header(cors_origin())
      .with_header(cors_exposed_headers()),
  );
  request.respond(response).ok();
  304
}

fn respond_html(request: tiny_http::Request, status: StatusCode, body: &str) -> u16 {
  let code = status.0;
  let response = force_content_length(
    Response::from_string(body)
      .with_status_code(status)
      .with_header(Header::from_bytes("Content-Type", "text/html; charset=utf-8").unwrap())
      .with_header(Header::from_bytes("Cache-Control", "public, max-age=3600").unwrap())
      .with_header(cors_origin()),
  );
  request.respond(response).ok();
  code
}