use std::io::{self, Write, Seek};
use zip::{ZipWriter, write::SimpleFileOptions, CompressionMethod};
pub struct ZipDocumentWriter<W: Write + Seek> {
pub(crate) zip: ZipWriter<W>,
pub(crate) manifest: serde_json::Value,
}
impl<W: Write + Seek> ZipDocumentWriter<W> {
pub fn new(w: W) -> io::Result<Self> {
let zip = ZipWriter::new(w);
let manifest = serde_json::json!({ "version": 1, "entries": [] });
Ok(Self { zip, manifest })
}
pub fn manifest(&self) -> &serde_json::Value {
&self.manifest
}
pub fn manifest_mut(&mut self) -> &mut serde_json::Value {
&mut self.manifest
}
pub fn set_manifest(
&mut self,
manifest: serde_json::Value,
) -> &mut Self {
self.manifest = manifest;
self
}
#[cfg_attr(feature = "dev-tracing", tracing::instrument(skip(self, value), fields(
crate_name = "file",
file_name = %name
)))]
pub fn add_json(
&mut self,
name: &str,
value: &serde_json::Value,
) -> io::Result<()> {
let opts = SimpleFileOptions::default()
.compression_method(CompressionMethod::Deflated);
self.zip.start_file(name, opts)?;
let data = serde_json::to_vec(value).map_err(io::Error::other)?;
if let Some(entries) =
self.manifest.get_mut("entries").and_then(|v| v.as_array_mut())
{
entries.push(serde_json::json!({
"name": name,
"kind": "json",
"logical_len": data.len(),
"compression": "deflate"
}));
}
self.zip.write_all(&data)
}
pub fn add_stored(
&mut self,
name: &str,
bytes: &[u8],
) -> io::Result<()> {
let opts = SimpleFileOptions::default()
.compression_method(CompressionMethod::Stored);
self.zip.start_file(name, opts)?;
if let Some(entries) =
self.manifest.get_mut("entries").and_then(|v| v.as_array_mut())
{
entries.push(serde_json::json!({
"name": name,
"kind": "binary",
"logical_len": bytes.len(),
"compression": "stored"
}));
}
self.zip.write_all(bytes)
}
pub fn add_plugin_state(
&mut self,
plugin_name: &str,
state_data: &[u8],
) -> io::Result<()> {
let plugin_file_path = format!("plugins/{plugin_name}");
let opts = SimpleFileOptions::default()
.compression_method(CompressionMethod::Deflated);
self.zip.start_file(&plugin_file_path, opts)?;
if let Some(entries) =
self.manifest.get_mut("entries").and_then(|v| v.as_array_mut())
{
entries.push(serde_json::json!({
"name": plugin_file_path,
"kind": "plugin_state",
"plugin": plugin_name,
"logical_len": state_data.len(),
"compression": "deflate"
}));
}
self.zip.write_all(state_data)
}
pub fn add_plugin_states<I>(
&mut self,
plugin_states: I,
) -> io::Result<()>
where
I: IntoIterator<Item = (String, Vec<u8>)>,
{
for (plugin_name, state_data) in plugin_states {
self.add_plugin_state(&plugin_name, &state_data)?;
}
Ok(())
}
pub fn add_deflated(
&mut self,
name: &str,
bytes: &[u8],
) -> io::Result<()> {
let opts = SimpleFileOptions::default()
.compression_method(CompressionMethod::Deflated);
self.zip.start_file(name, opts)?;
if let Some(entries) =
self.manifest.get_mut("entries").and_then(|v| v.as_array_mut())
{
entries.push(serde_json::json!({
"name": name,
"kind": "binary",
"logical_len": bytes.len(),
"compression": "deflate"
}));
}
self.zip.write_all(bytes)
}
pub fn finalize(mut self) -> io::Result<W> {
let opts = SimpleFileOptions::default()
.compression_method(CompressionMethod::Deflated);
self.zip.start_file("manifest.json", opts)?;
let data =
serde_json::to_vec(&self.manifest).map_err(io::Error::other)?;
self.zip.write_all(&data)?;
self.zip.finish().map_err(io::Error::other)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn test_plugin_state_export() {
let buffer = Vec::new();
let cursor = Cursor::new(buffer);
let mut writer = ZipDocumentWriter::new(cursor).unwrap();
writer.add_plugin_state("test_plugin", b"test state data").unwrap();
writer
.add_plugin_state("another_plugin", b"another state data")
.unwrap();
let result = writer.finalize().unwrap();
let final_data = result.into_inner();
assert!(!final_data.is_empty());
let cursor = Cursor::new(&final_data);
let mut reader = crate::zipdoc::ZipDocumentReader::new(cursor).unwrap();
let plugins = reader.list_plugins().unwrap();
assert_eq!(plugins.len(), 2);
assert!(plugins.contains(&"test_plugin".to_string()));
assert!(plugins.contains(&"another_plugin".to_string()));
let test_state =
reader.read_plugin_state("test_plugin").unwrap().unwrap();
assert_eq!(test_state, b"test state data");
}
#[test]
fn test_batch_plugin_states() {
let buffer = Vec::new();
let cursor = Cursor::new(buffer);
let mut writer = ZipDocumentWriter::new(cursor).unwrap();
let plugin_states = vec![
("plugin1".to_string(), b"state1".to_vec()),
("plugin2".to_string(), b"state2".to_vec()),
("plugin3".to_string(), b"state3".to_vec()),
];
writer.add_plugin_states(plugin_states.clone()).unwrap();
let result = writer.finalize().unwrap();
let final_data = result.into_inner();
let cursor = Cursor::new(&final_data);
let mut reader = crate::zipdoc::ZipDocumentReader::new(cursor).unwrap();
let all_states = reader.read_all_plugin_states().unwrap();
assert_eq!(all_states.len(), 3);
for (name, expected_data) in plugin_states {
let actual_data = all_states.get(&name).unwrap();
assert_eq!(actual_data, &expected_data);
}
}
}