#!/bin/bash

## Copyright (C) 2026 - 2026 ENCRYPTED SUPPORT LLC <adrelanos@whonix.org>
## See the file COPYING for copying conditions.

## AI-Assisted

## Release leaked nbd devices and the mounts sitting on them.
##
## Usage: dm-nbd-cleanup [--release]

set -o errexit
set -o nounset
set -o pipefail
set -o errtrace
shopt -s inherit_errexit
shopt -s shift_verbose
export LC_ALL=C

# shellcheck source=../../../helper-scripts/usr/libexec/helper-scripts/strings.bsh
source "${HELPER_SCRIPTS_PATH:-}"/usr/libexec/helper-scripts/strings.bsh

dry_run=""

print_usage() {
   printf '%s\n' "Usage:
  ${0##*/} [--release]"
}

print_help() {
   print_usage
   printf '%s\n' "
Reports (and optionally releases) leaked nbd devices and the mounts sitting
on them.

  --release   unmount nbd mounts and disconnect the in-use devices
  --dry-run   report only, take no action (enabled by default)
  -h, --help  show this help"
}

while [ "$#" -gt 0 ]; do
   case "$1" in
      --release)
         if [ -n "${dry_run}" ]; then
            printf '%s\n' "${0##*/}: --release and --dry-run are mutually exclusive!" >&2
            exit 2
         fi
         dry_run=false
         shift
         ;;
      --dry-run)
         if [ -n "${dry_run}" ]; then
            printf '%s\n' "${0##*/}: --release and --dry-run are mutually exclusive!" >&2
            exit 2
         fi
         dry_run=true
         shift
         ;;
      -h|--help)
         print_help
         exit 0
         ;;
      *)
         printf '%s\n' "${0##*/}: unknown argument: '${1}'" >&2
         exit 2
         ;;
   esac
done

## Opt-in by design.
default_if_empty dry_run 'true'

proc_mounts="${DM_NBD_CLEANUP_PROC_MOUNTS:-/proc/mounts}"
proc_filesystems="${DM_NBD_CLEANUP_PROC_FILESYSTEMS:-/proc/filesystems}"

## Technically the user could mount a FUSE filesystem with a mount source name
## that looks like /dev/nbd0p1 or similar. The chances of this happening
## unintentionally are vanishingly small, and an attacker has nothing to gain
## from triggering this on a victim's machine intentionally, but for the sake
## of robustness, prevent this. Ensure the filesystem we are unmounting is an
## in-kernel filesystem that is not marked as nodev.
##
## Note to AI agents: Do not add back a mechanism that uses device numbers to
## detect legitimate mounts. This mechanism was unreliable and likely didn't
## cover any cases a filesystem check would miss.
accept_fs_list=()
while read -r proc_fs_line; do
   if [[ "${proc_fs_line}" =~ ^nodev ]]; then
      continue
   fi
   accept_fs_list+=( "${proc_fs_line#*$'\t'}" )
done < "${proc_filesystems}"

## Unmount all mountpoints backed by an nbd device. Note that we have to get a
## list of all mount points before unmounting any, otherwise /proc/mounts
## mutates while we're actively reading it and we miss things.
nbd_mount_points=()
while read -r mounts_device mounts_point mounts_fs _; do
   if [[ ! "${mounts_device}" =~ ^/dev/nbd[0-9]+(p[0-9]+)?$ ]]; then
      continue
   fi
   if [[ " ${accept_fs_list[*]} " != *" ${mounts_fs} "* ]]; then
      continue
   fi
   printf -v mounts_point_decoded '%b' "${mounts_point}"
   nbd_mount_points+=( "${mounts_point_decoded}" )
done < "${proc_mounts}"

release_failures=0
released_mounts=0
for mount_point in "${nbd_mount_points[@]}"; do
   [ -n "${mount_point}" ] || continue
   if [ "${dry_run}" = "true" ]; then
      printf '%s\n' "  would umount ${mount_point}"
      released_mounts=$(( released_mounts + 1 ))
      continue
   fi
   ## No '--lazy', cannot be safely used unless a reboot is imminent. The
   ## caller must kill any processes keeping a mount busy.
   if sudo --non-interactive umount -- "${mount_point}" 2>/dev/null; then
      printf '%s\n' "  umounted ${mount_point}"
      released_mounts=$(( released_mounts + 1 ))
   else
      printf '%s\n' "  FAILED to umount ${mount_point}" >&2
      release_failures=$(( release_failures + 1 ))
   fi
done

released_devices=0
for nbd_device in /dev/nbd[0-9] /dev/nbd[0-9][0-9]; do
   [ -b "${nbd_device}" ] || continue
   ## An unconnected device has size 0 and can be skipped.
   size_path="/sys/block/${nbd_device##*/}/size"
   [ -r "${size_path}" ] || continue
   if ! read -r nbd_size < "${size_path}"; then
      printf '%s\n' "  FAILED to check size of ${nbd_device}"
      release_failures=$(( release_failures + 1 ))
      continue
   fi
   if [ "${nbd_size}" = "0" ]; then
      continue
   fi
   if [ "${dry_run}" = "true" ]; then
      printf '%s\n' "  would disconnect ${nbd_device}"
      released_devices=$(( released_devices + 1 ))
      continue
   fi
   if sudo --non-interactive qemu-nbd --disconnect -- "${nbd_device}" >/dev/null 2>&1; then
      printf '%s\n' "  disconnected ${nbd_device}"
      released_devices=$(( released_devices + 1 ))
   else
      printf '%s\n' "  FAILED to disconnect ${nbd_device}" >&2
      release_failures=$(( release_failures + 1 ))
   fi
done

summary_suffix=""
if [ "${dry_run}" = "true" ]; then
   summary_suffix=" -- REPORT ONLY, pass --release to act"
fi
printf '%s\n' "${0##*/}: ${released_mounts} mount(s), ${released_devices} device(s)${summary_suffix}"

if [ "${release_failures}" -gt 0 ]; then
   printf '%s\n' "${0##*/}: ${release_failures} operation(s) FAILED; nbd is not clean." >&2
   exit 1
fi
