From 5154fcf9da95471bec46b84d7fecfcfaac5c19f7 Mon Sep 17 00:00:00 2001
From: Aiden McClelland <me@drbonez.dev>
Date: Wed, 15 Sep 2021 11:33:23 -0600
Subject: [PATCH] docker image persistence

---
 appmgr/src/bin/embassy-init.rs |  2 ++
 appmgr/src/install/cleanup.rs  |  6 ++++
 appmgr/src/install/mod.rs      | 56 ++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/appmgr/src/bin/embassy-init.rs b/appmgr/src/bin/embassy-init.rs
index 39d9f5ce1..f0351c977 100644
--- a/appmgr/src/bin/embassy-init.rs
+++ b/appmgr/src/bin/embassy-init.rs
@@ -114,6 +114,8 @@ async fn init(cfg_path: Option<&str>) -> Result<(), Error> {
         .invoke(embassy::ErrorKind::Docker)
         .await?;
     log::info!("Mounted Docker Data");
+    embassy::install::load_images().await?;
+    log::info!("Loaded Docker Images");
     embassy::ssh::sync_keys_from_db(&secret_store, "/root/.ssh/authorized_keys").await?;
     log::info!("Synced SSH Keys");
     // todo!("sync wifi");
diff --git a/appmgr/src/install/cleanup.rs b/appmgr/src/install/cleanup.rs
index c8166d246..31d73d9be 100644
--- a/appmgr/src/install/cleanup.rs
+++ b/appmgr/src/install/cleanup.rs
@@ -1,9 +1,11 @@
 use std::collections::HashMap;
+use std::path::Path;
 
 use anyhow::anyhow;
 use bollard::image::ListImagesOptions;
 use patch_db::{DbHandle, PatchDbHandle};
 
+use super::PKG_DOCKER_DIR;
 use crate::context::RpcContext;
 use crate::db::model::{InstalledPackageDataEntry, PackageDataEntry};
 use crate::s9pk::manifest::PackageId;
@@ -83,6 +85,10 @@ pub async fn cleanup(ctx: &RpcContext, id: &PackageId, version: &Version) -> Res
         ctx.docker.remove_image(&image.id, None, None).await
     }))
     .await?;
+    let docker_path = Path::new(PKG_DOCKER_DIR).join(id).join(version.as_str());
+    if tokio::fs::metadata(&docker_path).await.is_ok() {
+        tokio::fs::remove_dir_all(&docker_path).await?;
+    }
     // TODO: delete public dir if not a dependency
 
     Ok(())
diff --git a/appmgr/src/install/mod.rs b/appmgr/src/install/mod.rs
index bbf175702..1e685b556 100644
--- a/appmgr/src/install/mod.rs
+++ b/appmgr/src/install/mod.rs
@@ -16,6 +16,7 @@ use rpc_toolkit::command;
 use tokio::fs::{File, OpenOptions};
 use tokio::io::{AsyncRead, AsyncSeek, AsyncSeekExt};
 use tokio::process::Command;
+use tokio_stream::wrappers::ReadDirStream;
 
 use self::cleanup::cleanup_failed;
 use crate::context::RpcContext;
@@ -670,3 +671,58 @@ pub async fn install_s9pk<R: AsyncRead + AsyncSeek + Unpin>(
 
     Ok(())
 }
+
+pub async fn load_images() -> Result<(), Error> {
+    if tokio::fs::metadata(PKG_DOCKER_DIR).await.is_ok() {
+        ReadDirStream::new(tokio::fs::read_dir(PKG_DOCKER_DIR).await?)
+            .map_err(|e| {
+                Error::new(
+                    anyhow::Error::from(e).context(PKG_DOCKER_DIR),
+                    crate::ErrorKind::Filesystem,
+                )
+            })
+            .try_for_each_concurrent(None, |pkg_id| async move {
+                ReadDirStream::new(tokio::fs::read_dir(pkg_id.path()).await?)
+                    .map_err(|e| {
+                        Error::new(
+                            anyhow::Error::from(e).context(pkg_id.path().display().to_string()),
+                            crate::ErrorKind::Filesystem,
+                        )
+                    })
+                    .try_for_each_concurrent(None, |version| async move {
+                        let mut load = Command::new("docker")
+                            .arg("load")
+                            .stdin(Stdio::piped())
+                            .stderr(Stdio::piped())
+                            .spawn()?;
+                        let load_in = load.stdin.take().ok_or_else(|| {
+                            Error::new(
+                                anyhow!("Could not write to stdin of docker load"),
+                                crate::ErrorKind::Docker,
+                            )
+                        })?;
+                        let mut docker_rdr = File::open(version.path().join("image.tar")).await?;
+                        copy_and_shutdown(&mut docker_rdr, load_in).await?;
+                        let res = load.wait_with_output().await?;
+                        if !res.status.success() {
+                            Err(Error::new(
+                                anyhow!(
+                                    "{}",
+                                    String::from_utf8(res.stderr).unwrap_or_else(|e| format!(
+                                        "Could not parse stderr: {}",
+                                        e
+                                    ))
+                                ),
+                                crate::ErrorKind::Docker,
+                            ))
+                        } else {
+                            Ok(())
+                        }
+                    })
+                    .await
+            })
+            .await
+    } else {
+        Ok(())
+    }
+}