jam-program-blob 0.1.26

use std::{
	ffi::OsString,
	fs::File,
	io::{BufRead, BufReader, Seek, SeekFrom, Write},
	path::{Path, PathBuf},
	process::ExitCode,
};

use clap::Parser;
use jam_program_blob_common::{
	read_metadata, write_metadata, ConventionalMetadata, CrateInfo, ProgramBlob,
};

#[derive(Parser)]
#[clap(about, arg_required_else_help = true)]
struct Args {
	/// File operation.
	#[command(subcommand)]
	command: Commands,
}

#[derive(clap::Subcommand, Clone, Debug)]
enum Commands {
	/// Show metadata.
	ShowMeta {
		/// JAM/CoreVM binary.
		#[clap(value_name = "FILE")]
		path: PathBuf,
	},
	/// Add/modify metadata.
	SetMeta {
		/// Program name.
		#[clap(long)]
		name: Option<String>,
		/// Program version.
		#[clap(long)]
		version: Option<String>,
		/// License.
		#[clap(long)]
		license: Option<String>,
		/// Authors (repeat the argument to specify more than one).
		#[clap(long = "author")]
		authors: Option<Vec<String>>,
		/// JAM/CoreVM binary.
		#[clap(value_name = "FILE")]
		path: PathBuf,
	},
	/// Get PVM binary from CoreVM binary.
	GetPvm {
		/// CoreVM binary.
		#[clap(value_name = "FILE")]
		path: PathBuf,
		/// Output file (PVM binary).
		#[clap()]
		output_file: PathBuf,
	},
	/// Builds a canonical null authorizer blob.
	BuildNullAuthorizer {
		/// The path to save the authorizer blob to.
		output: PathBuf,
	},
}

fn main() -> ExitCode {
	do_main().unwrap_or_else(|e| {
		eprintln!("{e}");
		ExitCode::FAILURE
	})
}

const NULL_AUTHORIZER_EXPECTED_HASH: &str =
	"f8d86b97d65319a078e5840f1614c296a5254217794dcc910e72ca174e3c2e86";

fn build_null_authorizer() -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> {
	let assembly = r"
		pub @is_authorized_ext:
		// Return empty auth output.
		//
		// Since the same registers are used for the inputs and the outputs
		// we need to clear these to zero.
		a0 = 0
		a1 = 0
		ret
	";

	// The assembler is currently not (yet) a public API, so we need to go through PolkaVM's
	// private `polkavm-common` crate to access it.
	let polkavm_blob = polkavm_common::assembler::assemble(assembly)
		.map_err(|err| format!("PolkaVM assembler failed: {err}"))?;
	let parts = polkavm_common::program::ProgramParts::from_bytes(polkavm_blob.into())
		.map_err(|err| format!("failed to split PolkaVM blob into parts: {err}"))?;

	let metadata = ConventionalMetadata::Info(CrateInfo {
		name: "NULL Authorizer".into(),
		version: "1".into(),
		license: "CC0".into(),
		authors: vec![],
	});
	let metadata = codec::Encode::encode(&metadata);

	ProgramBlob::from_pvm(&parts, metadata.into())
		.to_vec()
		.map_err(|err| format!("failed to serialize JAM blob: {err}").into())
}

fn do_main() -> Result<ExitCode, Box<dyn std::error::Error + Send + Sync>> {
	let args = Args::parse();
	match args.command {
		Commands::ShowMeta { path } => {
			let file = File::open(&path).map_err(|e| format!("Failed to open {path:?}: {e}"))?;
			let mut reader = BufReader::with_capacity(MAX_METADATA_LEN, file);
			reader.fill_buf()?;
			let ConventionalMetadata::Info(CrateInfo { name, version, license, authors }) =
				read_metadata(&mut reader.buffer())
					.map_err(|e| format!("Failed to read metadata: {e}"))?;
			println!("Name: {name}");
			println!("Version: {version}");
			println!("License: {license}");
			println!("Authors: {}", authors.join(", "));
		},
		Commands::SetMeta { name, version, license, authors, path } => {
			let input_file =
				File::open(&path).map_err(|e| format!("Failed to open {path:?}: {e}"))?;
			let mut reader = BufReader::with_capacity(MAX_METADATA_LEN, input_file);
			reader.fill_buf()?;
			let mut slice = reader.buffer();
			let old_len = slice.len();
			let (metadata, metadata_len) = match read_metadata(&mut slice) {
				Ok(ConventionalMetadata::Info(mut old_metadata)) => {
					let new_len = slice.len();
					let metadata_len = old_len - new_len;
					if let Some(name) = name {
						old_metadata.name = name;
					}
					if let Some(version) = version {
						old_metadata.version = version;
					}
					if let Some(license) = license {
						old_metadata.license = license;
					}
					if let Some(authors) = authors {
						old_metadata.authors = authors;
					}
					(ConventionalMetadata::Info(old_metadata), metadata_len)
				},
				Err(_) => {
					let metadata = ConventionalMetadata::Info(CrateInfo {
						name: name.unwrap_or_default(),
						version: version.unwrap_or_default(),
						license: license.unwrap_or_default(),
						authors: authors.unwrap_or_default(),
					});
					(metadata, 0)
				},
			};
			let mut input_file = reader.into_inner();
			input_file.seek(SeekFrom::Start(metadata_len as u64))?;
			let tmp_file_name = get_tmp_file_name(&path).expect("File name exists");
			let mut tmp_file = File::create(&tmp_file_name)
				.map_err(|e| format!("Failed to create {tmp_file_name:?}: {e}"))?;
			let mut metadata_bytes = Vec::new();
			write_metadata(&metadata, &mut metadata_bytes)?;
			tmp_file.write_all(&metadata_bytes)?;
			std::io::copy(&mut input_file, &mut tmp_file)?;
			std::fs::rename(&tmp_file_name, path)?;
		},
		Commands::GetPvm { path, output_file } => {
			if ProgramBlob::from_bytes(
				&std::fs::read(&path).map_err(|e| format!("Failed to read {path:?}: {e}"))?,
			)
			.is_some()
			{
				return Err("Can't get PVM blob from JAM binary".into());
			}
			let file = File::open(&path).map_err(|e| format!("Failed to open {path:?}: {e}"))?;
			let mut reader = BufReader::with_capacity(MAX_METADATA_LEN, file);
			reader.fill_buf()?;
			let slice = &mut reader.buffer();
			let old_len = slice.len();
			let _meta =
				read_metadata(slice).map_err(|e| format!("Failed to read metadata: {e}"))?;
			let new_len = slice.len();
			let metadata_len = old_len - new_len;
			let mut file = reader.into_inner();
			file.seek(SeekFrom::Start(metadata_len as u64))?;
			let mut out_file = File::create(&output_file)
				.map_err(|e| format!("Failed to create {output_file:?}: {e}"))?;
			std::io::copy(&mut file, &mut out_file)?;
		},
		Commands::BuildNullAuthorizer { output } => {
			let jam_blob = build_null_authorizer()?;

			let hash = blake2b_simd::Params::new().hash_length(32).hash(&jam_blob);
			eprintln!("Built canonical NULL authorizer; hash: 0x{}", hash.to_hex());

			if hash.to_hex().to_string() != NULL_AUTHORIZER_EXPECTED_HASH {
				eprintln!("ERROR: Hash mismatch: expected hash 0x{NULL_AUTHORIZER_EXPECTED_HASH}. This should never happen; please report this!");
				return Ok(ExitCode::FAILURE);
			}

			std::fs::write(&output, jam_blob)
				.map_err(|err| format!("failed to write to {}: {err}", output.display()))?;

			eprintln!("Canonical NULL authorizer successfully written to '{}'.", output.display());
		},
	}
	Ok(ExitCode::SUCCESS)
}

fn get_tmp_file_name(path: &Path) -> Option<PathBuf> {
	let mut buf = PathBuf::new();
	if let Some(dir) = path.parent() {
		buf.push(dir);
	}
	let mut file_name = OsString::new();
	file_name.push(".");
	file_name.push(path.file_name()?);
	file_name.push(".tmp");
	Some(file_name.into())
}

const MAX_METADATA_LEN: usize = 2 * 4096;

#[test]
fn test_build_null_authorizer() {
	let jam_blob = build_null_authorizer().unwrap();
	let hash = blake2b_simd::Params::new().hash_length(32).hash(&jam_blob);
	assert_eq!(hash.to_hex().to_string(), NULL_AUTHORIZER_EXPECTED_HASH);
}