From 067ce261779b60baf611bf6760397ae41c144791 Mon Sep 17 00:00:00 2001
From: Robert Hensing <robert@roberthensing.nl>
Date: Fri, 2 Oct 2020 11:34:40 +0200
Subject: [PATCH] Use dockerTools.streamLayeredImage if available

Technically this opens a new attack vector, but if you don't trust
the code you're deploying, you should already have taken precautions
because of nix-shell, direnv etc. This just adds arion to that list.
---
 src/haskell/lib/Arion/Images.hs        | 34 ++++++++++++++++++++++----
 src/nix/modules/composition/images.nix | 11 +++++++--
 src/nix/modules/service/image.nix      |  7 +++++-
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/src/haskell/lib/Arion/Images.hs b/src/haskell/lib/Arion/Images.hs
index b22e79b..40e0108 100644
--- a/src/haskell/lib/Arion/Images.hs
+++ b/src/haskell/lib/Arion/Images.hs
@@ -17,7 +17,8 @@ import Control.Lens
 import Data.Aeson.Lens
 
 data Image = Image
-  { image :: Text -- ^ file path
+  { image :: Maybe Text -- ^ image tar.gz file path
+  , imageExe :: Maybe Text -- ^ path to exe producing image tar
   , imageName :: Text
   , imageTag :: Text
   } deriving (Generic, Aeson.ToJSON, Aeson.FromJSON, Show)
@@ -38,10 +39,11 @@ loadImages fp = do
 
     isNew i = (imageName i <> ":" <> imageTag i) `notElem` loaded
 
-  traverse_ loadImage . map (toS . image) . filter isNew $ images
+  traverse_ loadImage . filter isNew $ images
 
-loadImage :: FilePath -> IO ()
-loadImage imgPath = withFile (imgPath) ReadMode $ \fileHandle -> do
+loadImage :: Image -> IO ()
+loadImage (Image { image = Just imgPath, imageName = name }) =
+  withFile (toS imgPath) ReadMode $ \fileHandle -> do
   let procSpec = (Process.proc "docker" [ "load" ]) {
           Process.std_in = Process.UseHandle fileHandle
         }
@@ -49,7 +51,29 @@ loadImage imgPath = withFile (imgPath) ReadMode $ \fileHandle -> do
     e <- Process.waitForProcess procHandle 
     case e of
       ExitSuccess -> pass
-      ExitFailure code -> panic $ "docker load (" <> show code <> ") failed for " <> toS imgPath
+      ExitFailure code ->
+        panic $ "docker load failed with exit code " <> show code <> " for image " <> name <> " from path " <> imgPath
+
+loadImage (Image { imageExe = Just imgExe, imageName = name }) = do
+  let loadSpec = (Process.proc "docker" [ "load" ]) { Process.std_in = Process.CreatePipe }
+  Process.withCreateProcess loadSpec $ \(Just inHandle) _out _err loadProcHandle -> do
+    let streamSpec = Process.proc (toS imgExe) []
+    Process.withCreateProcess streamSpec { Process.std_out = Process.UseHandle inHandle } $ \_ _ _ streamProcHandle ->
+      withAsync (Process.waitForProcess loadProcHandle) $ \loadExitAsync ->
+        withAsync (Process.waitForProcess streamProcHandle) $ \streamExitAsync -> do
+          r <- waitEither loadExitAsync streamExitAsync
+          case r of
+            Right (ExitFailure code) -> panic $ "image producer for image " <> name <> " failed with exit code " <> show code <> " from executable " <> imgExe
+            Right ExitSuccess -> pass
+            Left _ -> pass
+          loadExit <- wait loadExitAsync
+          case loadExit of
+            ExitFailure code -> panic $ "docker load failed with exit code " <> show code <> " for image " <> name <> " produced by executable " <> imgExe
+            _ -> pass
+          pass
+
+loadImage (Image { imageName = name }) = do
+  panic $ "image " <> name <> " doesn't specify an image file or imageExe executable"
 
 
 getDockerImages :: IO [TaggedImage]
diff --git a/src/nix/modules/composition/images.nix b/src/nix/modules/composition/images.nix
index dc93691..0377100 100644
--- a/src/nix/modules/composition/images.nix
+++ b/src/nix/modules/composition/images.nix
@@ -16,13 +16,20 @@ let
     (let
       inherit (service) build;
     in {
-      image = build.image.outPath;
       imageName = build.imageName or service.image.name;
       imageTag =
                  if build.image.imageTag != ""
                  then build.image.imageTag
                  else lib.head (lib.strings.splitString "-" (baseNameOf build.image.outPath));
-    });
+    } // (if build.image.isExe or false
+        then {
+          imageExe = build.image.outPath;
+        }
+        else {
+          image = build.image.outPath;
+        }
+      )
+    );
 in
 {
   options = {
diff --git a/src/nix/modules/service/image.nix b/src/nix/modules/service/image.nix
index 88b75d2..df78ac2 100644
--- a/src/nix/modules/service/image.nix
+++ b/src/nix/modules/service/image.nix
@@ -9,7 +9,12 @@ let
     (pkgs.writeText "dummy-config.json" (builtins.toJSON config.image.rawConfig))
   ];
 
-  builtImage = pkgs.dockerTools.buildLayeredImage {
+  buildOrStreamLayeredImage = args:
+    if pkgs.dockerTools?streamLayeredImage
+    then pkgs.dockerTools.streamLayeredImage args // { isExe = true; }
+    else pkgs.dockerTools.buildLayeredImage args;
+
+  builtImage = buildOrStreamLayeredImage {
     inherit (config.image)
       name
       contents