1 files changed, 21 insertions, 14 deletions

diff --git a/src/epub.rs b/src/epub.rs
index e71e3ca..e993c31 100644
--- a/src/epub.rs
+++ b/src/epub.rs
@@ -1,5 +1,5 @@
 use crate::models::FileEntry;
-use anyhow::Result;
+use anyhow::{Context, Result};
 use reqwest::Client;
 use std::{io::Write, path::Path};
 use tokio::{
@@ -9,15 +9,7 @@ use tokio::{
 use zip::{CompressionMethod, ZipWriter, write::FileOptions};
 
 /// Creates and writes container.xml.
-pub async fn write_container_xml(dest_root: &Path, opf_full_path: &str) -> Result<()> {
-    // Create destination directory.
-    let dest_dir = dest_root.join("META-INF");
-    fs::create_dir_all(&dest_dir).await?;
-
-    // Create distination file.
-    let dest_path = dest_dir.join("container.xml");
-    let mut file = File::create(dest_path).await?;
-
+fn write_container_xml_to_zip(zip: &mut ZipWriter<std::fs::File>, opf_full_path: &str) -> Result<()> {
     // Prepare file contents.
     let contents = format!(
         r#"<?xml version="1.0" encoding="UTF-8"?>
@@ -30,7 +22,10 @@ pub async fn write_container_xml(dest_root: &Path, opf_full_path: &str) -> Resul
     );
 
     // Write down the file.
-    file.write_all(contents.as_bytes()).await?;
+    let options: FileOptions<()> =
+        FileOptions::default().compression_method(CompressionMethod::Deflated);
+    zip.start_file("META-INF/container.xml", options)?;
+    zip.write_all(contents.as_bytes())?;
     Ok(())
 }
 
@@ -61,14 +56,26 @@ pub async fn download_all_files(
 }
 
 /// Creates the EPUB archive (creates zip and includes all files in it).
-pub fn create_epub_archive(epub_root: &Path, output_epub: &Path) -> Result<()> {
+pub fn create_epub_archive(
+    epub_root: &Path,
+    output_epub: &Path,
+    file_entries: &[FileEntry],
+) -> Result<()> {
     let out_file = std::fs::File::create(output_epub)?;
     let mut zip = ZipWriter::new(out_file);
 
-    let mimetype_options: FileOptions<()> =
+    // Write mimetype to zip first. It must be uncompressed.
+    let options: FileOptions<()> =
         FileOptions::default().compression_method(CompressionMethod::Stored);
-    zip.start_file("mimetype", mimetype_options)?;
+    zip.start_file("mimetype", options)?;
     zip.write_all(b"application/epub+zip")?;
 
+    // Find the OPF file entry to reference it in container.xml
+    let opf_entry = file_entries
+        .iter()
+        .find(|f| f.filename_ext == ".opf" && f.media_type == "application/oebps-package+xml")
+        .context("No OPF file with the correct MIME type was found.")?;
+    write_container_xml_to_zip(&mut zip, &opf_entry.full_path)?;
+
     Ok(())
 }