Modify a nixos image to use a BTRFS root

modify_nix_image.sh

#!/bin/bash

# NB: I am no bash pro, so this could be dangerous.
# Please read through it and make sure you double check my work.

# Assuming nixos image nixos-sd-image-21.05pre-git-aarch64-linux.img and sdcard at /dev/sdx
# Example usage: `sudo ./modify_image.sh nixos-sd-image-21.05pre-git-aarch64-linux.img /dev/sdx`
# NB: This will overwrite /dev/sdx!

set -Eeuf -o pipefail
set -x

loopdev=
readonly mpoint=/tmp/nixbtrfs orig=/tmp/nixorig

cleanup() {
  umount -R "${mpoint}" "${orig}"
  losetup -d "${loopdev}"
}

trap cleanup EXIT

msg() {
  echo "$1" > /dev/stderr
}

err() {
  msg "$1"
  exit 1
}

main() {
  if [ "${EUID}" -ne 0 ]; then
    err "Run as root."
  fi

  local image dev
  image=$1
  device=$2

  if [ ! "${device:0:5}" = "/dev/" ]; then
    err "Second argument should be the block device."
  fi

  while :; do
    lsblk "${device}"
    read -r -n 1 -p "Using image ${image} and device ${device}. Are you sure? y/n" response
    case "${response}" in
      Y | y)
        echo
        break
        ;;
      N | n)
        echo
        exit 0
        ;;
      *)
        echo
        msg "Unrecognized response"
        ;;
    esac
  done

  sfdisk --dump "${image}" | sfdisk "${device}"

  loopdev=$(losetup --show --find --partscan "${image}")

  dd if="${loopdev}p1" of="${device}1" bs=4M status=progress

  local btrfsdev
  btrfsdev=${device}2

  mkfs.btrfs --force --data=dup --metadata=dup --label=NIXOS_SD "${btrfsdev}"

  for mp in "${mpoint}" "${orig}"; do
    if mountpoint -q "${mp}"; then
      err "Is something already mounted on ${mp}?"
    fi
    mkdir -p "${mp}"
  done

  local SUBVOLS OPTIONS comma_separated_options
  SUBVOLS=(@ @boot @var @home @snapshots @nix @swap)
  OPTIONS=(noatime ssd_spread compress-force=zstd autodefrag)

  comma_separated_options=$(
    IFS=,
    printf "%s" "${OPTIONS[*]}"
  )

  mount -t btrfs -o "${comma_separated_options}" "${btrfsdev}" "${mpoint}"

  for subvol in "${SUBVOLS[@]}"; do
    btrfs subvolume create "${mpoint}/${subvol}"
  done
  umount -R "${mpoint}"

  mount -t btrfs -o "${comma_separated_options},subvol=@" "${btrfsdev}" "${mpoint}"

  mkdir -p "${mpoint}"/.snapshots

  for subvol in "${SUBVOLS[@]}"; do
    for skipsubvol in "@" "@snapshots" "@boot"; do
      [ "${subvol}" = "${skipsubvol}" ] && continue 2
    done

    # trim the @
    dest=${mpoint}/${subvol:1}

    mkdir -p "${dest}"
    mount -t btrfs -o "${comma_separated_options},subvol=${subvol}" "${btrfsdev}" "${dest}"
  done

  mount "${loopdev}p2" "${orig}"
  rsync --archive --hard-links --acls --xattrs --sparse --partial --progress --verbose --info=progress2 -- "${orig}/" "${mpoint}/"
  sync
}
main "$@"