qrock 0.2.2

//! Utilities to manage built-in static content.
//!
//! For rocket applications that have static content that can be built-in to
//! the application binary, this module provides a few helpers to serve those
//! static files.
//!
//! Typically the static files reside in the filesystem somewhere and are
//! included into the program using the `include_bytes!` macro.  This module
//! supports serving gzip encoded content as well (sending the raw gzip'd data
//! to the client if it supports it, or decomressing it on the server if it
//! doesn't).
//!
//! # Usage
//! 1. Create a [`StaticFiles`] object.
//! 2. Add built-in files to the `StaticFiles` object.
//! 3. Tell `Rocket` instance to manage the `StaticFiles` object.
//! 4. Tell `Rocket` instance to mount `file_server` function.
//!
//! ## `file_server` example
//! ```
//! use std::path::Path;
//! use rocket::routes;
//! use qrock::sc::bisc::{StaticFiles, Entry, file_server};
//!
//! const RED_CSS: &[u8] = include_bytes!("../../static/red.css");
//!
//! let sfiles = StaticFiles::new()
//!   .add(Path::new("css/red.css"), Entry::new(None, RED_CSS));
//! rocket::build()
//!   .manage(sfiles)
//!   .mount("/static", routes![file_server]);
//! ```
//!
//! ## `BuiltinFiles` example
//! ```
//! use std::path::Path;
//! use qrock::sc::bisc::{StaticFiles, Entry, BuiltinFiles};
//!
//! const RED_CSS: &[u8] = include_bytes!("../../static/red.css");
//!
//! let sfiles = StaticFiles::new()
//!   .add(Path::new("css/red.css"), Entry::new(None, RED_CSS));
//! let bis = BuiltinFiles::new(sfiles.into());
//! rocket::build()
//!   .mount("/static", bis);
//! ```

use std::{
  ffi::OsStr,
  io::Cursor,
  io::Read,
  path::{Path, PathBuf},
  sync::Arc
};

use rocket::{
  get,
  http::{ContentType, Header, Method, Status},
  outcome::IntoOutcome,
  request::Request,
  response::{self, status::Custom, Responder, Response},
  route::{Handler, Outcome},
  Data, Route, State
};

use flate2::bufread::GzDecoder;

use hashbrown::HashMap;

use sha2::{Digest, Sha256};


#[derive(Default)]
#[repr(transparent)]
pub struct StaticFiles(HashMap<&'static Path, Entry>);

impl StaticFiles {
  #[must_use]
  pub fn new() -> Self {
    Self::default()
  }

  #[must_use]
  pub fn add(mut self, p: &'static Path, e: Entry) -> Self {
    self.0.insert(p, e);
    self
  }

  /// # Errors
  /// Returns [`Custom`] status responder that returns `Status::NotFound` if
  /// the path does not exist.
  pub fn get(&self, file: &Path) -> Result<StaticFile, Custom<&'static str>> {
    // Look up raw file data
    let Some(entry) = self.0.get(file) else {
      return Err(Custom(Status::NotFound, "File not found"));
    };

    // Look up file hash (for ETag)
    let content_hash = match self.0.get(file) {
      Some(h) => &h.hash,
      None => {
        return Err(Custom(Status::NotFound, "File (hash) not found"));
      }
    };

    // Get filename extension
    let Some(ext) = file.extension().and_then(OsStr::to_str) else {
      return Err(Custom(Status::BadRequest, "Unknown extension"));
    };

    // Look up content-type given the file's extension
    let Some(content_type) = ContentType::from_extension(ext) else {
      return Err(Custom(Status::BadRequest, "Unknown content type"));
    };

    Ok(StaticFile {
      ctype: content_type,
      hash: content_hash.clone(),
      encoding: entry.encoding,
      payload: entry.data
    })
  }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Encoding {
  Gzip
}

pub struct StaticFile {
  ctype: ContentType,
  hash: String,
  encoding: Option<Encoding>,
  payload: &'static [u8]
}

impl<'r> Responder<'r, 'static> for StaticFile {
  fn respond_to(self, req: &'r Request<'_>) -> response::Result<'static> {
    // Boldly assuming there will never be more than one of these.
    if let Some(etag) = req.headers().get_one("If-None-Match") {
      if self.hash == etag {
        // File has not been modified since it was last sent to client; return
        // 304 Not Modified to explicitly tell client to use the cached copy.
        //println!("Returning NotModified");
        return Response::build().status(Status::NotModified).ok();
      }
    }

    // Prepare response, including an etag so we can do quick hash checks in
    // the future.
    let mut resp = Response::build();
    let etag = Header::new("ETag", self.hash);
    resp.header(self.ctype).header(etag);

    if let Some(encoding) = self.encoding {
      match encoding {
        Encoding::Gzip => {
          // Local file is gzip'd.
          // If the client accepts gzip'd content, then send it as-is.
          // Otherwise decompress it before sending.
          let headers = req.headers();
          if headers.contains("Accept")
            && headers
              .get("Accept-Encoding")
              .any(|e| e.to_lowercase() == "gzip")
          {
            // The client accepts gzip-encoded content.  Send the compressed
            // content.
            resp
              .raw_header("Content-Encoding", "gzip")
              .sized_body(self.payload.len(), Cursor::new(self.payload))
              .ok()
          } else {
            // If we end up here it's because we're holding data that is
            // encoded in a format that the client does not accept,
            // which sucks.  So we have to decompress it on the server and send
            // the uncompressed original over the wire to the client.
            let mut gz = GzDecoder::new(self.payload);
            let mut bytes: Vec<u8> = Vec::new();
            gz.read_to_end(&mut bytes).unwrap();

            resp.sized_body(bytes.len(), Cursor::new(bytes)).ok()
          }
        }
      }
    } else {
      // Local file is stored in its original/raw format, so just send it
      // as-is
      resp
        .sized_body(self.payload.len(), Cursor::new(self.payload))
        .ok()
    }
  }
}


pub struct Entry {
  pub encoding: Option<Encoding>,
  pub data: &'static [u8],
  pub hash: String
}

impl Entry {
  #[must_use]
  pub fn new(encoding: Option<Encoding>, data: &'static [u8]) -> Self {
    let mut hasher = Sha256::new();
    hasher.update(data);
    let result = hasher.finalize();

    // No need to use the full hash
    let hash = hex::encode(&result[..16]);

    Self {
      encoding,
      data,
      hash: format!(r#""{hash}""#)
    }
  }
}


/// Router end-point for loading built-in static content.
///
/// # Errors
/// Returns [`Custom`] status responder that returns `Status::NotFound` if
/// the path does not exist.
#[get("/<file..>")]
#[allow(clippy::needless_pass_by_value)]
pub fn file_server(
  file: PathBuf,
  files: &State<StaticFiles>
) -> Result<StaticFile, Custom<&'static str>> {
  files.get(&file)
}


/// Router end-point for loading built-in static content.
///
/// # Errors
/// Returns [`Custom`] status responder that returns `Status::NotFound` if
/// the path does not exist.
#[get("/<file..>")]
#[allow(clippy::needless_pass_by_value)]
pub fn arc_file_server(
  file: PathBuf,
  files: &State<Arc<StaticFiles>>
) -> Result<StaticFile, Custom<&'static str>> {
  files.get(&file)
}


#[derive(Clone)]
pub struct BuiltinFiles {
  sfiles: Arc<StaticFiles>,
  rank: isize
}

impl BuiltinFiles {
  const DEFAULT_RANK: isize = 10;

  #[must_use]
  pub const fn new(sfiles: Arc<StaticFiles>) -> Self {
    Self {
      sfiles,
      rank: Self::DEFAULT_RANK
    }
  }

  #[must_use]
  pub const fn rank(mut self, rank: isize) -> Self {
    self.rank = rank;
    self
  }
}

impl From<BuiltinFiles> for Vec<Route> {
  fn from(server: BuiltinFiles) -> Self {
    let mut route =
      Route::ranked(server.rank, Method::Get, "/<path..>", server);
    route.name = Some("BuiltinFiles".to_string().into());
    vec![route]
  }
}

#[rocket::async_trait]
impl Handler for BuiltinFiles {
  async fn handle<'r>(
    &self,
    req: &'r Request<'_>,
    data: Data<'r>
  ) -> Outcome<'r> {
    use rocket::http::uri::{fmt::Path, Segments};

    let pth = req
      .segments::<Segments<'_, Path>>(0..)
      .ok()
      .map(|s| s.to_path_buf(false));

    let Some(Ok(pth)) = pth else {
      return Outcome::forward(data, Status::NotFound);
    };

    if let Ok(entry) = self.sfiles.get(&pth) {
      entry.respond_to(req).or_forward((data, Status::NotFound))
    } else {
      Outcome::forward(data, Status::NotFound)
    }
  }
}

// vim: set ft=rust et sw=2 ts=2 sts=2 cinoptions=2 tw=79 :