From bb1f777078c3ad403b899d1244b9809630a27434 Mon Sep 17 00:00:00 2001
From: Matthieu Gautier <mgautier@kymeria.fr>
Date: Tue, 16 Oct 2018 17:40:32 +0200
Subject: [PATCH] Store the aria2 session and recover from it.

---
 include/downloader.h |  4 +++-
 src/aria2.cpp        | 50 +++++++++++++++++++++++++++++++++++++++++++-
 src/aria2.h          |  6 ++++--
 src/downloader.cpp   |  5 +++++
 4 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/include/downloader.h b/include/downloader.h
index 066a0483c..a70f6af72 100644
--- a/include/downloader.h
+++ b/include/downloader.h
@@ -43,7 +43,9 @@ class Downloader
 {
  public:
   Downloader();
-  ~Downloader();
+  virtual ~Downloader();
+
+  void close();
 
   /**
    * Download a content.
diff --git a/src/aria2.cpp b/src/aria2.cpp
index 2a5fa4661..0079327bf 100644
--- a/src/aria2.cpp
+++ b/src/aria2.cpp
@@ -22,6 +22,9 @@ Aria2::Aria2():
 
   std::string rpc_port = "--rpc-listen-port=" + std::to_string(m_port);
   std::string download_dir = "--dir=" + getDataDirectory();
+  std::string session_file = appendToDirectory(getDataDirectory(), "kiwix.session");
+  std::string session = "--save-session=" + session_file;
+  std::string inputFile = "--input-file=" + session_file;
 //  std::string log_dir = "--log=\"" + logDir + "\"";
 #ifdef _WIN32
   int pid = GetCurrentProcessId();
@@ -37,7 +40,12 @@ Aria2::Aria2():
   callCmd.push_back(rpc_secret.c_str());
   callCmd.push_back(rpc_port.c_str());
   callCmd.push_back(download_dir.c_str());
+  if (fileExists(session_file)) {
+    callCmd.push_back(inputFile.c_str());
+  }
+  callCmd.push_back(session.c_str());
 //  callCmd.push_back(log_dir.c_str());
+  callCmd.push_back("--auto-save-interval=10");
   callCmd.push_back(stop_with_pid.c_str());
   callCmd.push_back("--allow-overwrite=true");
   callCmd.push_back("--dht-entry-point=router.bittorrent.com:6881");
@@ -66,7 +74,13 @@ Aria2::Aria2():
 Aria2::~Aria2()
 {
   curl_easy_cleanup(mp_curl);
-};
+}
+
+void Aria2::close()
+{
+  saveSession();
+  shutdown();
+}
 
 size_t write_callback_to_iss(char* ptr, size_t size, size_t nmemb, void* userdata)
 {
@@ -128,4 +142,38 @@ std::string Aria2::tellStatus(const std::string& gid, const std::vector<std::str
   return doRequest(methodCall);
 }
 
+std::vector<std::string> Aria2::tellActive()
+{
+  MethodCall methodCall("aria2.tellActive", m_secret);
+  auto statusArray = methodCall.newParamValue().getArray();
+  statusArray.addValue().set(std::string("gid"));
+  statusArray.addValue().set(std::string("following"));
+  auto responseContent = doRequest(methodCall);
+  MethodResponse response(responseContent);
+  std::vector<std::string> activeGID;
+  int index = 0;
+  while(true) {
+    try {
+      auto structNode = response.getParamValue(0).getArray().getValue(index++).getStruct();
+      auto gidNode = structNode.getMember("gid");
+      activeGID.push_back(gidNode.getValue().getAsS());
+    } catch (InvalidRPCNode& e) { break; }
+  }
+  return activeGID;
+}
+
+void Aria2::saveSession()
+{
+  MethodCall methodCall("aria2.saveSession", m_secret);
+  doRequest(methodCall);
+  std::cout << "session saved" << std::endl;
+}
+
+void Aria2::shutdown()
+{
+  MethodCall methodCall("aria2.shutdown", m_secret);
+  doRequest(methodCall);
+}
+
+
 } // end namespace kiwix
diff --git a/src/aria2.h b/src/aria2.h
index d94a62508..2c30174f6 100644
--- a/src/aria2.h
+++ b/src/aria2.h
@@ -31,11 +31,13 @@ class Aria2
   public:
     Aria2();
     virtual ~Aria2();
+    void close();
 
     std::string addUri(const std::vector<std::string>& uri);
     std::string tellStatus(const std::string& gid, const std::vector<std::string>& statusKey);
-
-
+    std::vector<std::string> tellActive();
+    void saveSession();
+    void shutdown();
 };
 
 }; //end namespace kiwix
diff --git a/src/downloader.cpp b/src/downloader.cpp
index 9edb29c58..20112e71a 100644
--- a/src/downloader.cpp
+++ b/src/downloader.cpp
@@ -45,6 +45,11 @@ Downloader::~Downloader()
 {
 }
 
+void Downloader::close()
+{
+  mp_aria->close();
+}
+
 pugi::xml_node find_member_in_struct(pugi::xml_node struct_node, std::string member_name) {
   for(auto member=struct_node.first_child(); member; member=member.next_sibling()) {
     std::string _member_name = member.child("name").text().get();