use base64::{Engine as B64Engine, engine::general_purpose::STANDARD as b64};
use sha2::{Digest, Sha256};
use std::fmt::{Display, Formatter};
use swc_common::{FileName, FilePathMapping, SourceMap};
use swc_html_ast::Child;
use swc_html_parser::parse_file_as_document;
use swc_html_parser::parser::ParserConfig;
pub struct CspValues {
include_self: bool,
has_wasm: bool,
pub script_src_inline_hashes: Vec<String>,
pub style_src_inline_hashes: Vec<String>,
}
impl Default for CspValues {
fn default() -> Self {
Self::new(false, false)
}
}
impl CspValues {
#[must_use]
pub fn new(include_self: bool, has_wasm: bool) -> Self {
Self {
include_self,
has_wasm,
script_src_inline_hashes: vec![],
style_src_inline_hashes: vec![],
}
}
pub(crate) fn has_any(&self) -> bool {
if self.script_src_inline_hashes.is_empty() {
!self.style_src_inline_hashes.is_empty()
} else {
true
}
}
}
impl Display for CspValues {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
if !self.style_src_inline_hashes.is_empty() {
write!(f, "style-src")?;
if self.include_self {
write!(f, " 'self'")?;
}
for hash in &self.style_src_inline_hashes {
write!(f, " '{hash}'")?;
}
if !self.script_src_inline_hashes.is_empty() {
write!(f, "; ")?;
}
}
if !self.script_src_inline_hashes.is_empty() {
write!(f, "script-src")?;
if self.include_self {
write!(f, " 'self'")?;
}
if self.has_wasm {
write!(f, " 'wasm-unsafe-eval'")?;
}
for hash in &self.script_src_inline_hashes {
write!(f, " '{hash}'")?;
}
}
write!(f, "")
}
}
#[allow(clippy::similar_names)]
pub(crate) fn walk_document(children: Vec<Child>, hashes: &mut CspValues) {
for child in children {
if let Some(element) = child.element() {
if element.tag_name == "script" {
let mut is_inline = true;
for attr in element.attributes {
if attr.name == "src" {
is_inline = false;
}
}
if is_inline && let Some(Child::Text(text)) = element.children.first() {
let mut hasher = Sha256::new();
hasher.update(text.data.as_bytes());
let hash = hasher.finalize().to_vec();
let mut b64_hash = b64.encode(hash);
b64_hash.insert_str(0, "sha256-");
hashes.script_src_inline_hashes.push(b64_hash);
}
} else if element.tag_name == "style" {
if let Some(Child::Text(text)) = element.children.first() {
let mut hasher = Sha256::new();
hasher.update(text.data.as_bytes());
let hash = hasher.finalize().to_vec();
let mut b64_hash = b64.encode(hash);
b64_hash.insert_str(0, "sha256-");
hashes.style_src_inline_hashes.push(b64_hash);
}
} else {
walk_document(element.children, hashes);
}
}
}
}
#[must_use]
pub fn for_html(html: &str, include_self: bool, has_wasm: bool) -> String {
let cm = SourceMap::new(FilePathMapping::empty());
let fm = cm.new_source_file(FileName::Anon.into(), html.to_string());
let mut errors = Vec::new();
let mut csp_values = CspValues::new(include_self, has_wasm);
if let Ok(document) = parse_file_as_document(&fm, ParserConfig::default(), &mut errors) {
walk_document(document.children.clone(), &mut csp_values);
}
format!("{csp_values}")
}