2 files changed, 22 insertions, 22 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(())
 }
diff --git a/src/main.rs b/src/main.rs
index db1fdd4..8c5c712 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,7 +5,7 @@ mod models;
 use std::collections::HashMap;
 use std::path::Path;
 
-use crate::epub::{download_all_files, write_container_xml};
+use crate::epub::{download_all_files,};
 use crate::http_client::build_authenticated_client;
 use crate::models::{Chapter, EpubResponse, FileEntry, Paginated, SpineItem, TocNode};
 use anyhow::{Context, Result, ensure};
@@ -112,16 +112,9 @@ async fn main() -> Result<()> {
     let spine_items: Vec<SpineItem> = fetch_all_pages(&client, epub_data.spine.clone()).await?;
     let toc_vec: Vec<TocNode> = fetch_direct_array(&client, &epub_data.table_of_contents).await?;
 
-    // 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.")?;
-
     let dest_root = format!("Books/{}/epub_root", args.bookid);
     let dest_root = Path::new(&dest_root);
     download_all_files(&client, &file_entries, dest_root).await?;
-    write_container_xml(dest_root, &opf_entry.full_path).await?;
 
     // Sanity check: Every entry in spine exists in chapters.
     let chapters: HashMap<String, Chapter> =