2 files changed, 61 insertions, 3 deletions
diff --git a/nixos/modules/virtualisation/amazon-image.nix b/nixos/modules/virtualisation/amazon-image.nix
index bf5c04543a7..fe248a94488 100644
--- a/nixos/modules/virtualisation/amazon-image.nix
+++ b/nixos/modules/virtualisation/amazon-image.nix
@@ -41,17 +41,23 @@ in
 
     boot.growPartition = cfg.hvm;
 
-    fileSystems."/" = {
+    fileSystems."/" = mkIf (!cfg.zfs.enable) {
       device = "/dev/disk/by-label/nixos";
       fsType = "ext4";
       autoResize = true;
     };
 
-    fileSystems."/boot" = mkIf cfg.efi {
+    fileSystems."/boot" = mkIf (cfg.efi || cfg.zfs.enable) {
+      # The ZFS image uses a partition labeled ESP whether or not we're
+      # booting with EFI.
       device = "/dev/disk/by-label/ESP";
       fsType = "vfat";
     };
 
+    services.zfs.expandOnBoot = mkIf cfg.zfs.enable "all";
+
+    boot.zfs.devNodes = mkIf cfg.zfs.enable "/dev/";
+
     boot.extraModulePackages = [
       config.boot.kernelPackages.ena
     ];
diff --git a/nixos/modules/virtualisation/amazon-options.nix b/nixos/modules/virtualisation/amazon-options.nix
index 2e807131e93..698edcd835a 100644
--- a/nixos/modules/virtualisation/amazon-options.nix
+++ b/nixos/modules/virtualisation/amazon-options.nix
@@ -1,7 +1,46 @@
 { config, lib, pkgs, ... }:
-{
+let
+  inherit (lib) types;
+in {
   options = {
     ec2 = {
+      zfs = {
+        enable = lib.mkOption {
+          default = false;
+          internal = true;
+          description = ''
+            Whether the EC2 instance uses a ZFS root.
+          '';
+        };
+
+        datasets = lib.mkOption {
+          description = ''
+            Datasets to create under the `tank` and `boot` zpools.
+
+            **NOTE:** This option is used only at image creation time, and
+            does not attempt to declaratively create or manage datasets
+            on an existing system.
+          '';
+
+          default = {};
+
+          type = types.attrsOf (types.submodule {
+            options = {
+              mount = lib.mkOption {
+                description = "Where to mount this dataset.";
+                type = types.nullOr types.string;
+                default = null;
+              };
+
+              properties = lib.mkOption {
+                description = "Properties to set on this dataset.";
+                type = types.attrsOf types.string;
+                default = {};
+              };
+            };
+          });
+        };
+      };
       hvm = lib.mkOption {
         default = lib.versionAtLeast config.system.stateVersion "17.03";
         internal = true;
@@ -18,4 +57,17 @@
       };
     };
   };
+
+  config = lib.mkIf config.ec2.zfs.enable {
+    networking.hostId = lib.mkDefault "00000000";
+
+    fileSystems = let
+      mountable = lib.filterAttrs (_: value: ((value.mount or null) != null)) config.ec2.zfs.datasets;
+    in lib.mapAttrs'
+      (dataset: opts: lib.nameValuePair opts.mount {
+        device = dataset;
+        fsType = "zfs";
+      })
+      mountable;
+  };
 }