From bd4dfccbf91143d98af5d896776c4bc1de75c12b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lukas=20Backstr=C3=B6m?= <lukas@erlang.org>
Date: Mon, 17 Feb 2025 09:27:56 +0100
Subject: [PATCH] zstd: Fix copy of dict binary to environment

If we don't make a copy of the binary before referencing it in
the dict and it is a heap binary, the binary will be released
upon the next GC and the dict will segfault.
---
 erts/emulator/nifs/common/zstd_nif.c | 57 +++++++++++++++++-----------
 erts/emulator/zstd/zstd.mk           |  9 ++++-
 lib/stdlib/test/zstd_SUITE.erl       |  3 ++
 3 files changed, 44 insertions(+), 25 deletions(-)

diff --git a/erts/emulator/nifs/common/zstd_nif.c b/erts/emulator/nifs/common/zstd_nif.c
index 92b7c7243b6f..9ddaa8c13699 100644
--- a/erts/emulator/nifs/common/zstd_nif.c
+++ b/erts/emulator/nifs/common/zstd_nif.c
@@ -820,29 +820,34 @@ static ERL_NIF_TERM create_cdict_nif(
     ErlNifBinary bin;
     int level;
 
-    if (!enif_inspect_iolist_as_binary(env, argv[0], &bin) ||
-        !enif_get_int(env, argv[1], &level)) {
+    if (!enif_is_binary(env, argv[0]) || !enif_get_int(env, argv[1], &level)) {
         return enif_make_badarg(env);
     } else {
-        ERL_NIF_TERM result;
-        ZstdDict *dict;
-        ZSTD_CDict *cdict = ZSTD_createCDict_advanced(bin.data, bin.size,
+        ERL_NIF_TERM result, binary_copy;
+        ZstdDict *dictp, dict;
+
+        dict.env = enif_alloc_env();
+        binary_copy = enif_make_copy(dict.env, argv[0]);
+
+        (void)enif_inspect_binary(dict.env, binary_copy, &bin);
+        
+        dict.c = ZSTD_createCDict_advanced(bin.data, bin.size,
                     ZSTD_dlm_byRef, ZSTD_dct_auto,
                     ZSTD_getCParams(level, 0, bin.size),
                     zstd_customMem);
 
-        if (!cdict)
+        if (!dict.c) {
+            enif_free_env(dict.env);
             return enif_make_tuple2(env, am_error,
                 enif_make_atom(env, "invalid_compress_dict"));
+        }
 
-        dict = enif_alloc_resource(compress_dict_type, sizeof(ZstdDict));
+        dictp = enif_alloc_resource(compress_dict_type, sizeof(ZstdDict));
 
-        dict->c = cdict;
-        dict->env = enif_alloc_env();
-        enif_make_binary(dict->env, &bin);
+        *dictp = dict;
 
-        result = enif_make_resource(env, (void *)dict);
-        enif_release_resource((void *)dict);
+        result = enif_make_resource(env, (void *)dictp);
+        enif_release_resource((void *)dictp);
         return result;
     }
 }
@@ -857,26 +862,32 @@ static ERL_NIF_TERM create_ddict_nif(
     ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[]) {
     ErlNifBinary bin;
 
-    if (!enif_inspect_iolist_as_binary(env, argv[0], &bin)) {
+    if (!enif_is_binary(env, argv[0])) {
         return enif_make_badarg(env);
     } else {
-        ERL_NIF_TERM result;
-        ZstdDict *dict;
-        ZSTD_DDict *ddict = ZSTD_createDDict_advanced(bin.data, bin.size,
+        ERL_NIF_TERM result, binary_copy;
+        ZstdDict *dictp, dict;
+        
+        dict.env = enif_alloc_env();
+        binary_copy = enif_make_copy(dict.env, argv[0]);
+
+        (void)enif_inspect_binary(dict.env, binary_copy, &bin);
+        
+        dict.d = ZSTD_createDDict_advanced(bin.data, bin.size,
                     ZSTD_dlm_byRef, ZSTD_dct_auto, zstd_customMem);
 
-        if (!ddict)
+        if (!dict.d) {
+            enif_free_env(dict.env);
             return enif_make_tuple2(env, am_error,
                 enif_make_atom(env, "invalid_decompress_dict"));
+        }
 
-        dict = enif_alloc_resource(decompress_dict_type, sizeof(ZstdDict));
+        dictp = enif_alloc_resource(decompress_dict_type, sizeof(ZstdDict));
 
-        dict->d = ddict;
-        dict->env = enif_alloc_env();
-        enif_make_binary(dict->env, &bin);
+        *dictp = dict;
 
-        result = enif_make_resource(env, (void *)dict);
-        enif_release_resource((void *)dict);
+        result = enif_make_resource(env, (void *)dictp);
+        enif_release_resource((void *)dictp);
         return result;
     }
 }
diff --git a/erts/emulator/zstd/zstd.mk b/erts/emulator/zstd/zstd.mk
index 518d656698ac..2362be95e21f 100644
--- a/erts/emulator/zstd/zstd.mk
+++ b/erts/emulator/zstd/zstd.mk
@@ -45,12 +45,17 @@ endif
 
 ifeq ($(TYPE),gcov)
 ZSTD_CFLAGS = -O0 -fprofile-arcs -ftest-coverage $(DEBUG_CFLAGS) $(DEFS) $(THR_DEFS)
-else  # gcov
+else  # !gcov
 ifeq ($(TYPE),debug)
 ## DEBUGLEVEL=1 enables asserts, see common/debug.h for details
 ZSTD_CFLAGS = -DDEBUGLEVEL=1 $(DEBUG_CFLAGS) $(DEFS) $(THR_DEFS)
-else # debug
+else # !debug && !gcov
+
 ZSTD_CFLAGS = $(subst -O2, -O3, $(CONFIGURE_CFLAGS) $(DEFS) $(THR_DEFS))
+ifeq ($(TYPE), asan)
+ZSTD_CFLAGS += -DZSTD_ADDRESS_SANITIZER
+endif # asan
+
 endif # debug
 endif # gcov
 
diff --git a/lib/stdlib/test/zstd_SUITE.erl b/lib/stdlib/test/zstd_SUITE.erl
index 41d320e69b60..ed6e34d307f0 100644
--- a/lib/stdlib/test/zstd_SUITE.erl
+++ b/lib/stdlib/test/zstd_SUITE.erl
@@ -437,6 +437,9 @@ dict_api(Config) ->
     {ok, DCtx} = zstd:context(decompress, #{ dictionary => DDict }),
     {'EXIT', _} = catch zstd:set_parameter(DCtx, dictionary, CDict),
 
+    {'EXIT', _} = catch zstd:dict(compress, [1,2,3]),
+    {'EXIT', _} = catch zstd:dict(decompress, [1,2,3]),
+
     ok.