From 95b1d5b4b87aa0996fd11ba28458435f8de32147 Mon Sep 17 00:00:00 2001
From: Naim A <227396+naim94a@users.noreply.github.com>
Date: Sun, 28 Oct 2018 08:41:52 +0200
Subject: [PATCH] Ability to save compressed database.

---
 Cargo.lock     | 21 +++++++++++++++++++++
 Cargo.toml     |  1 +
 src/main.rs    |  1 +
 src/tracker.rs | 24 ++++++++++++++++++++++++
 4 files changed, 47 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 3d19156..2a20095 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -253,6 +253,24 @@ dependencies = [
  "iovec 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
+[[package]]
+name = "bzip2"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "bzip2-sys 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)",
+ "libc 0.2.43 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
+name = "bzip2-sys"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "cc 1.0.25 (registry+https://github.com/rust-lang/crates.io-index)",
+ "libc 0.2.43 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "cc"
 version = "1.0.25"
@@ -1441,6 +1459,7 @@ dependencies = [
  "actix-web 0.7.13 (registry+https://github.com/rust-lang/crates.io-index)",
  "binascii 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "bincode 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "bzip2 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "clap 2.32.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "fern 0.5.6 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1623,6 +1642,8 @@ dependencies = [
 "checksum build_const 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "39092a32794787acd8525ee150305ff051b0aa6cc2abaf193924f5ab05425f39"
 "checksum byteorder 1.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "90492c5858dd7d2e78691cfb89f90d273a2800fc11d98f60786e5d87e2f83781"
 "checksum bytes 0.4.10 (registry+https://github.com/rust-lang/crates.io-index)" = "0ce55bd354b095246fc34caf4e9e242f5297a7fd938b090cadfea6eee614aa62"
+"checksum bzip2 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)" = "42b7c3cbf0fa9c1b82308d57191728ca0256cb821220f4e2fd410a72ade26e3b"
+"checksum bzip2-sys 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)" = "2c5162604199bbb17690ede847eaa6120a3f33d5ab4dcc8e7c25b16d849ae79b"
 "checksum cc 1.0.25 (registry+https://github.com/rust-lang/crates.io-index)" = "f159dfd43363c4d08055a07703eb7a3406b0dac4d0584d96965a3262db3c9d16"
 "checksum cfg-if 0.1.6 (registry+https://github.com/rust-lang/crates.io-index)" = "082bb9b28e00d3c9d39cc03e64ce4cea0f1bb9b3fde493f0cbc008472d22bdf4"
 "checksum clap 2.32.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b957d88f4b6a63b9d70d5f454ac8011819c6efa7727858f458ab71c756ce2d3e"
diff --git a/Cargo.toml b/Cargo.toml
index 81c245b..6bf4b17 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,3 +19,4 @@ log = "0.4.5"
 fern = "0.5.6"
 num_cpus = "1.8.0"
 serde_json = "1.0.32"
+bzip2 = "0.3.3"
diff --git a/src/main.rs b/src/main.rs
index b785a60..f0fe2cc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,6 +11,7 @@ extern crate toml;
 extern crate fern;
 extern crate num_cpus;
 extern crate serde_json;
+extern crate bzip2;
 
 mod server;
 mod tracker;
diff --git a/src/tracker.rs b/src/tracker.rs
index b449b92..538fd13 100644
--- a/src/tracker.rs
+++ b/src/tracker.rs
@@ -1,6 +1,7 @@
 use std;
 use serde;
 use binascii;
+use serde_json;
 
 use server::Events;
 
@@ -298,6 +299,18 @@ impl TorrentTracker {
     pub (crate) fn get_database(&self) -> std::sync::RwLockReadGuard<std::collections::BTreeMap<InfoHash, TorrentEntry>>{
         self.database.torrent_peers.read().unwrap()
     }
+
+    pub fn save_database<W: std::io::Write>(&self, writer: &mut W) -> serde_json::Result<()> {
+        use bzip2;
+
+        let compressor = bzip2::write::BzEncoder::new(writer, bzip2::Compression::Best);
+
+        let db_lock = self.database.torrent_peers.read().unwrap();
+
+        let db = &*db_lock;
+
+        serde_json::to_writer(compressor, &db)
+    }
 }
 
 #[cfg(test)]
@@ -317,6 +330,17 @@ mod tests {
         is_sync::<TorrentTracker>();
     }
 
+    #[test]
+    fn test_save_db() {
+        let tracker = TorrentTracker::new(TrackerMode::DynamicMode);
+        tracker.add_torrent(&[0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0].into());
+
+        let mut out = Vec::new();
+
+        tracker.save_database(&mut out).unwrap();
+        assert!(out.len() > 0);
+    }
+
     #[test]
     fn test_infohash_de() {
         use serde_json;