From 9ac58819e623c9b9ca1c1b52c5b7fb5c56c7b1ff Mon Sep 17 00:00:00 2001
From: xinyangli <lixinyang411@gmail.com>
Date: Tue, 11 Jun 2024 18:19:43 +0800
Subject: [PATCH] dolmite: support bandwagon and lightsail

---
 machines/dolomite/bandwagon.nix |   7 +--
 machines/dolomite/default.nix   |   8 +--
 machines/dolomite/lightsail.nix | 103 ++++++++++++++++++++++++++++++--
 3 files changed, 105 insertions(+), 13 deletions(-)

diff --git a/machines/dolomite/bandwagon.nix b/machines/dolomite/bandwagon.nix
index 853f8d8..32d2b9f 100644
--- a/machines/dolomite/bandwagon.nix
+++ b/machines/dolomite/bandwagon.nix
@@ -10,7 +10,7 @@ in
     isBandwagon = lib.mkEnableOption "Bandwagon instance";
   };
 
-  config = lib.mkIf cfg.isBandwagon {
+  config = lib.mkIf cfg {
     boot.initrd.availableKernelModules = [ "ata_piix" "xhci_pci" "virtio_pci" "virtio_scsi" "sd_mod" "sr_mod" ];
     boot.initrd.kernelModules = [ ];
     boot.kernelModules = [ ];
@@ -28,9 +28,8 @@ in
 
     swapDevices = [ ];
 
-    boot.loader.grub.enable = lib.mkForce true;
-    boot.loader.grub.version = lib.mkForce 2;
-    boot.loader.grub.device = lib.mkForce "/dev/sda";
+    boot.loader.grub.enable = true;
+    boot.loader.grub.device = "/dev/sda";
     networking.useDHCP = false;
     networking.interfaces.ens18.useDHCP = true;
     networking.interfaces.ens19.useDHCP = true;
diff --git a/machines/dolomite/default.nix b/machines/dolomite/default.nix
index 15f7e2e..e8b2797 100644
--- a/machines/dolomite/default.nix
+++ b/machines/dolomite/default.nix
@@ -1,13 +1,13 @@
-{ inputs, config, pkgs, lib, modulesPath, ... }:
+{ config, lib, ... }:
 let
-  awsHosts = [ "sgp-00" "tok-00 "];
+  awsHosts = [ "tok-00 "];
   bwgHosts = [ "la-00" ];
 in
 {
   imports = [
     ../sops.nix
-   ./bandwagon.nix
-   ./lightsail.nix
+    ./bandwagon.nix
+    ./lightsail.nix
   ];
 
 
diff --git a/machines/dolomite/lightsail.nix b/machines/dolomite/lightsail.nix
index 187c6ff..a71c460 100644
--- a/machines/dolomite/lightsail.nix
+++ b/machines/dolomite/lightsail.nix
@@ -1,13 +1,106 @@
 { config, lib, pkgs, modulesPath, ... }:
+with lib;
 let
-  cfg = config.isLightsail;
+  cfg = config.ec2;
 in
 {
-  imports = [ "${modulesPath}/virtualisation/amazon-image.nix" ];
+  imports = [ 
+    "${modulesPath}/profiles/headless.nix"
+    # Note: While we do use the headless profile, we also explicitly
+    # turn on the serial console on ttyS0 below. This is because
+    # AWS does support accessing the serial console:
+    # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configure-access-to-serial-console.html
+    "${modulesPath}/virtualisation/ec2-data.nix"
+    "${modulesPath}/virtualisation/amazon-init.nix"
+  ];
+
   options = {
-    isLightsail = lib.mkEnableOption "Lightsail instance";
+    isLightsail = mkEnableOption "Lightsail instance";
   };
-  config = lib.mkIf cfg.isLightsail{
-    boot.loader.grub.device = lib.mkForce "/dev/nvme0n1";
+
+  config = mkIf config.isLightsail {
+    boot.loader.grub.device = "/dev/nvme0n1";
+
+    # from nixpkgs amazon-image.nix
+    assertions = [ ];
+
+    boot.growPartition = true;
+
+    fileSystems."/" = mkIf (!cfg.zfs.enable) {
+      device = "/dev/disk/by-label/nixos";
+      fsType = "ext4";
+      autoResize = true;
+    };
+
+    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
+    ];
+    boot.initrd.kernelModules = [ "xen-blkfront" ];
+    boot.initrd.availableKernelModules = [ "nvme" ];
+    boot.kernelParams = [ "console=ttyS0,115200n8" "random.trust_cpu=on" ];
+
+    # Prevent the nouveau kernel module from being loaded, as it
+    # interferes with the nvidia/nvidia-uvm modules needed for CUDA.
+    # Also blacklist xen_fbfront to prevent a 30 second delay during
+    # boot.
+    boot.blacklistedKernelModules = [ "nouveau" "xen_fbfront" ];
+
+    boot.loader.grub.efiSupport = cfg.efi;
+    boot.loader.grub.efiInstallAsRemovable = cfg.efi;
+    boot.loader.timeout = 1;
+    boot.loader.grub.extraConfig = ''
+      serial --unit=0 --speed=115200 --word=8 --parity=no --stop=1
+      terminal_output console serial
+      terminal_input console serial
+    '';
+
+    systemd.services.fetch-ec2-metadata = {
+      wantedBy = [ "multi-user.target" ];
+      wants = [ "network-online.target" ];
+      after = ["network-online.target"];
+      path = [ pkgs.curl ];
+      script = builtins.readFile ./ec2-metadata-fetcher.sh;
+      serviceConfig.Type = "oneshot";
+      serviceConfig.StandardOutput = "journal+console";
+    };
+
+    # Amazon-issued AMIs include the SSM Agent by default, so we do the same.
+    # https://docs.aws.amazon.com/systems-manager/latest/userguide/ami-preinstalled-agent.html
+    services.amazon-ssm-agent.enable = true;
+
+    # Allow root logins only using the SSH key that the user specified
+    # at instance creation time.
+    services.openssh.enable = true;
+    services.openssh.settings.PermitRootLogin = "prohibit-password";
+
+    # Enable the serial console on ttyS0
+    systemd.services."serial-getty@ttyS0".enable = true;
+
+    # Creates symlinks for block device names.
+    services.udev.packages = [ pkgs.amazon-ec2-utils ];
+
+    # Force getting the hostname from EC2.
+    # networking.hostName = mkDefault "";
+
+    # Always include cryptsetup so that Charon can use it.
+    environment.systemPackages = [ pkgs.cryptsetup ];
+
+    # EC2 has its own NTP server provided by the hypervisor
+    networking.timeServers = [ "169.254.169.123" ];
+
+    # udisks has become too bloated to have in a headless system
+    # (e.g. it depends on GTK).
+    services.udisks2.enable = false;
   };
 }