Assign missing subordinate IDs during corporate setup

6 files changed, 371 insertions, 11 deletions

diff --git a/docs/canonical-laptop.md b/docs/canonical-laptop.md
index 7479b46..8b60195 100644
--- a/docs/canonical-laptop.md
+++ b/docs/canonical-laptop.md
@@ -112,7 +112,7 @@ just canonical-setup
 This installs declared packages, builds the locked Home-Manager profile,
 deploys the corporate home files, loads two program-specific AppArmor profiles,
 connects Thunderbird's GPG interface, permits Mattermost to use GNOME Keyring,
-and installs the GNOME extensions.
+assigns missing subordinate ID ranges, and installs the GNOME extensions.
 It does not remove packages or switch to another source when installation fails.
 
 Log out and back in through GDM. Then run:
@@ -183,16 +183,28 @@ machine. Do not expose the private key or full agent logs in support requests.
 Rootless Podman needs subordinate UID and GID ranges for the final authd user:
 
 ```sh
-getsubids "$USER"
-getsubids -g "$USER"
+/usr/bin/getsubids "$USER"
+/usr/bin/getsubids -g "$USER"
 podman info
 podman run --rm docker.io/library/alpine:latest id
 ```
 
-If either range is missing, have a free, non-overlapping range assigned through
-the system's account administration method. Do not copy another user's ranges
-or assume that an authd user can be changed with `usermod`. The setup does not
-rewrite `/etc/subuid`, `/etc/subgid`, or company account data.
+`just canonical-system` assigns missing ranges for the current account through
+Ubuntu's system Python and account lookup. It preserves existing allocations,
+including ranges assigned by numeric UID. New ranges follow `/etc/login.defs`
+and contain at least 65,536 IDs. Allocation excludes existing subordinate ranges
+and user/group IDs returned by the system account database. An external `subid`
+provider stops setup without changing either file.
+
+The helper locks account administration while it reads and updates `/etc/subuid`
+and `/etc/subgid`. It does not modify `/etc/passwd`, authd, or the login UID.
+Each file replacement is atomic; if an I/O error interrupts setup between the
+two files, rerun it after fixing the error. Existing ranges remain unchanged.
+Do not copy another user's ranges. Coordinate allocation with IT if the company
+reserves additional ID ranges that are not visible in these databases.
+
+If Podman was used before ranges were assigned, stop its containers and run
+`podman system migrate` as your work user before testing again.
 
 The repo loads `dotfiles-nix-bwrap` and `dotfiles-nix-podman` AppArmor
 profiles. Global user-namespace restrictions stay enabled. Test `aibox -p` and
diff --git a/scripts/canonical-system.sh b/scripts/canonical-system.sh
index cd85879..179bce8 100644
--- a/scripts/canonical-system.sh
+++ b/scripts/canonical-system.sh
@@ -8,6 +8,7 @@ source just-lib.sh
   . /etc/os-release
   echo "$ID"
 ) == ubuntu ]]
+sudo /usr/bin/python3 scripts/canonical_subids.py "$(/usr/bin/id -un)"
 # Parse before replacing the installed profile.
 sudo apparmor_parser --skip-kernel-load --skip-cache canonical/apparmor/dotfiles-nix
 sudo install -m 644 canonical/apparmor/dotfiles-nix /etc/apparmor.d/dotfiles-nix
diff --git a/scripts/canonical.py b/scripts/canonical.py
index db7f773..6fd83d6 100644
--- a/scripts/canonical.py
+++ b/scripts/canonical.py
@@ -108,15 +108,20 @@ def check(lab: bool = False) -> None:
         ["gnome-extensions", "list", "--enabled"],
         ["flatpak", "list", "--user", "--app"],
         *[["flatpak", "info", "--user", app] for app in packages("flatpak")],
-        ["getent", "passwd", str(os.getuid())],
-        ["getsubids", os.environ.get("USER", "")],
-        ["getsubids", "-g", os.environ.get("USER", "")],
+        ["/usr/bin/getent", "passwd", str(os.getuid())],
+        ["/usr/bin/getsubids", os.environ.get("USER", "")],
+        ["/usr/bin/getsubids", "-g", os.environ.get("USER", "")],
     ]
     failed = False
     for command in commands:
         print("\n> " + " ".join(command), flush=True)
         try:
-            failed |= subprocess.run(command, check=False).returncode != 0
+            result = subprocess.run(command, check=False)
+            failed |= result.returncode != 0
+            if result.returncode and command[0] == "/usr/bin/getsubids":
+                print(
+                    "Subordinate IDs unavailable. Run just canonical-system to configure local ranges."
+                )
         except FileNotFoundError:
             print(f"Missing: {command[0]}")
             failed = True
diff --git a/scripts/canonical_subids.py b/scripts/canonical_subids.py
new file mode 100644
index 0000000..c99891e
--- /dev/null
+++ b/scripts/canonical_subids.py
@@ -0,0 +1,151 @@
+"""Assign missing local subordinate IDs without changing existing allocations."""
+
+import argparse
+import ctypes
+import grp
+import os
+import pwd
+import subprocess
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+
+
+@contextmanager
+def account_lock():
+    libc = ctypes.CDLL(None, use_errno=True)
+    if libc.lckpwdf() != 0:
+        raise OSError("Cannot lock the account files; try again later.")
+    try:
+        yield
+    finally:
+        libc.ulckpwdf()
+
+
+def allocation_policy(etc: Path) -> dict[str, int]:
+    for line in (etc / "nsswitch.conf").read_text().splitlines():
+        key, _, value = line.split("#", 1)[0].partition(":")
+        if key.strip() == "subid" and value.split() != ["files"]:
+            raise ValueError(
+                "Subordinate IDs use an external provider; no files changed."
+            )
+    policy = {}
+    for line in (etc / "login.defs").read_text().splitlines():
+        fields = line.split("#", 1)[0].split()
+        if fields and fields[0] in {
+            f"SUB_{kind}_{limit}"
+            for kind in ("UID", "GID")
+            for limit in ("MIN", "MAX", "COUNT")
+        }:
+            key, value = fields
+            policy[key] = int(value)
+    return policy
+
+
+def read_ranges(text: str) -> list[tuple[str, int, int]]:
+    ranges = []
+    for line in text.splitlines():
+        if not line.strip() or line.lstrip().startswith("#"):
+            continue
+        owner, first, size = line.split(":")
+        start, count = int(first), int(size)
+        if not owner or start < 0 or count <= 0 or start + count > 2**32 - 1:
+            raise ValueError("Invalid subordinate ID range; no files changed.")
+        ranges.append((owner, start, count))
+    return ranges
+
+
+def free_range(
+    occupied: list[tuple[int, int]], minimum: int, maximum: int, count: int
+) -> int:
+    start = minimum
+    for first, end in sorted(occupied):
+        if start + count <= first:
+            break
+        start = max(start, end)
+    if start + count - 1 > maximum:
+        raise ValueError("No free subordinate ID range within login.defs limits.")
+    return start
+
+
+def allocation_limits(policy: dict[str, int], kind: str) -> tuple[int, int, int]:
+    minimum = policy.get(f"SUB_{kind}_MIN", 100000)
+    maximum = policy.get(f"SUB_{kind}_MAX", 600100000)
+    count = policy.get(f"SUB_{kind}_COUNT", 65536)
+    if count == 0:
+        raise ValueError(f"Subordinate {kind} allocation is disabled in login.defs.")
+    if not (0 < minimum <= maximum < 2**32 - 1) or count < 65536:
+        raise ValueError(f"Invalid or insufficient SUB_{kind} limits in login.defs.")
+    return minimum, maximum, count
+
+
+def plan_allocations(
+    etc: Path, name: str, uid: int, user_ids: set[int], group_ids: set[int]
+) -> dict[Path, str]:
+    policy = allocation_policy(etc)
+    changes = {}
+    for kind, used in (("UID", user_ids), ("GID", group_ids)):
+        path = etc / f"sub{kind.lower()}"
+        if path.is_symlink():
+            raise ValueError(f"Refusing to replace symlink: {path}")
+        text = path.read_text() if path.exists() else ""
+        ranges = read_ranges(text)
+        if any(owner in {name, str(uid)} for owner, _, _ in ranges):
+            continue
+        minimum, maximum, count = allocation_limits(policy, kind)
+        occupied = [(first, first + size) for _, first, size in ranges]
+        occupied.extend((number, number + 1) for number in used | {uid})
+        start = free_range(occupied, minimum, maximum, count)
+        separator = "\n" if text and not text.endswith("\n") else ""
+        changes[path] = f"{text}{separator}{name}:{start}:{count}\n"
+    return changes
+
+
+def write_atomic(path: Path, text: str) -> None:
+    info = path.stat() if path.exists() else None
+    with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as stream:
+        temporary = Path(stream.name)
+        try:
+            stream.write(text)
+            stream.flush()
+            if info:
+                os.fchown(stream.fileno(), info.st_uid, info.st_gid)
+            os.fchmod(stream.fileno(), (info.st_mode & 0o777) if info else 0o644)
+            os.fsync(stream.fileno())
+            os.replace(temporary, path)
+        finally:
+            temporary.unlink(missing_ok=True)
+
+
+def configure(name: str) -> None:
+    account = pwd.getpwnam(name)
+    if account.pw_uid == 0 or any(char in name for char in ":\n\r"):
+        raise ValueError("Select a non-root login account.")
+    with account_lock():
+        changes = plan_allocations(
+            Path("/etc"),
+            name,
+            account.pw_uid,
+            {entry.pw_uid for entry in pwd.getpwall()},
+            {entry.gr_gid for entry in grp.getgrall()} | {account.pw_gid},
+        )
+        for path, text in changes.items():
+            write_atomic(path, text)
+            print(f"Assigned a subordinate ID range in {path} for {name}.")
+    for flags in ([], ["-g"]):
+        subprocess.run(["/usr/bin/getsubids", *flags, name], check=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("user")
+    args = parser.parse_args()
+    if os.geteuid() != 0:
+        raise SystemExit(
+            "Run this command through sudo with Ubuntu's /usr/bin/python3."
+        )
+    configure(args.user)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_canonical.py b/tests/test_canonical.py
index 27fd199..14b7a3c 100644
--- a/tests/test_canonical.py
+++ b/tests/test_canonical.py
@@ -17,6 +17,23 @@ SPEC.loader.exec_module(canonical)
 
 
 class PackageTests(unittest.TestCase):
+    def test_subid_failure_remains_fatal_with_setup_guidance(self):
+        def result(command, **kwargs):
+            return subprocess.CompletedProcess(
+                command, int(command[0] == "/usr/bin/getsubids")
+            )
+
+        with (
+            patch.object(canonical.subprocess, "run", side_effect=result),
+            patch("builtins.print") as printed,
+            self.assertRaises(SystemExit) as failure,
+        ):
+            canonical.check()
+        self.assertEqual(failure.exception.code, 1)
+        self.assertTrue(
+            any("just canonical-system" in str(call) for call in printed.call_args_list)
+        )
+
     def test_install_does_not_enable_experimental_snap_features(self):
         with patch.object(canonical.subprocess, "run") as command:
             canonical.install()
diff --git a/tests/test_canonical_subids.py b/tests/test_canonical_subids.py
new file mode 100644
index 0000000..3cf16f6
--- /dev/null
+++ b/tests/test_canonical_subids.py
@@ -0,0 +1,174 @@
+import tempfile
+import unittest
+from contextlib import nullcontext
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from scripts import canonical_subids as subids
+
+
+class SubidTests(unittest.TestCase):
+    def setUp(self):
+        self.directory = tempfile.TemporaryDirectory()
+        self.addCleanup(self.directory.cleanup)
+        self.etc = Path(self.directory.name)
+        for name, text in {
+            "nsswitch.conf": "passwd: files authd\ngroup: files authd\n",
+            "login.defs": "SUB_UID_MIN 1000000\nSUB_GID_MIN 1000000\n",
+            "subuid": "ruifm:1000000:65536\n",
+            "subgid": "ruifm:1000000:65536\n",
+        }.items():
+            (self.etc / name).write_text(text)
+        self.uid = 10000
+        self.name = "rui.marques@canonical.com"
+
+    def plan(self):
+        return subids.plan_allocations(self.etc, self.name, self.uid, {10000}, {10000})
+
+    def test_allocates_after_existing_ranges(self):
+        changes = self.plan()
+        for name in ("subuid", "subgid"):
+            self.assertEqual(
+                changes[self.etc / name],
+                f"ruifm:1000000:65536\n{self.name}:1065536:65536\n",
+            )
+
+    def test_preserves_existing_name_and_numeric_allocations(self):
+        (self.etc / "subuid").write_text(f"{self.name}:2000000:1000\n")
+        (self.etc / "subgid").write_text(f"{self.uid}:3000000:65536\n")
+        self.assertEqual(self.plan(), {})
+
+    def test_fills_only_missing_file_and_is_idempotent(self):
+        (self.etc / "subgid").write_text(f"{self.name}:2000000:65536\n")
+        changes = self.plan()
+        self.assertEqual(list(changes), [self.etc / "subuid"])
+        for path, text in changes.items():
+            subids.write_atomic(path, text)
+        self.assertEqual(self.plan(), {})
+
+    def test_skips_real_user_and_group_ids(self):
+        changes = subids.plan_allocations(
+            self.etc, self.name, self.uid, {1065536}, {1065537}
+        )
+        self.assertIn(f"{self.name}:1065537:65536", changes[self.etc / "subuid"])
+        self.assertIn(f"{self.name}:1065538:65536", changes[self.etc / "subgid"])
+
+    def test_rejects_external_provider(self):
+        (self.etc / "nsswitch.conf").write_text("subid: sss\n")
+        with self.assertRaisesRegex(ValueError, "provider"):
+            self.plan()
+
+    def test_rejects_exhausted_range(self):
+        with (self.etc / "login.defs").open("a") as stream:
+            stream.write("SUB_UID_MAX 1065535\n")
+        with self.assertRaisesRegex(ValueError, "free"):
+            self.plan()
+
+    def test_rejects_disabled_allocation(self):
+        with (self.etc / "login.defs").open("a") as stream:
+            stream.write("SUB_UID_COUNT 0\n")
+        with self.assertRaisesRegex(ValueError, "disabled"):
+            self.plan()
+
+    def test_rejects_malformed_ranges_before_writing(self):
+        (self.etc / "subgid").write_text("broken entry\n")
+        with self.assertRaises(ValueError):
+            self.plan()
+        self.assertEqual((self.etc / "subuid").read_text(), "ruifm:1000000:65536\n")
+
+    def test_handles_missing_file_and_missing_final_newline(self):
+        (self.etc / "subuid").unlink()
+        (self.etc / "subgid").write_text("ruifm:1000000:65536")
+        changes = self.plan()
+        self.assertEqual(changes[self.etc / "subuid"], f"{self.name}:1000000:65536\n")
+        self.assertIn("65536\n" + self.name, changes[self.etc / "subgid"])
+
+    def test_atomic_write_preserves_mode_and_leaves_no_temporary_file(self):
+        target = self.etc / "subuid"
+        target.chmod(0o640)
+        subids.write_atomic(target, "replacement\n")
+        self.assertEqual(target.read_text(), "replacement\n")
+        self.assertEqual(target.stat().st_mode & 0o777, 0o640)
+        self.assertEqual(len(list(self.etc.iterdir())), 4)
+
+    def test_failed_replace_preserves_file(self):
+        target = self.etc / "subuid"
+        with (
+            patch.object(subids.os, "replace", side_effect=OSError("failure")),
+            self.assertRaises(OSError),
+        ):
+            subids.write_atomic(target, "replacement\n")
+        self.assertEqual(target.read_text(), "ruifm:1000000:65536\n")
+        self.assertEqual(len(list(self.etc.iterdir())), 4)
+
+    def test_rejects_symlink_without_touching_its_target(self):
+        path = self.etc / "subuid"
+        path.unlink()
+        path.symlink_to(self.etc / "subgid")
+        with self.assertRaisesRegex(ValueError, "symlink"):
+            self.plan()
+        self.assertEqual((self.etc / "subgid").read_text(), "ruifm:1000000:65536\n")
+
+    def test_respects_larger_configured_range_count(self):
+        with (self.etc / "login.defs").open("a") as stream:
+            stream.write("SUB_UID_COUNT 131072\n")
+        self.assertIn(f"{self.name}:1065536:131072", self.plan()[self.etc / "subuid"])
+
+    def test_lock_is_released_on_error(self):
+        libc = MagicMock()
+        libc.lckpwdf.return_value = 0
+        with (
+            patch.object(subids.ctypes, "CDLL", return_value=libc),
+            self.assertRaises(ValueError),
+            subids.account_lock(),
+        ):
+            raise ValueError("failure")
+        libc.ulckpwdf.assert_called_once_with()
+
+    def test_failed_lock_does_not_enter_or_unlock(self):
+        libc = MagicMock()
+        libc.lckpwdf.return_value = -1
+        with (
+            patch.object(subids.ctypes, "CDLL", return_value=libc),
+            self.assertRaises(OSError),
+            subids.account_lock(),
+        ):
+            self.fail("Entered without the account lock")
+        libc.ulckpwdf.assert_not_called()
+
+    def test_configure_verifies_both_ranges_with_system_commands(self):
+        with (
+            patch.object(
+                subids.pwd,
+                "getpwnam",
+                return_value=SimpleNamespace(pw_uid=10000, pw_gid=10000),
+            ),
+            patch.object(subids.pwd, "getpwall", return_value=[]),
+            patch.object(subids.grp, "getgrall", return_value=[]),
+            patch.object(subids, "account_lock", return_value=nullcontext()),
+            patch.object(subids, "plan_allocations", return_value={}) as plan,
+            patch.object(subids, "write_atomic") as write,
+            patch.object(subids.subprocess, "run") as run,
+        ):
+            subids.configure(self.name)
+        plan.assert_called_once_with(Path("/etc"), self.name, 10000, set(), {10000})
+        write.assert_not_called()
+        self.assertEqual(run.call_count, 2)
+        run.assert_any_call(["/usr/bin/getsubids", self.name], check=True)
+        run.assert_any_call(["/usr/bin/getsubids", "-g", self.name], check=True)
+
+    def test_root_account_is_rejected_before_file_access(self):
+        with (
+            patch.object(
+                subids.pwd, "getpwnam", return_value=SimpleNamespace(pw_uid=0)
+            ),
+            patch.object(subids, "account_lock") as lock,
+            self.assertRaises(ValueError),
+        ):
+            subids.configure("root")
+        lock.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()