Add the ability to process quota config

Author: ben-z

src/index.mjs

@@ -5,7 +5,7 @@ import './patch.mjs';
 import path from "path";
 import { $, stdin, argv, question } from "zx";
 import { readFile } from "node:fs/promises";
-import { getExistingDirectory, parseConfig, diffProperties, deepEqual, getSSHKeys } from "./utils.mjs";
+import { getExistingDirectory, parseConfig, diffProperties, deepEqual, getSSHKeys, getDiskQuota, unique } from "./utils.mjs";
 import { validateConfig } from "./schema.mjs";
 
 
@@ -25,6 +25,10 @@ if (!argv.config) {
   printUsageAndExit(1);
 }
 
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Load, validate, and parse config
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 console.time("readConfig");
 let config;
 if (argv.config === "-") {
@@ -44,6 +48,24 @@ if (!validateConfig(config)) {
 }
 console.timeLog("validateConfig");
 
+console.log("Parsing config");
+console.time("parseConfig");
+const {
+  configGroups,
+  configUsers,
+  configPasswords,
+  configSSHKeys,
+  configUpdatePassword,
+  configLinger,
+  configDiskUserQuotaOverrides,
+  xfsDefaultDiskUserQuota,
+} = parseConfig(config);
+console.timeLog("parseConfig");
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Load existing directory
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 console.log("Loading existing directory...");
 console.time("getExistingDirectory");
 const { users, passwords, groups, lingerStates } = await getExistingDirectory();
@@ -61,10 +83,24 @@ console.time("getSSHKeys");
 const sshKeys = await getSSHKeys(Object.values(users), config.user_ssh_key_base_dir);
 console.timeLog("getSSHKeys");
 
-console.log("Parsing config");
-console.time("parseConfig");
-const { configGroups, configUsers, configPasswords, configSSHKeys, configUpdatePassword, configLinger } = parseConfig(config);
-console.timeLog("parseConfig");
+console.log("Loading existing disk quota");
+console.time("getDiskQuota");
+const diskQuota = await getDiskQuota(unique([...config.xfs_default_user_quota.map(q => q.path), ...Object.keys(configDiskUserQuotaOverrides)]));
+console.timeLog("getDiskQuota");
+
+// TODO: default user quota changes (per path)
+console.log("diskQuota", diskQuota);
+console.log(diskQuota['/var/lib/cluster'][0])
+console.log(xfsDefaultDiskUserQuota['/var/lib/cluster'])
+
+// TODO: user quota changes (per path)
+// new: in config but not existing
+// delete: existing (within the managed uid range) but not in config
+// update: in config and existing, but different
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Calculate changes
+/////////////////////////////////////////////////////////////////////////////////////////////////
 
 console.log("Calculating changes");
 console.time("calculateChanges");

src/schema.mjs

@@ -102,12 +102,12 @@ export const configSchema = {
       description: "Base directory for user SSH keys. Supports templating with %u (username) and %U (uid)",
       default: "/home/%u/.ssh",
     },
-    // sudo xfs_quota <path> -x -c "limit -d bsoft=<bsoft> bhard=<bhard> isoft=<isoft> ihard=<ihard>"
-    xfs_default_quotas: {
+    // XFS default quota can be set using `sudo xfs_quota <path> -x -c "limit -d bsoft=<bsoft> bhard=<bhard> isoft=<isoft> ihard=<ihard>"`
+    xfs_default_user_quota: {
       type: "array",
       items: quotaSpec,
       default: [],
-    }
+    },
   },
   required: ["users", "groups", "managed_uid_range", "managed_gid_range"],
   additionalProperties: false,

src/utils.mjs

@@ -1,11 +1,82 @@
 import { readFile, readdir } from 'node:fs/promises';
 
+// Converts a human-readable size string to a number of bytes
+// E.g. "1Ki" -> 1024, "1Mi" -> 1048576, "1Gi" -> 1073741824
+export function parseIECSize(s) {
+  const match = s.match(/^([0-9]+)([kKmMgGtT])?i?$/);
+  if (!match) {
+    throw new Error(`Invalid size string: ${s}`);
+  }
+  const size = Number(match[1]);
+  const unit = match[2]?.toLowerCase();
+  if (!unit) {
+    return size;
+  }
+  const unitMap = {
+    k: 1024,
+    m: Math.pow(1024, 2),
+    g: Math.pow(1024, 3),
+    t: Math.pow(1024, 4),
+  };
+  return size * unitMap[unit];
+}
+
+// Converts a human-readable size string to a number
+// E.g. "1K" -> 1000, "1M" -> 1000000, "1G" -> 1000000000
+export function parseSISize(s) {
+  const match = s.match(/^([0-9]+)([kKmMgGtT])?$/);
+  if (!match) {
+    throw new Error(`Invalid size string: ${s}`);
+  }
+  const size = Number(match[1]);
+  const unit = match[2]?.toLowerCase();
+  if (!unit) {
+    return size;
+  }
+  const unitMap = {
+    k: 1000,
+    m: Math.pow(1000, 2),
+    g: Math.pow(1000, 3),
+    t: Math.pow(1000, 4),
+  };
+  return size * unitMap[unit];
+}
+
+export function normalizeDiskQuota(q) {
+  return {
+    path: q.path,
+    bytes_soft_limit: parseIECSize(q.bytes_soft_limit),
+    bytes_hard_limit: parseIECSize(q.bytes_hard_limit),
+    inodes_soft_limit: parseSISize(q.inodes_soft_limit),
+    inodes_hard_limit: parseSISize(q.inodes_hard_limit),
+  };
+}
+
 export function parseConfig(config) {
   const configGroups = Object.fromEntries(config.groups.map((g) => [g.groupname, g]));
   const configUpdatePassword = Object.fromEntries(config.users.map((u) => [u.username, u.update_password]));
   const configPasswords = Object.fromEntries(config.users.map((u) => [u.username, u.password]));
   const configSSHKeys = Object.fromEntries(config.users.map((u) => [u.username, u.ssh_authorized_keys]));
   const configLinger = Object.fromEntries(config.users.map((u) => [u.username, u.linger]));
+  // Object of the form { <path>: { <uid>: { ...quotaConfig } } }
+  const configDiskUserQuotaOverrides = objectMap(
+    groupBy(
+      config.users.flatMap((u) =>
+        u.disk_quota_overrides.map(normalizeDiskQuota).map((d) => {
+          const { path, ...quotaConfig } = d;
+          return [path, u.uid, quotaConfig];
+        })
+      ),
+      // group by path
+      (x) => x[0]
+    ),
+    // convert to object of the form { <uid>: { ...quotaConfig } }
+    (v) => Object.fromEntries(v.map((x) => [x[1], x[2]]))
+  );
+  const xfsDefaultDiskUserQuota = Object.fromEntries(config.xfs_default_user_quota.map(normalizeDiskQuota).map((d) => {
+    const { path, ...quotaConfig } = d
+    return [path, quotaConfig];
+  }));
 
   const configUsers = config.users.reduce((out, u) => {
     const {
@@ -14,6 +85,7 @@ export function parseConfig(config) {
       update_password: _update_password,
       ssh_authorized_keys: _ssh_authorized_keys,
       linger: _linger,
+      disk_quota_overrides,
       ...rest
     } = u;
     out[u.username] = {
@@ -30,6 +102,8 @@ export function parseConfig(config) {
     configSSHKeys,
     configUpdatePassword,
     configLinger,
+    configDiskUserQuotaOverrides,
+    xfsDefaultDiskUserQuota,
   };
 }
 
@@ -157,6 +231,45 @@ export function diffProperties(obj1, obj2) {
   return out;
 }
 
+/**
+ * Maps over the properties of an object and applies a function to each value.
+ * Returns a new object with the transformed values.
+ * Derived from: https://stackoverflow.com/a/14810722
+ *
+ * @param {Object} obj - The object to map over.
+ * @param {Function} fn - The function to apply to each value.
+ * @returns {Object} - The new object with transformed values.
+ *
+ * @example
+ * const obj = { a: 1, b: 2, c: 3 };
+ * const double = (value) => value * 2;
+ * const transformedObj = mapObjectValues(obj, double);
+ * // transformedObj: { a: 2, b: 4, c: 6 }
+ *
+ * @example
+ * const obj = { name: 'John', age: 30 };
+ * const capitalize = (value) => value.toUpperCase();
+ * const transformedObj = mapObjectValues(obj, capitalize);
+ * // transformedObj: { name: 'JOHN', age: '30' }
+ */
+export function objectMap(obj, fn) {
+  return Object.fromEntries(
+    Object.entries(obj).map(
+      ([k, v], i) => [k, fn(v, k, i)]
+    )
+  )
+}
+
+// Implementation of Object.groupBy according to https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Object/groupBy
+export function groupBy(arr, fn) {
+  return arr.reduce((out, x) => {
+    const key = fn(x);
+    if (!(key in out)) out[key] = [];
+    out[key].push(x);
+    return out;
+  }, {})
+}
+
 export async function getSSHKeys(users, baseDir) {
   const sshKeyFiles = await Promise.all(
     users.map(async (u) => {
@@ -168,4 +281,38 @@ export async function getSSHKeys(users, baseDir) {
   );
 
   return Object.fromEntries(sshKeyFiles);
+}
+
+/**
+ * Retrieves the disk quota information for a given path.
+ * @param {string} path - The path for which to retrieve the disk quota.
+ * @returns {Promise<Object>} - A promise that resolves to an object of the form { <uid>: { bytes_soft_limit, bytes_hard_limit, inodes_soft_limit, inodes_hard_limit } }
+ */
+export async function getDiskQuotaForPath(path) {
+  // outputs <uid> <block soft> <block hard> <inode soft> <inode hard> where block is in 1KiB units
+  const repquotaResult = await $`repquota ${path} --user --no-names --raw-grace | grep '^#' | awk '{print $1,$4,$5,$8,$9}' | cut -c2-`;
+  const lines = repquotaResult.stdout
+    .split("\n")
+    .filter((l) => l)
+    .map((l) => l.split(" "));
+
+  return Object.fromEntries(lines.map((l) => [Number(l[0]), {
+    bytes_soft_limit: Number(l[1]),
+    bytes_hard_limit: Number(l[2]),
+    inodes_soft_limit: Number(l[3]),
+    inodes_hard_limit: Number(l[4]),
+  }]));
+}
+
+/**
+ * Retrieves the disk quota for multiple paths.
+ * @param {string[]} paths - An array of paths for which to retrieve the disk quota.
+ * @returns {Promise<Object>} - A promise that resolves to an object of the form { <path>: { <uid>: { bytes_soft_limit, bytes_hard_limit, inodes_soft_limit, inodes_hard_limit } } }
+ */
+export async function getDiskQuota(paths) {
+  return Object.fromEntries(await Promise.all(paths.map(async (p) => [p, await getDiskQuotaForPath(p)])));
+}
+
+export function unique(arr) {
+  return [...new Set(arr)];
 }

src/utils.test.mjs

@@ -1,5 +1,5 @@
 import fsPromises from "node:fs/promises";
-import { deepEqual, getSSHKeys, diffProperties, getExistingDirectory } from "./utils";
+import { deepEqual, getSSHKeys, diffProperties, getExistingDirectory, parseIECSize, objectMap } from "./utils";
 
 describe("deepEqual", () => {
   // Test case: Comparing two identical objects
@@ -215,3 +215,64 @@ describe("getExistingDirectory", () => {
     await expect(getExistingDirectory()).rejects.toThrow("user and password lists don't match up! users: user1, passwords: user1,user2");
   });
 });
+
+describe("objectMap", () => {
+  // Test case: Mapping empty object
+  test("should return an empty object when mapping an empty object", () => {
+    const obj = {};
+    const result = objectMap(obj, (value) => value);
+    const expected = {};
+    expect(result).toEqual(expected);
+  });
+
+  // Test case: Mapping object values
+  test("should map object values using the provided function", () => {
+    const obj = { a: 1, b: 2, c: 3 };
+    const result = objectMap(obj, (value) => value * 2);
+    const expected = { a: 2, b: 4, c: 6 };
+    expect(result).toEqual(expected);
+  });
+});
+
+describe("parseIECSize", () => {
+  // Test case: Valid size string without unit
+  test("should return the parsed size when given a valid size string without unit", () => {
+    const sizeString = "100";
+    const expectedSize = 100;
+    expect(parseIECSize(sizeString)).toBe(expectedSize);
+  });
+
+  // Test case: Valid size string with unit 'k'
+  test("should return the parsed size when given a valid size string with unit 'k'", () => {
+    const sizeString = "10k";
+    const expectedSize = 10240;
+    expect(parseIECSize(sizeString)).toBe(expectedSize);
+  });
+
+  // Test case: Valid size string with unit 'm'
+  test("should return the parsed size when given a valid size string with unit 'm'", () => {
+    const sizeString = "5m";
+    const expectedSize = 5242880;
+    expect(parseIECSize(sizeString)).toBe(expectedSize);
+  });
+
+  // Test case: Valid size string with unit 'g'
+  test("should return the parsed size when given a valid size string with unit 'g'", () => {
+    const sizeString = "2g";
+    const expectedSize = 2147483648;
+    expect(parseIECSize(sizeString)).toBe(expectedSize);
+  });
+
+  // Test case: Valid size string with unit 't'
+  test("should return the parsed size when given a valid size string with unit 't'", () => {
+    const sizeString = "1t";
+    const expectedSize = 1099511627776;
+    expect(parseIECSize(sizeString)).toBe(expectedSize);
+  });
+
+  // Test case: Invalid size string
+  test("should throw an error when given an invalid size string", () => {
+    const sizeString = "abc";
+    expect(() => parseIECSize(sizeString)).toThrow("Invalid size string: abc");
+  });
+});