From 6f7c8a3e9ee1cc25be151993dc1d3da7bdeb363e Mon Sep 17 00:00:00 2001
From: Michael Santos <michael.santos@gmail.com>
Date: Sat, 24 Dec 2016 00:43:43 -0500
Subject: [PATCH] ezlib_drv: avoid overflow when OOM

---
 c_src/ezlib_drv.c   | 61 +++++++++++++++++++++++++++++++++++++--------
 rebar.config.script |  7 +++++-
 2 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/c_src/ezlib_drv.c b/c_src/ezlib_drv.c
index b60ab8b..9d22bbe 100644
--- a/c_src/ezlib_drv.c
+++ b/c_src/ezlib_drv.c
@@ -38,6 +38,45 @@ typedef struct {
       z_stream *i_stream;
 } ezlib_data;
 
+/* Wrappers around driver_alloc() that check  */
+/* for OOM.                                   */
+#ifdef HAS_ERTS_EXIT
+void erts_exit(int n, char* v, ...);
+#define erl_exit erts_exit
+#else
+void erl_exit(int n, char* v, ...) { abort(); }
+#endif
+
+void *ezlib_alloc(ErlDrvSizeT size);
+ErlDrvBinary *ezlib_alloc_binary(ErlDrvSizeT size);
+ErlDrvBinary *ezlib_realloc_binary(ErlDrvBinary *bin,
+    ErlDrvSizeT size);
+
+void *ezlib_alloc(ErlDrvSizeT size) {
+    void *p = driver_alloc(size);
+    if (p == NULL) {
+        erl_exit(1, "ezlib: Can't allocate %lu bytes of memory\n",
+            size);
+    }
+    return p;
+}
+
+ErlDrvBinary *ezlib_alloc_binary(ErlDrvSizeT size) {
+    ErlDrvBinary *p = driver_alloc_binary(size);
+    if (p == NULL) {
+        erl_exit(1, "ezlib: Can't allocate %lu binary\n", size);
+    }
+    return p;
+}
+
+ErlDrvBinary *ezlib_realloc_binary(ErlDrvBinary *bin, ErlDrvSizeT size) {
+    ErlDrvBinary *p = driver_realloc_binary(bin, size);
+    if (p == NULL) {
+        erl_exit(1, "ezlib: Can't reallocate %lu binary\n", size);
+    }
+    return p;
+}
+
 static void* zlib_alloc(void* data, unsigned int items, unsigned int size)
 {
     return (void*) driver_alloc(items*size);
@@ -51,10 +90,10 @@ static void zlib_free(void* data, void* addr)
 static ErlDrvData ezlib_drv_start(ErlDrvPort port, char *buff)
 {
    ezlib_data *d =
-      (ezlib_data *)driver_alloc(sizeof(ezlib_data));
+      ezlib_alloc(sizeof(ezlib_data));
    d->port = port;
 
-   d->d_stream = (z_stream *)driver_alloc(sizeof(z_stream));
+   d->d_stream = ezlib_alloc(sizeof(z_stream));
 
    d->d_stream->zalloc = zlib_alloc;
    d->d_stream->zfree = zlib_free;
@@ -62,7 +101,7 @@ static ErlDrvData ezlib_drv_start(ErlDrvPort port, char *buff)
 
    deflateInit(d->d_stream, Z_DEFAULT_COMPRESSION);
 
-   d->i_stream = (z_stream *)driver_alloc(sizeof(z_stream));
+   d->i_stream = ezlib_alloc(sizeof(z_stream));
 
    d->i_stream->zalloc = zlib_alloc;
    d->i_stream->zfree = zlib_free;
@@ -96,7 +135,7 @@ static void ezlib_drv_stop(ErlDrvData handle)
 	 if (!(cond))						\
 	 {							\
 	    rlen = strlen(errstr) + 1;				\
-	    b = driver_realloc_binary(b, rlen);			\
+	    b = ezlib_realloc_binary(b, rlen);			\
 	    b->orig_bytes[0] = 1;				\
 	    strncpy(b->orig_bytes + 1, errstr, rlen - 1);	\
 	    *rbuf = (char *)b;					\
@@ -119,7 +158,7 @@ static ErlDrvSSizeT ezlib_drv_control(ErlDrvData handle,
       case DEFLATE:
 	 size = BUF_SIZE + 1;
 	 rlen = 1;
-	 b = driver_alloc_binary(size);
+	 b = ezlib_alloc_binary(size);
 	 b->orig_bytes[0] = 0;
 
 	 d->d_stream->next_in = (unsigned char *)buf;
@@ -145,15 +184,15 @@ static ErlDrvSSizeT ezlib_drv_control(ErlDrvData handle,
 
 	    rlen += (BUF_SIZE - d->d_stream->avail_out);
 	    size += (BUF_SIZE - d->d_stream->avail_out);
-	    b = driver_realloc_binary(b, size);
+	    b = ezlib_realloc_binary(b, size);
 	 }
-	 b = driver_realloc_binary(b, rlen);
+	 b = ezlib_realloc_binary(b, rlen);
 	 *rbuf = (char *)b;
 	 return rlen;
       case INFLATE:
 	 size = BUF_SIZE + 1;
 	 rlen = 1;
-	 b = driver_alloc_binary(size);
+	 b = ezlib_alloc_binary(size);
 	 b->orig_bytes[0] = 0;
 
 	 if (len > 0) {
@@ -180,15 +219,15 @@ static ErlDrvSSizeT ezlib_drv_control(ErlDrvData handle,
 
 	       rlen += (BUF_SIZE - d->i_stream->avail_out);
 	       size += (BUF_SIZE - d->i_stream->avail_out);
-	       b = driver_realloc_binary(b, size);
+	       b = ezlib_realloc_binary(b, size);
 	    }
 	 }
-	 b = driver_realloc_binary(b, rlen);
+	 b = ezlib_realloc_binary(b, rlen);
 	 *rbuf = (char *)b;
 	 return rlen;
    }
 
-   b = driver_alloc_binary(1);
+   b = ezlib_alloc_binary(1);
    b->orig_bytes[0] = 0;
    *rbuf = (char *)b;
    return 1;
diff --git a/rebar.config.script b/rebar.config.script
index a7a91ae..cf1baff 100644
--- a/rebar.config.script
+++ b/rebar.config.script
@@ -46,7 +46,12 @@ ModCfg0 = fun(F, Cfg, [Key|Tail], Op, Default) ->
 ModCfg = fun(Cfg, Keys, Op, Default) -> ModCfg0(ModCfg0, Cfg, Keys, Op, Default) end.
 ModCfgS = fun(Cfg, Keys, Val) -> ModCfg0(ModCfg0, Cfg, Keys, fun(_V) -> Val end, "") end.
 
-Cfg0 = ModCfg(CONFIG, [port_env, "CFLAGS"], fun(V) -> V ++ " " ++ CfgCFlags end, "$CFLAGS"),
+SysVersion = lists:map(fun erlang:list_to_integer/1,
+               string:tokens(erlang:system_info(version), ".")),
+
+ExitFlag = case SysVersion >= [7, 3] of true -> "-DHAS_ERTS_EXIT"; _ -> "" end.
+
+Cfg0 = ModCfg(CONFIG, [port_env, "CFLAGS"], fun(V) -> V ++ " " ++ ExitFlag ++ " " ++ CfgCFlags end, "$CFLAGS"),
 Cfg00 = ModCfg(Cfg0, [port_env, "LDFLAGS"], fun(V) -> V ++ " " ++ CfgLDFlags end, "$LDFLAGS"),
 Cfg1 = case CfgWithGCov of
            "true" ->