use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use helm_schema_json_schema_walk::try_visit_subschemas_mut;
use jsonschema::{Retrieve, Uri};
use referencing::uri;
use serde_json::{Map, Value};
use tracing::instrument;
use url::Url;
use crate::error::{CliError, EngineResult};
use crate::fetch_policy::FetchPolicy;
use crate::load_budget::{LoadBudget, read_to_end_capped};
#[instrument(skip_all)]
pub fn flatten_refs(
schema: &Value,
base_dir: &Path,
fetch_policy: FetchPolicy,
load_budget: LoadBudget,
) -> EngineResult<Value> {
let base_uri = directory_file_uri(base_dir)?;
let retriever = FsHttpRetrieve::new(fetch_policy, load_budget);
flatten_with_retriever(schema, &base_uri, retriever)
}
#[instrument(skip_all)]
pub fn flatten_prepared_refs(schema: &Value, base_dir: &Path) -> EngineResult<Value> {
let base_uri = directory_file_uri(base_dir)?;
flatten_with_retriever(schema, &base_uri, NoExternalRetrieve)
}
#[instrument(skip_all)]
pub fn bundle_refs(
schema: Value,
base_dir: &Path,
fetch_policy: FetchPolicy,
load_budget: LoadBudget,
) -> EngineResult<Value> {
let base_uri = directory_file_uri(base_dir)?;
let retriever = FsHttpRetrieve::new(fetch_policy, load_budget);
bundle_with_retriever(schema, &base_uri, retriever)
}
#[instrument(skip_all)]
pub fn bundle_prepared_refs(schema: Value, base_dir: &Path) -> EngineResult<Value> {
let base_uri = directory_file_uri(base_dir)?;
bundle_with_retriever(schema, &base_uri, NoExternalRetrieve)
}
#[instrument(skip_all)]
pub fn flatten_with_retriever(
schema: &Value,
base_uri: &str,
retriever: impl Retrieve + 'static,
) -> EngineResult<Value> {
let dereferenced = jsonschema::options()
.with_base_uri(base_uri.to_string())
.with_retriever(retriever)
.dereference(schema)?;
Ok(dereferenced)
}
#[instrument(skip_all)]
pub fn bundle_with_retriever(
mut schema: Value,
base_uri: &str,
retriever: impl Retrieve,
) -> EngineResult<Value> {
let root_document_uri = document_uri(&uri::from_str(base_uri)?)?;
let root_base_uri = effective_base_uri(&schema, &root_document_uri)?;
let root_document_uris = BTreeSet::from([
root_document_uri.as_str().to_string(),
root_base_uri.as_str().to_string(),
]);
let existing_definition_names = existing_definition_names(&schema);
let mut state = BundleState::new(retriever, root_document_uris, existing_definition_names);
state.bundle_schema(&mut schema, &root_document_uri)?;
state.insert_definitions(&mut schema)?;
Ok(schema)
}
struct BundleState<R> {
retriever: R,
root_document_uris: BTreeSet<String>,
names_by_target_uri: BTreeMap<String, String>,
definitions: BTreeMap<String, Value>,
existing_definition_names: BTreeSet<String>,
next_definition_id: usize,
}
impl<R: Retrieve> BundleState<R> {
fn new(
retriever: R,
root_document_uris: BTreeSet<String>,
existing_definition_names: BTreeSet<String>,
) -> Self {
Self {
retriever,
root_document_uris,
names_by_target_uri: BTreeMap::new(),
definitions: BTreeMap::new(),
existing_definition_names,
next_definition_id: 1,
}
}
fn bundle_schema(
&mut self,
schema: &mut Value,
current_document_uri: &Uri<String>,
) -> EngineResult<()> {
let current_document_uri = effective_base_uri(schema, current_document_uri)?;
if let Some(reference) = schema_reference(schema) {
let target_uri = uri::resolve_against(¤t_document_uri.borrow(), &reference)?;
if self.should_preserve_reference(&target_uri, ¤t_document_uri)? {
return Ok(());
}
let definition_name = self.definition_name_for_target(&target_uri)?;
*schema = definition_ref(&definition_name);
return Ok(());
}
try_visit_subschemas_mut(schema, &mut |subschema| {
self.bundle_schema(subschema, ¤t_document_uri)
})
}
fn should_preserve_reference(
&self,
target_uri: &Uri<String>,
current_document_uri: &Uri<String>,
) -> EngineResult<bool> {
let target_document_uri = document_uri(target_uri)?;
Ok(self.is_root_document(&target_document_uri)
&& self.is_root_document(current_document_uri))
}
fn definition_name_for_target(&mut self, target_uri: &Uri<String>) -> EngineResult<String> {
let target_key = target_uri.as_str().to_string();
if let Some(name) = self.names_by_target_uri.get(&target_key) {
return Ok(name.clone());
}
let name = self.next_definition_name();
self.names_by_target_uri.insert(target_key, name.clone());
let target_document_uri = document_uri(target_uri)?;
let mut target_schema = self.resolve_target_schema(target_uri, &target_document_uri)?;
self.bundle_schema(&mut target_schema, &target_document_uri)?;
self.definitions.insert(name.clone(), target_schema);
Ok(name)
}
fn resolve_target_schema(
&self,
target_uri: &Uri<String>,
target_document_uri: &Uri<String>,
) -> EngineResult<Value> {
if self.is_root_document(target_document_uri) {
return Err(CliError::RefBundling(format!(
"cannot bundle non-local ref back to root document: {target_uri}"
)));
}
let document = self
.retriever
.retrieve(target_document_uri)
.map_err(|err| {
CliError::RefBundling(format!("retrieve {target_document_uri}: {err}"))
})?;
select_fragment(document, target_uri)
}
fn next_definition_name(&mut self) -> String {
loop {
let name = format!("schema{}", self.next_definition_id);
self.next_definition_id += 1;
if self.existing_definition_names.insert(name.clone()) {
return name;
}
}
}
fn insert_definitions(self, schema: &mut Value) -> EngineResult<()> {
if self.definitions.is_empty() {
return Ok(());
}
let Value::Object(root) = schema else {
return Err(CliError::RefBundling(
"cannot insert bundled definitions into non-object root schema".to_string(),
));
};
let entry = root
.entry("$defs".to_string())
.or_insert_with(|| Value::Object(Map::new()));
let Value::Object(existing) = entry else {
return Err(CliError::RefBundling(
"cannot insert bundled definitions because root $defs is not an object".to_string(),
));
};
for (name, definition) in self.definitions {
existing.insert(name, definition);
}
Ok(())
}
fn is_root_document(&self, document_uri: &Uri<String>) -> bool {
self.root_document_uris.contains(document_uri.as_str())
}
}
struct FsHttpRetrieve {
fetch_policy: FetchPolicy,
load_budget: LoadBudget,
agent: ureq::Agent,
}
impl FsHttpRetrieve {
fn new(fetch_policy: FetchPolicy, load_budget: LoadBudget) -> Self {
Self {
fetch_policy,
load_budget,
agent: ureq::Agent::new_with_defaults(),
}
}
}
impl Retrieve for FsHttpRetrieve {
fn retrieve(
&self,
uri: &Uri<String>,
) -> Result<Value, Box<dyn std::error::Error + Send + Sync>> {
let scheme = uri.scheme().as_str().to_ascii_lowercase();
match scheme.as_str() {
"file" => {
let host = uri
.authority()
.map(|authority| authority.host())
.unwrap_or("");
self.fetch_policy
.validate_file_host(host)
.map_err(|err| format!("$ref to {uri} but {err}"))?;
let path = file_uri_path(uri)?;
let mut file = std::fs::File::open(&path)
.map_err(|error| format!("open {}: {error}", path.display()))?;
let bytes = read_to_end_capped(
&mut file,
self.load_budget.max_schema_document_bytes,
path.display().to_string(),
)
.map_err(|e| e.to_string())?;
let value: Value = serde_json::from_slice(&bytes)
.map_err(|error| format!("parse {}: {error}", path.display()))?;
Ok(value)
}
"http" | "https" => {
let host = uri.authority().map(|authority| authority.host());
self.fetch_policy
.validate_network_host(host)
.map_err(|err| format!("$ref to {uri} but {err}"))?;
let resp = self
.agent
.get(uri.as_str())
.call()
.map_err(|e| format!("fetch {uri}: {e}"))?;
let mut body = resp.into_body();
let mut reader = body.as_reader();
let bytes = read_to_end_capped(
&mut reader,
self.load_budget.max_schema_document_bytes,
uri.as_str().to_string(),
)
.map_err(|e| e.to_string())?;
let value: Value =
serde_json::from_slice(&bytes).map_err(|e| format!("parse {uri}: {e}"))?;
Ok(value)
}
other => Err(format!("unsupported $ref scheme: {other} (uri={uri})").into()),
}
}
}
struct NoExternalRetrieve;
impl Retrieve for NoExternalRetrieve {
fn retrieve(
&self,
uri: &Uri<String>,
) -> Result<Value, Box<dyn std::error::Error + Send + Sync>> {
Err(format!("external $ref remained after input preparation: {uri}").into())
}
}
fn schema_reference(schema: &Value) -> Option<String> {
let Value::Object(object) = schema else {
return None;
};
object
.get("$ref")
.and_then(Value::as_str)
.map(str::to_string)
}
fn document_uri(uri: &Uri<String>) -> EngineResult<Uri<String>> {
let document = uri.strip_fragment().as_str().to_string();
Uri::parse(document)
.map_err(|err| CliError::RefBundling(format!("parse document uri for {uri}: {err:?}")))
}
fn effective_base_uri(
schema: &Value,
current_document_uri: &Uri<String>,
) -> EngineResult<Uri<String>> {
let Some(id) = schema
.as_object()
.and_then(|object| object.get("$id"))
.and_then(Value::as_str)
else {
return Ok(current_document_uri.clone());
};
let resolved = uri::resolve_against(¤t_document_uri.borrow(), id)?;
document_uri(&resolved)
}
fn select_fragment(document: Value, target_uri: &Uri<String>) -> EngineResult<Value> {
let Some(fragment) = target_uri.fragment() else {
return Ok(document);
};
let pointer = fragment.decode().to_string().map_err(|_| {
CliError::RefBundling(format!("decode json pointer fragment for {target_uri}"))
})?;
if pointer.is_empty() {
return Ok(document);
}
if !pointer.starts_with('/') {
return Err(CliError::RefBundling(format!(
"unsupported non-json-pointer fragment in {target_uri}"
)));
}
document.pointer(&pointer).cloned().ok_or_else(|| {
CliError::RefBundling(format!("json pointer {pointer} not found in {target_uri}"))
})
}
fn existing_definition_names(schema: &Value) -> BTreeSet<String> {
schema
.get("$defs")
.and_then(Value::as_object)
.map(|definitions| definitions.keys().cloned().collect())
.unwrap_or_default()
}
fn definition_ref(name: &str) -> Value {
Value::Object(Map::from_iter([(
"$ref".to_string(),
Value::String(format!("#/$defs/{name}")),
)]))
}
fn directory_file_uri(path: &Path) -> EngineResult<String> {
let path = if path.as_os_str().is_empty() {
Path::new(".")
} else {
path
};
let absolute = path.canonicalize().or_else(|_| std::path::absolute(path))?;
let uri = Url::from_directory_path(&absolute)
.map_err(|()| CliError::InvalidFileUriPath { path: absolute })?;
Ok(uri.into())
}
fn file_uri_path(uri: &Uri<String>) -> EngineResult<PathBuf> {
let invalid_uri = || CliError::InvalidFileUri {
uri: uri.as_str().to_string(),
};
Url::parse(uri.as_str())
.map_err(|_| invalid_uri())?
.to_file_path()
.map_err(|()| invalid_uri())
}
#[cfg(test)]
#[path = "tests/flatten.rs"]
mod tests;