[PATCH] commands: add new command memsize

[Xinhui Yang]

Greetings all,

I rewrote the program mentioned previously in an email[1], and it is
mostly complete, and it is useful for a number of scenarios, so I am
sending it here for upstreaming.

- This command traverses the entire GRUB memory map, adds the size of each
  region together. The result is the total amount of system memory
  available to GRUB in bytes.

- It stores the vaule to the environment as the name `memsize' by default,
  but it is customizable. The unit of the value is also customizable
  too.

- It checks if the vaule is too large for the `test' command to process,
  whose limit is INT_MAX.

- This command is useful with `test' command to implement mechanisms
  such as booting the live media with live filesystem loaded into RAN or
  not, instead of waiting initrd to fail:

  insmod memsize
  memsize -M --set=mem_avail
  if [ $mem_avail -gt 4000 ] ; then
    menuentry 'Live OS (Load into RAM)' {
      linux /boot/vmlinuz rd.live.ram=1 ...
      initrd /boot/initrd.img
    }
  fi

  It is also possible to implement mechanisms to boot into different
  mode (e.g. graphical or multi-user), depending on RAM available.

[1]: https://lists.gnu.org/archive/html/grub-devel/2023-07/msg00033.html

Regards,
Xinhui

* docs/grub.texi: add documentation for memsize

Signed-off-by: Xinhui Yang <cyan@cyano.uk>
---
 docs/grub.texi               |  63 ++++++++++
 grub-core/Makefile.core.def  |   5 +
 grub-core/commands/memsize.c | 227 +++++++++++++++++++++++++++++++++++
 3 files changed, 295 insertions(+)
 create mode 100644 grub-core/commands/memsize.c

diff --git a/docs/grub.texi b/docs/grub.texi
index 975e521d1..1dcd11b65 100644
--- a/docs/grub.texi
+++ b/docs/grub.texi
@@ -4372,6 +4372,7 @@ you forget a command, you can run the command 
@command{help}
 * ls::                          List devices or files
 * lsfonts::                     List loaded fonts
 * lsmod::                       Show loaded modules
+* memsize::                     Get the amount of available RAM
 * md5sum::                      Compute or check MD5 hash
 * module::                      Load module for multiboot kernel
 * multiboot::                   Load multiboot compliant kernel
@@ -5142,6 +5143,68 @@ List loaded fonts.
 Show list of loaded modules.
 @end deffn
 
+
+@node memsize
+@subsection memsize
+
+@deffn Command memsize 
[@option{-b}|@option{-k}|@option{-m}|@option{-g}|@option{-K}|@option{-M}|@option{-G}]
 [@option{-q}] [@option{--set} VARNAME]
+Iterate through GRUB memory map to get the amount of system RAM available to
+GRUB. The final integer value of the unit specified will be exported to the
+environment as @code{VARNAME} to make good use of the @pxref{test} command,
+and optionally be printed out to the console.
+
+The final value is checked to make sure it does not exceed the limit of the
+@pxref{test} command, which is INT_MAX. Otherwise, it raises an error.
+
+The default unit used for the value is MiB (@option{-M}).
+
+The default variable name is @code{memsize}.
+
+@option{-q} disables the output.
+
+@code{VARNAME} should be a valid identifier and should not exceed 63
+characters.
+
+The units can be:
+
+@table @code
+@item @option{-b}
+Display and export the value in number of bytes.
+
+@item @option{-k}
+Display and export the vaule in kilobytes (@code{kB}).
+
+@item @option{-m}
+Display and export the vaule in megabytes (@code{MB}).
+
+@item @option{-g}
+Display and export the vaule in gigabytes (@code{GB}).
+
+@item @option{-K}
+Display and export the vaule in kibibytes (@code{KiB}).
+
+@item @option{-M}
+Display and export the vaule in mebibytes (@code{MiB}).
+
+@item @option{-G}
+Display and export the vaule in gibibytes (@code{GiB}).
+@end table
+
+One can make use of this command to implement mechanisms like e.g. booting
+into different modes depending on available RAM:
+
+@example
+memsize -M --set=mem_avail
+if [ "$mem_avail" -lt 512 ] ; then
+  set boot_target="multi-user.target"
+else
+  set boot_target="graphical.target"
+fi
+@end example
+
+@end deffn
+
+
 @node md5sum
 @subsection md5sum
 
diff --git a/grub-core/Makefile.core.def b/grub-core/Makefile.core.def
index d2cf29584..3a938ad0e 100644
--- a/grub-core/Makefile.core.def
+++ b/grub-core/Makefile.core.def
@@ -1033,6 +1033,11 @@ module = {
   common = commands/memrw.c;
 };
 
+module = {
+  name = memsize;
+  common = commands/memsize.c;
+};
+
 module = {
   name = minicmd;
   common = commands/minicmd.c;
diff --git a/grub-core/commands/memsize.c b/grub-core/commands/memsize.c
new file mode 100644
index 000000000..05bed5c45
--- /dev/null
+++ b/grub-core/commands/memsize.c
@@ -0,0 +1,227 @@
+/*
+ *  memsize.c - Iterate through GRUB memory map to get the total amount of
+ *  available RAM. The memory map does NOT reflect the total physical
+ *  amount, so use with care.
+ *  It can export the value to the environment, to make great use of tests.
+ */
+/*
+ *  GRUB  --  GRand Unified Bootloader
+ *  Copyright (C) 2003,2007  Free Software Foundation, Inc.
+ *  Copyright (C) 2003  NIIBE Yutaka <gniibe@m17n.org>
+ *
+ *  GRUB is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  GRUB is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with GRUB.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <grub/types.h>
+#include <grub/misc.h>
+#include <grub/memory.h>
+#include <grub/err.h>
+#include <grub/env.h>
+#include <grub/dl.h>
+#include <grub/extcmd.h>
+#include <grub/i18n.h>
+
+GRUB_MOD_LICENSE ("GPLv3+");
+
+const char *default_varname = "memsize";
+grub_extcmd_t cmd;
+grub_uint64_t grub_avail_mem_bytes = 0ULL;
+
+static const struct grub_arg_option grub_options_memsize[] =
+  {
+    {"byte", 'b', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Print and export the vaule as number of bytes."), 0, 0},
+    {"kilo", 'k', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Print and export the value as kilobytes."), 0, 0},
+    {"mega", 'm', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Print and export the value as megabytes."), 0, 0},
+    {"giga", 'g', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Print and export the value as gigabytes."), 0, 0},
+    {"kibi", 'K', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Print and export the value as kibibytes."), 0, 0},
+    {"mebi", 'M', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Print and export the value as mebibytes."), 0, 0},
+    {"gibi", 'G', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Print and export the value as gibibytes."), 0, 0},
+    {"quiet", 'q', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Do not print anyting, just set a variable quietly."), 0, 0},
+    {"set", 's', GRUB_ARG_OPTION_OPTIONAL,
+      N_("Set a variable to the vaule as the unit specified."),
+      N_("VARNAME"), ARG_TYPE_STRING},
+    {0, 0, 0, 0, 0, 0},
+};
+
+/* Corresponding strings of the units */
+const char* unit_strs[] = {
+  "B", "kB", "MB", "GB",
+  "KiB", "MiB", "GiB",
+};
+
+/* Corresponding vaules to convert */
+const int conversions[] = {
+  1, 1e3, 1e6, 1e9,
+  (1 << 10), (1 << 20), (1 << 30),
+};
+
+enum grub_options_memsize
+  {
+    UNIT_BYTES,
+    UNIT_SI_KB,
+    UNIT_SI_MB,
+    UNIT_SI_GB,
+    UNIT_BIN_KIB,
+    UNIT_BIN_MIB,
+    UNIT_BIN_GIB,
+    FLAG_QUIET,
+    HAS_CUSTOM_VARIABLE,
+};
+
+static int
+memsize_hook (grub_uint64_t addr __attribute__ ((unused)),
+              grub_uint64_t size, grub_memory_type_t type, void *data)
+{
+  /*
+   * Iterate through GRUB memory map.
+   * FIXME: Which type of memory region should we consider to add up?
+   */
+  grub_uint64_t *total = (grub_uint64_t *) data;
+  if (type == GRUB_MEMORY_AVAILABLE) {
+      *total += size;
+  }
+  return GRUB_ERR_NONE;
+}
+
+static grub_err_t
+grub_cmd_memsize (grub_extcmd_context_t ctxt,
+                  int argc __attribute__ ((unused)),
+                  char **args __attribute__ ((unused)))
+{
+  /* Parsed argument list. */
+  struct grub_arg_list *state = ctxt->state;
+  /* Length of the variable name. */
+  grub_size_t len = 0;
+  /* Buffer of variable name. */
+  char buf[64] = {0};
+  /* Requested unit to convert. */
+  enum grub_options_memsize unit = UNIT_BIN_MIB;
+  enum grub_options_memsize opt = UNIT_BYTES;
+  /* Flag to make sure only one unit is specified. */
+  int unit_switch_is_set = 0;
+  /* Total amount of available RAM. */
+  grub_uint64_t avail_mem = 0ULL;
+  /* The converted value. */
+  grub_uint64_t converted_avail_mem = 0ULL;
+  char *converted_avail_mem_str = 0;
+  /* Record errors during grub_machine_mmap_iterate. */
+  grub_err_t err;
+  /* Sanity checks */
+  /* Only one unit switch is allowed. */
+  for (; opt < FLAG_QUIET; opt++) {
+    if (state[opt].set != 0) {
+      if (unit_switch_is_set != 0) {
+        return grub_error (GRUB_ERR_BAD_ARGUMENT,
+         N_("Only one unit switch is allowed"));
+      } else {
+        unit = opt;
+        unit_switch_is_set = 1;
+      }
+    }
+  }
+  /* Check if the specified variable name is valid. */
+  if (state[HAS_CUSTOM_VARIABLE].set != 0) {
+    char *str = state[HAS_CUSTOM_VARIABLE].arg;
+    len = grub_strlen (str);
+    if (len > 63) {
+      return grub_error (GRUB_ERR_BAD_ARGUMENT,
+        N_("Variable names should not exceed 63 characters"));
+    } else if (len == 0) {
+      return grub_error (GRUB_ERR_BAD_ARGUMENT,
+        N_("Variable name expected but not specified"));
+    }
+    /* [0-9a-zA-Z_] for variable names */
+    for (char *c = str; c < str + len; c++) {
+      if (   !(*c >= '0' && *c <= '9')
+          && !(*c >= 'a' && *c <= 'z')
+          && !(*c >= 'A' && *c <= 'Z')
+          && !(*c == '_')) {
+        return grub_error(GRUB_ERR_BAD_ARGUMENT,
+          N_("Invalid variable name"));
+      }
+    }
+    grub_strncpy (buf, str, len);
+  } else {
+    len = grub_strlen (default_varname);
+    grub_strncpy (buf, default_varname, len);
+  }
+  /* Iterate through GRUB's memory map. */
+#ifndef GRUB_MACHINE_EMU
+  err = grub_machine_mmap_iterate (memsize_hook, (void *) &avail_mem);
+  if (err != GRUB_ERR_NONE) {
+    return err;
+  }
+#endif
+  /*
+   * Convert the vaule to the specified unit, and export it to the
+   * environment.
+   */
+  converted_avail_mem = avail_mem / conversions[unit];
+  /*
+   * Unfortunately the test command can only perform comparisons across
+   * signed ints, or the test command will fail. We need to check this in
+   * order to make it usable to test command.
+   */
+  if (converted_avail_mem > GRUB_INT_MAX) {
+    return grub_error (GRUB_ERR_BAD_NUMBER,
+      N_("Value is too large (%"
+        PRIuGRUB_UINT64_T " > %" PRIdGRUB_INT32_T ")"),
+      converted_avail_mem, GRUB_INT_MAX);
+  }
+  if (state[FLAG_QUIET].set == 0) {
+    grub_printf (N_("The amount of available RAM is %"
+                    PRIuGRUB_UINT64_T " %s.\n"),
+                 converted_avail_mem, unit_strs[unit]);
+  }
+  converted_avail_mem_str =
+    grub_xasprintf ("%" PRIuGRUB_UINT64_T,
+                    converted_avail_mem);
+  if (!converted_avail_mem_str) {
+    return grub_error (GRUB_ERR_OUT_OF_MEMORY,
+        N_("Can't allocate space for value string"));
+  }
+
+  grub_env_set (buf, converted_avail_mem_str);
+  grub_env_export (buf);
+
+  return GRUB_ERR_NONE;
+}
+
+GRUB_MOD_INIT(memsize)
+{
+  cmd =
+    grub_register_extcmd ("memsize", grub_cmd_memsize,
+     GRUB_COMMAND_FLAG_EXTCMD,
+     N_("[-b|-k|-K|-m|-M|-g|-G] [-q] [--set VARNAME]"),
+     N_("Get the amount of system RAM available to GRUB, "
+        "and export the integer value to the environment. "
+        "The vaule will be printed out by default. "
+        "-q disables the output. "
+        "The default variable name is `memsize'. "
+        "The unit can be b for bytes, k, m, g for SI units, "
+        "or K, M, G for binary units. The default unit is MiB."),
+     grub_options_memsize);
+}
+
+GRUB_MOD_FINI(memsize) {
+  grub_unregister_extcmd (cmd);
+}
-- 
2.39.1