[PATCH] i386-pc: build btrfs zstd support into separate module

[Michael Chang]

The zstd support in btrfs brings significant size increment to the
on-disk image that it can no longer fit into btrfs bootloader area and
short mbr gap.

In order to support grub update on outstanding i386-pc setup with these
size constraints remain in place, here we build the zstd suppprt of
btrfs into a separate module, named btrfs_zstd, to alleviate the size
change. Please note this only makes it's way to i386-pc, other
architecture is not affected.

Therefore if the system has enough space of embedding area for grub then
zstd support for btrfs will be enabled automatically in the process of
running grub-install through inserting btrfs_zstd module to the on-disk
image, otherwise a warning will be logged on screen to indicate user
that zstd support for btrfs is disabled due to the size limit.

Signed-off-by: Michael Chang <mchang@suse.com>
---
 Makefile.util.def           |   1 +
 grub-core/Makefile.core.def |  11 ++++
 grub-core/fs/btrfs.c        | 114 +++++---------------------------
 grub-core/fs/btrfs_zstd.c   |  36 +++++++++++
 grub-core/lib/zstd.c        | 126 ++++++++++++++++++++++++++++++++++++
 include/grub/btrfs.h        |   6 ++
 include/grub/lib/zstd.h     |  26 ++++++++
 include/grub/util/install.h |   2 +
 util/grub-install.c         |  20 ++++++
 util/setup.c                |  43 ++++++++++++
 10 files changed, 288 insertions(+), 97 deletions(-)
 create mode 100644 grub-core/fs/btrfs_zstd.c
 create mode 100644 grub-core/lib/zstd.c
 create mode 100644 include/grub/lib/zstd.h

diff --git a/Makefile.util.def b/Makefile.util.def
index 2f2881cb7..ac2b6aab1 100644
--- a/Makefile.util.def
+++ b/Makefile.util.def
@@ -176,6 +176,7 @@ library = {
   common = grub-core/lib/zstd/xxhash.c;
   common = grub-core/lib/zstd/zstd_common.c;
   common = grub-core/lib/zstd/zstd_decompress.c;
+  common = grub-core/lib/zstd.c;
 };
 
 program = {
diff --git a/grub-core/Makefile.core.def b/grub-core/Makefile.core.def
index dd1777bd3..b5328d7a0 100644
--- a/grub-core/Makefile.core.def
+++ b/grub-core/Makefile.core.def
@@ -1370,10 +1370,21 @@ module = {
   name = btrfs;
   common = fs/btrfs.c;
   common = lib/crc.c;
+  nopc = lib/zstd.c;
   cflags = '$(CFLAGS_POSIX) -Wno-undef';
+  i386_pc_cppflags = '-I$(srcdir)/lib/posix_wrap -I$(srcdir)/lib/minilzo 
-DMINILZO_HAVE_CONFIG_H';
   cppflags = '-I$(srcdir)/lib/posix_wrap -I$(srcdir)/lib/minilzo 
-I$(srcdir)/lib/zstd -DMINILZO_HAVE_CONFIG_H';
 };
 
+module = {
+  name = btrfs_zstd;
+  common = fs/btrfs_zstd.c;
+  common = lib/zstd.c;
+  cflags = '$(CFLAGS_POSIX) -Wno-undef';
+  cppflags = '-I$(srcdir)/lib/posix_wrap -I$(srcdir)/lib/zstd';
+  enable = i386_pc;
+};
+
 module = {
   name = archelp;
   common = fs/archelp.c;
diff --git a/grub-core/fs/btrfs.c b/grub-core/fs/btrfs.c
index c908b460f..9625bdf16 100644
--- a/grub-core/fs/btrfs.c
+++ b/grub-core/fs/btrfs.c
@@ -17,14 +17,6 @@
  *  along with GRUB.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-/*
- * Tell zstd to expose functions that aren't part of the stable API, which
- * aren't safe to use when linking against a dynamic library. We vendor in a
- * specific zstd version, so we know what we're getting. We need these unstable
- * functions to provide our own allocator, which uses grub_malloc(), to zstd.
- */
-#define ZSTD_STATIC_LINKING_ONLY
-
 #include <grub/err.h>
 #include <grub/file.h>
 #include <grub/mm.h>
@@ -35,7 +27,9 @@
 #include <grub/lib/crc.h>
 #include <grub/deflate.h>
 #include <minilzo.h>
-#include <zstd.h>
+#ifndef GRUB_MACHINE_PCBIOS
+#include <grub/lib/zstd.h>
+#endif
 #include <grub/i18n.h>
 #include <grub/btrfs.h>
 #include <grub/crypto.h>
@@ -61,12 +55,15 @@ GRUB_MOD_LICENSE ("GPLv3+");
 #define GRUB_BTRFS_LZO_BLOCK_MAX_CSIZE (GRUB_BTRFS_LZO_BLOCK_SIZE + \
                                     (GRUB_BTRFS_LZO_BLOCK_SIZE / 16) + 64 + 3)
 
-#define ZSTD_BTRFS_MAX_WINDOWLOG 17
-#define ZSTD_BTRFS_MAX_INPUT     (1 << ZSTD_BTRFS_MAX_WINDOWLOG)
-
 typedef grub_uint8_t grub_btrfs_checksum_t[0x20];
 typedef grub_uint16_t grub_btrfs_uuid_t[8];
 
+#ifdef GRUB_MACHINE_PCBIOS
+grub_ssize_t (*grub_btrfs_zstd_decompress_func) (char *ibuf,
+       grub_size_t isize, grub_off_t off,
+       char *obuf, grub_size_t osize) = NULL;
+#endif
+
 struct grub_btrfs_device
 {
   grub_uint64_t device_id;
@@ -1392,94 +1389,17 @@ grub_btrfs_read_inode (struct grub_btrfs_data *data,
   return grub_btrfs_read_logical (data, elemaddr, inode, sizeof (*inode), 0);
 }
 
-static void *grub_zstd_malloc (void *state __attribute__((unused)), size_t 
size)
-{
-  return grub_malloc (size);
-}
-
-static void grub_zstd_free (void *state __attribute__((unused)), void *address)
-{
-  return grub_free (address);
-}
-
-static ZSTD_customMem grub_zstd_allocator (void)
-{
-  ZSTD_customMem allocator;
-
-  allocator.customAlloc = &grub_zstd_malloc;
-  allocator.customFree = &grub_zstd_free;
-  allocator.opaque = NULL;
-
-  return allocator;
-}
-
-static grub_ssize_t
+static inline grub_ssize_t
 grub_btrfs_zstd_decompress (char *ibuf, grub_size_t isize, grub_off_t off,
                            char *obuf, grub_size_t osize)
 {
-  void *allocated = NULL;
-  char *otmpbuf = obuf;
-  grub_size_t otmpsize = osize;
-  ZSTD_DCtx *dctx = NULL;
-  grub_size_t zstd_ret;
-  grub_ssize_t ret = -1;
-
-  /*
-   * Zstd will fail if it can't fit the entire output in the destination
-   * buffer, so if osize isn't large enough, allocate a temporary buffer.
-   */
-  if (otmpsize < ZSTD_BTRFS_MAX_INPUT)
-    {
-      allocated = grub_malloc (ZSTD_BTRFS_MAX_INPUT);
-      if (!allocated)
-       {
-         grub_error (GRUB_ERR_OUT_OF_MEMORY, "failed allocate a zstd buffer");
-         goto err;
-       }
-      otmpbuf = (char *) allocated;
-      otmpsize = ZSTD_BTRFS_MAX_INPUT;
-    }
-
-  /* Create the ZSTD_DCtx. */
-  dctx = ZSTD_createDCtx_advanced (grub_zstd_allocator ());
-  if (!dctx)
-    {
-      /* ZSTD_createDCtx_advanced() only fails if it is out of memory. */
-      grub_error (GRUB_ERR_OUT_OF_MEMORY, "failed to create a zstd context");
-      goto err;
-    }
-
-  /*
-   * Get the real input size, there may be junk at the
-   * end of the frame.
-   */
-  isize = ZSTD_findFrameCompressedSize (ibuf, isize);
-  if (ZSTD_isError (isize))
-    {
-      grub_error (GRUB_ERR_BAD_COMPRESSED_DATA, "zstd data corrupted");
-      goto err;
-    }
-
-  /* Decompress and check for errors. */
-  zstd_ret = ZSTD_decompressDCtx (dctx, otmpbuf, otmpsize, ibuf, isize);
-  if (ZSTD_isError (zstd_ret))
-    {
-      grub_error (GRUB_ERR_BAD_COMPRESSED_DATA, "zstd data corrupted");
-      goto err;
-    }
-
-  /*
-   * Move the requested data into the obuf. obuf may be equal
-   * to otmpbuf, which is why grub_memmove() is required.
-   */
-  grub_memmove (obuf, otmpbuf + off, osize);
-  ret = osize;
-
-err:
-  grub_free (allocated);
-  ZSTD_freeDCtx (dctx);
-
-  return ret;
+#ifdef GRUB_MACHINE_PCBIOS
+  if (!grub_btrfs_zstd_decompress_func)
+    return -1;
+  return grub_btrfs_zstd_decompress_func (ibuf, isize, off, obuf, osize);
+#else
+  return grub_zstd_decompress (ibuf, isize, off, obuf, osize);
+#endif
 }
 
 static grub_ssize_t
diff --git a/grub-core/fs/btrfs_zstd.c b/grub-core/fs/btrfs_zstd.c
new file mode 100644
index 000000000..d5d1e013c
--- /dev/null
+++ b/grub-core/fs/btrfs_zstd.c
@@ -0,0 +1,36 @@
+/*
+ *  GRUB  --  GRand Unified Bootloader
+ *  Copyright (C) 2008  Free Software Foundation, Inc.
+ *
+ *  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/dl.h>
+/* For NULL.  */
+#include <grub/mm.h>
+#include <grub/btrfs.h>
+#include <grub/lib/zstd.h>
+
+GRUB_MOD_LICENSE ("GPLv3+");
+
+GRUB_MOD_INIT (btrfs_zstd)
+{
+  grub_btrfs_zstd_decompress_func = grub_zstd_decompress;
+}
+
+GRUB_MOD_FINI (btrfs_zstd)
+{
+  grub_btrfs_zstd_decompress_func = NULL;
+}
diff --git a/grub-core/lib/zstd.c b/grub-core/lib/zstd.c
new file mode 100644
index 000000000..643e90d84
--- /dev/null
+++ b/grub-core/lib/zstd.c
@@ -0,0 +1,126 @@
+/*
+ *  GRUB  --  GRand Unified Bootloader
+ *  Copyright (C) 2008  Free Software Foundation, Inc.
+ *
+ *  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/>.
+ */
+
+/*
+ * Tell zstd to expose functions that aren't part of the stable API, which
+ * aren't safe to use when linking against a dynamic library. We vendor in a
+ * specific zstd version, so we know what we're getting. We need these unstable
+ * functions to provide our own allocator, which uses grub_malloc(), to zstd.
+ */
+#define ZSTD_STATIC_LINKING_ONLY
+
+#include <grub/types.h>
+#include <grub/misc.h>
+#include <grub/err.h>
+#include <grub/mm.h>
+#include <grub/lib/zstd.h>
+#include <zstd.h>
+
+#define ZSTD_MAX_WINDOWLOG 17
+#define ZSTD_MAX_INPUT     (1 << ZSTD_MAX_WINDOWLOG)
+
+static void *grub_zstd_malloc (void *state __attribute__((unused)), size_t 
size)
+{
+  return grub_malloc (size);
+}
+
+static void grub_zstd_free (void *state __attribute__((unused)), void *address)
+{
+  return grub_free (address);
+}
+
+static ZSTD_customMem grub_zstd_allocator (void)
+{
+  ZSTD_customMem allocator;
+
+  allocator.customAlloc = &grub_zstd_malloc;
+  allocator.customFree = &grub_zstd_free;
+  allocator.opaque = NULL;
+
+  return allocator;
+}
+
+grub_ssize_t
+grub_zstd_decompress (char *ibuf, grub_size_t isize, grub_off_t off,
+                char *obuf, grub_size_t osize)
+{
+  void *allocated = NULL;
+  char *otmpbuf = obuf;
+  grub_size_t otmpsize = osize;
+  ZSTD_DCtx *dctx = NULL;
+  grub_size_t zstd_ret;
+  grub_ssize_t ret = -1;
+
+  /*
+   * Zstd will fail if it can't fit the entire output in the destination
+   * buffer, so if osize isn't large enough, allocate a temporary buffer.
+   */
+  if (otmpsize < ZSTD_MAX_INPUT)
+    {
+      allocated = grub_malloc (ZSTD_MAX_INPUT);
+      if (!allocated)
+       {
+         grub_error (GRUB_ERR_OUT_OF_MEMORY, "failed allocate a zstd buffer");
+         goto err;
+       }
+      otmpbuf = (char *) allocated;
+      otmpsize = ZSTD_MAX_INPUT;
+    }
+
+  /* Create the ZSTD_DCtx. */
+  dctx = ZSTD_createDCtx_advanced (grub_zstd_allocator ());
+  if (!dctx)
+    {
+      /* ZSTD_createDCtx_advanced() only fails if it is out of memory. */
+      grub_error (GRUB_ERR_OUT_OF_MEMORY, "failed to create a zstd context");
+      goto err;
+    }
+
+  /*
+   * Get the real input size, there may be junk at the
+   * end of the frame.
+   */
+  isize = ZSTD_findFrameCompressedSize (ibuf, isize);
+  if (ZSTD_isError (isize))
+    {
+      grub_error (GRUB_ERR_BAD_COMPRESSED_DATA, "zstd data corrupted");
+      goto err;
+    }
+
+  /* Decompress and check for errors. */
+  zstd_ret = ZSTD_decompressDCtx (dctx, otmpbuf, otmpsize, ibuf, isize);
+  if (ZSTD_isError (zstd_ret))
+    {
+      grub_error (GRUB_ERR_BAD_COMPRESSED_DATA, "zstd data corrupted");
+      goto err;
+    }
+
+  /*
+   * Move the requested data into the obuf. obuf may be equal
+   * to otmpbuf, which is why grub_memmove() is required.
+   */
+  grub_memmove (obuf, otmpbuf + off, osize);
+  ret = osize;
+
+err:
+  grub_free (allocated);
+  ZSTD_freeDCtx (dctx);
+
+  return ret;
+}
+
diff --git a/include/grub/btrfs.h b/include/grub/btrfs.h
index 234ad9767..f8e551982 100644
--- a/include/grub/btrfs.h
+++ b/include/grub/btrfs.h
@@ -69,4 +69,10 @@ struct grub_btrfs_inode_ref
   char name[0];
 };
 
+#ifdef GRUB_MACHINE_PCBIOS
+extern grub_ssize_t (*EXPORT_VAR (grub_btrfs_zstd_decompress_func)) (char 
*ibuf,
+               grub_size_t isize, grub_off_t off,
+               char *obuf, grub_size_t osize);
+#endif
+
 #endif
diff --git a/include/grub/lib/zstd.h b/include/grub/lib/zstd.h
new file mode 100644
index 000000000..0867b0c34
--- /dev/null
+++ b/include/grub/lib/zstd.h
@@ -0,0 +1,26 @@
+/*
+ *  GRUB  --  GRand Unified Bootloader
+ *  Copyright (C) 2008  Free Software Foundation, Inc.
+ *
+ *  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/>.
+ */
+
+#ifndef GRUB_ZSTD_HEADER
+#define GRUB_ZSTD_HEADER       1
+
+grub_ssize_t
+grub_zstd_decompress (char *ibuf, grub_size_t isize, grub_off_t off,
+                char *obuf, grub_size_t osize);
+
+#endif /* ! GRUB_ZSTD_HEADER */
diff --git a/include/grub/util/install.h b/include/grub/util/install.h
index faf9ff1ab..9e83e1339 100644
--- a/include/grub/util/install.h
+++ b/include/grub/util/install.h
@@ -208,6 +208,8 @@ grub_util_sparc_setup (const char *dir,
                       const char *dest, int force,
                       int fs_probe, int allow_floppy,
                       int add_rs_codes, int warn_short_mbr_gap);
+int
+grub_util_try_partmap_embed (const char *dest, unsigned int *nsec);
 
 char *
 grub_install_get_image_targets_string (void);
diff --git a/util/grub-install.c b/util/grub-install.c
index 891e3ced8..a2286b3dd 100644
--- a/util/grub-install.c
+++ b/util/grub-install.c
@@ -1482,6 +1482,26 @@ main (int argc, char *argv[])
        }
     }
 
+  if (install_drive
+      && platform == GRUB_INSTALL_PLATFORM_I386_PC
+      && grub_strcmp (grub_fs->name, "btrfs") == 0)
+    {
+      char *mod;
+      mod = grub_util_path_concat (2, grub_install_source_directory, 
"btrfs_zstd.mod");
+      if (grub_util_is_regular (mod))
+       {
+          unsigned int nsec = GRUB_MIN_RECOMMENDED_MBR_GAP;
+         int ret = grub_util_try_partmap_embed (install_drive, &nsec);
+         if (ret == 0)
+           grub_install_push_module ("btrfs_zstd");
+         else if (ret == 1)
+           grub_util_warn ("%s", _("btrfs zstd compression is disabled, please 
change install device to disk"));
+         else
+           grub_util_warn ("%s", _("btrfs zstd compression is disabled due to 
not enough space to embed image"));
+       }
+      grub_free (mod);
+    }
+
   if (!have_abstractions)
     {
       if ((disk_module && grub_strcmp (disk_module, "biosdisk") != 0)
diff --git a/util/setup.c b/util/setup.c
index 980e74374..6604d86bb 100644
--- a/util/setup.c
+++ b/util/setup.c
@@ -241,6 +241,49 @@ identify_partmap (grub_disk_t disk __attribute__ 
((unused)),
   return 1;
 }
 
+#ifdef GRUB_SETUP_BIOS
+int
+grub_util_try_partmap_embed (const char *dest, unsigned int *nsec)
+{
+  grub_device_t dest_dev;
+
+  dest_dev = grub_device_open (dest);
+  if (! dest_dev)
+    grub_util_error ("%s", grub_errmsg);
+
+  struct identify_partmap_ctx ctx = {
+    .dest_partmap = NULL,
+    .container = dest_dev->disk->partition,
+    .multiple_partmaps = 0
+  };
+
+  grub_partition_iterate (dest_dev->disk, identify_partmap, &ctx);
+
+  if (ctx.dest_partmap && ctx.dest_partmap->embed)
+    {
+      grub_err_t err;
+
+      grub_disk_addr_t *sectors = NULL;
+
+      err = ctx.dest_partmap->embed (dest_dev->disk, nsec, *nsec,
+                       GRUB_EMBED_PCBIOS, &sectors, 0);
+      if (err)
+       {
+         grub_errno = GRUB_ERR_NONE;
+          goto no_embed;
+       }
+      grub_free (sectors);
+      return 0;
+    }
+
+no_embed:
+  grub_device_close (dest_dev);
+  if (ctx.container)
+    return 1;
+  return 2;
+}
+#endif
+
 #ifdef GRUB_SETUP_BIOS
 #define SETUP grub_util_bios_setup
 #elif GRUB_SETUP_SPARC64
-- 
2.31.1