From bb7b4196f2571d806407173092dc779b3f5dfa35 Mon Sep 17 00:00:00 2001
From: Yorhel <git@yorhel.nl>
Date: Tue, 23 Jul 2019 11:01:48 +0200
Subject: [PATCH] Handle malloc failure by pausing any activity

This is a best-effort approach to save ncdu state when memory is low.
There's likely allocation in libraries that isn't being checked
(ncurses, printf).

Fixes #132 (it actually doesn't, that needs a 64bit static binary too,
but I'll get to that)
---
 src/dir_common.c |  6 +++---
 src/dir_import.c |  4 ++--
 src/dir_mem.c    |  2 +-
 src/dir_scan.c   |  6 +++---
 src/dirlist.c    |  2 +-
 src/exclude.c    |  8 ++++----
 src/global.h     |  3 +++
 src/main.c       | 15 ++++++++++-----
 src/path.c       | 24 ++++++++++++------------
 src/path.h       |  1 -
 src/util.c       | 25 ++++++++++++++++++++++---
 src/util.h       |  9 +++++++--
 12 files changed, 68 insertions(+), 37 deletions(-)

diff --git a/src/dir_common.c b/src/dir_common.c
index ee80c81..3a2a9d0 100644
--- a/src/dir_common.c
+++ b/src/dir_common.c
@@ -44,7 +44,7 @@ static int lasterrl; /* ^ of lasterr */
 static void curpath_resize(int s) {
   if(curpathl < s) {
     curpathl = s < 128 ? 128 : s < curpathl*2 ? curpathl*2 : s;
-    dir_curpath = realloc(dir_curpath, curpathl);
+    dir_curpath = xrealloc(dir_curpath, curpathl);
   }
 }
 
@@ -85,7 +85,7 @@ void dir_setlasterr(const char *path) {
   int req = strlen(path)+1;
   if(lasterrl < req) {
     lasterrl = req;
-    lasterr = realloc(lasterr, lasterrl);
+    lasterr = xrealloc(lasterr, lasterrl);
   }
   strcpy(lasterr, path);
 }
@@ -99,7 +99,7 @@ void dir_seterr(const char *fmt, ...) {
 
   va_list va;
   va_start(va, fmt);
-  dir_fatalerr = malloc(1024); /* Should be enough for everything... */
+  dir_fatalerr = xmalloc(1024); /* Should be enough for everything... */
   vsnprintf(dir_fatalerr, 1023, fmt, va);
   dir_fatalerr[1023] = 0;
   va_end(va);
diff --git a/src/dir_import.c b/src/dir_import.c
index 13ebe64..01b987d 100644
--- a/src/dir_import.c
+++ b/src/dir_import.c
@@ -596,12 +596,12 @@ int dir_import_init(const char *fn) {
   else if((stream = fopen(fn, "r")) == NULL)
     return 1;
 
-  ctx = malloc(sizeof(struct ctx));
+  ctx = xmalloc(sizeof(struct ctx));
   ctx->stream = stream;
   ctx->line = 1;
   ctx->byte = ctx->eof = ctx->items = 0;
   ctx->buf = ctx->lastfill = ctx->readbuf;
-  ctx->buf_dir = malloc(dir_memsize(""));
+  ctx->buf_dir = xmalloc(dir_memsize(""));
   ctx->readbuf[0] = 0;
 
   dir_curpath_set(fn);
diff --git a/src/dir_mem.c b/src/dir_mem.c
index f1a0742..54783e8 100644
--- a/src/dir_mem.c
+++ b/src/dir_mem.c
@@ -123,7 +123,7 @@ static int item(struct dir *dir, const char *name, struct dir_ext *ext) {
 
   if(!extended_info)
     dir->flags &= ~FF_EXT;
-  item = malloc(dir->flags & FF_EXT ? dir_ext_memsize(name) : dir_memsize(name));
+  item = xmalloc(dir->flags & FF_EXT ? dir_ext_memsize(name) : dir_memsize(name));
   memcpy(item, dir, offsetof(struct dir, name));
   strcpy(item->name, name);
   if(dir->flags & FF_EXT)
diff --git a/src/dir_scan.c b/src/dir_scan.c
index 442c4a2..6f65aab 100644
--- a/src/dir_scan.c
+++ b/src/dir_scan.c
@@ -100,7 +100,7 @@ static char *dir_read(int *err) {
     return NULL;
   }
 
-  buf = malloc(buflen);
+  buf = xmalloc(buflen);
   errno = 0;
 
   while((item = readdir(dir)) != NULL) {
@@ -109,7 +109,7 @@ static char *dir_read(int *err) {
     int req = off+3+strlen(item->d_name);
     if(req > buflen) {
       buflen = req < buflen*2 ? buflen*2 : req;
-      buf = realloc(buf, buflen);
+      buf = xrealloc(buf, buflen);
     }
     strcpy(buf+off, item->d_name);
     off += strlen(item->d_name)+1;
@@ -313,6 +313,6 @@ void dir_scan_init(const char *path) {
   dir_seterr(NULL);
   dir_process = process;
   if (!buf_dir)
-    buf_dir = malloc(dir_memsize(""));
+    buf_dir = xmalloc(dir_memsize(""));
   pstate = ST_CALC;
 }
diff --git a/src/dirlist.c b/src/dirlist.c
index 2f8f323..86cf4fe 100644
--- a/src/dirlist.c
+++ b/src/dirlist.c
@@ -217,7 +217,7 @@ void dirlist_open(struct dir *d) {
   /* set the reference to the parent dir */
   if(d->parent) {
     if(!parent_alloc)
-      parent_alloc = calloc(1, dir_memsize(".."));
+      parent_alloc = xcalloc(1, dir_memsize(".."));
     dirlist_parent = parent_alloc;
     strcpy(dirlist_parent->name, "..");
     dirlist_parent->next = head;
diff --git a/src/exclude.c b/src/exclude.c
index f417c92..460543a 100644
--- a/src/exclude.c
+++ b/src/exclude.c
@@ -23,7 +23,7 @@
 
 */
 
-#include "exclude.h"
+#include "global.h"
 
 #include <stdio.h>
 #include <stdlib.h>
@@ -45,8 +45,8 @@ void exclude_add(char *pat) {
   while(*n != NULL)
     n = &((*n)->next);
 
-  *n = (struct exclude *) calloc(1, sizeof(struct exclude));
-  (*n)->pattern = (char *) malloc(strlen(pat)+1);
+  *n = (struct exclude *) xcalloc(1, sizeof(struct exclude));
+  (*n)->pattern = (char *) xmalloc(strlen(pat)+1);
   strcpy((*n)->pattern, pat);
 }
 
@@ -125,7 +125,7 @@ int has_cachedir_tag(const char *name) {
     /* We don't need to copy the content of `path`, so it's more efficient to
      * use `free` + `malloc`. */
     free(path);
-    path = malloc(path_l);
+    path = xmalloc(path_l);
   }
   snprintf(path, path_l, "%s/%s", name, CACHEDIR_TAG_FILENAME);
   f = fopen(path, "rb");
diff --git a/src/global.h b/src/global.h
index 90b804a..b544ea7 100644
--- a/src/global.h
+++ b/src/global.h
@@ -109,6 +109,9 @@ extern int follow_symlinks;
 /* handle input from keyboard and update display */
 int input_handle(int);
 
+/* de-initialize ncurses */
+void close_nc();
+
 
 /* import all other global functions and variables */
 #include "browser.h"
diff --git a/src/main.c b/src/main.c
index a9f4bdc..6ce951d 100644
--- a/src/main.c
+++ b/src/main.c
@@ -280,6 +280,15 @@ static void init_nc() {
 }
 
 
+void close_nc() {
+  if(ncurses_init) {
+    erase();
+    refresh();
+    endwin();
+  }
+}
+
+
 /* main program */
 int main(int argc, char **argv) {
   read_locale();
@@ -309,11 +318,7 @@ int main(int argc, char **argv) {
       break;
   }
 
-  if(ncurses_init) {
-    erase();
-    refresh();
-    endwin();
-  }
+  close_nc();
   exclude_clear();
 
   return 0;
diff --git a/src/path.c b/src/path.c
index 87227f3..0a5f637 100644
--- a/src/path.c
+++ b/src/path.c
@@ -23,7 +23,7 @@
 
 */
 
-#include "path.h"
+#include "global.h"
 
 #include <string.h>
 #include <stdio.h>
@@ -64,8 +64,8 @@ static int path_split(char *cur, char ***res) {
     }
 
   /* create array of the components */
-  old = malloc((j+1)*sizeof(char *));
-  *res = malloc((j+1)*sizeof(char *));
+  old = xmalloc((j+1)*sizeof(char *));
+  *res = xmalloc((j+1)*sizeof(char *));
   for(i=j=0; i<n; i++)
     if(i == 0 || (cur[i-1] == 0 && cur[i] != 0))
       old[j++] = cur+i;
@@ -98,11 +98,11 @@ static char *path_absolute(const char *path) {
   /* not an absolute path? prepend cwd */
   if(path[0] != '/') {
     n = RPATH_CNKSZ;
-    ret = malloc(n);
+    ret = xmalloc(n);
     errno = 0;
     while(!getcwd(ret, n) && errno == ERANGE) {
       n += RPATH_CNKSZ;
-      ret = realloc(ret, n);
+      ret = xrealloc(ret, n);
       errno = 0;
     }
     if(errno) {
@@ -112,12 +112,12 @@ static char *path_absolute(const char *path) {
 
     i = strlen(path) + strlen(ret) + 2;
     if(i > n)
-      ret = realloc(ret, i);
+      ret = xrealloc(ret, i);
     strcat(ret, "/");
     strcat(ret, path);
   /* otherwise, just make a copy */
   } else {
-    ret = malloc(strlen(path)+1);
+    ret = xmalloc(strlen(path)+1);
     strcpy(ret, path);
   }
   return ret;
@@ -133,7 +133,7 @@ static char *path_real_rec(char *cur, int *links) {
   char **arr, *tmp, *lnk = NULL, *ret = NULL;
 
   tmpl = strlen(cur)+1;
-  tmp = malloc(tmpl);
+  tmp = xmalloc(tmpl);
 
   /* split path */
   i = path_split(cur, &arr);
@@ -142,7 +142,7 @@ static char *path_real_rec(char *cur, int *links) {
   strcpy(tmp, "/");
   if(i > 0) {
     lnkl = RPATH_CNKSZ;
-    lnk = malloc(lnkl);
+    lnk = xmalloc(lnkl);
     if(chdir("/") < 0)
       goto path_real_done;
   }
@@ -153,7 +153,7 @@ static char *path_real_rec(char *cur, int *links) {
     /* check for symlink */
     while((n = readlink(arr[i], lnk, lnkl)) == lnkl || (n < 0 && errno == ERANGE)) {
       lnkl += RPATH_CNKSZ;
-      lnk = realloc(lnk, lnkl);
+      lnk = xrealloc(lnk, lnkl);
     }
     if(n < 0 && errno != EINVAL)
       goto path_real_done;
@@ -168,7 +168,7 @@ static char *path_real_rec(char *cur, int *links) {
         n += strlen(tmp);
       if(tmpl < n) {
         tmpl = n;
-        tmp = realloc(tmp, tmpl);
+        tmp = xrealloc(tmp, tmpl);
       }
       if(lnk[0] != '/')
         strcat(tmp, lnk);
@@ -179,7 +179,7 @@ static char *path_real_rec(char *cur, int *links) {
         n += strlen(arr[i])+1;
         if(tmpl < n) {
           tmpl = n;
-          tmp = realloc(tmp, tmpl);
+          tmp = xrealloc(tmp, tmpl);
         }
         strcat(tmp, "/");
         strcat(tmp, arr[i]);
diff --git a/src/path.h b/src/path.h
index 6ed9a35..e9d91cc 100644
--- a/src/path.h
+++ b/src/path.h
@@ -30,7 +30,6 @@
   - path_real uses chdir(), so it's not thread safe
   - Process requires +x access for all directory components
   - Potentionally slow
-  - Doesn't check return value of malloc() and realloc()
   - path_real doesn't check for the existance of the last component
   - cwd is unreliable after path_real
 */
diff --git a/src/util.c b/src/util.c
index d99a42d..c62e9bd 100644
--- a/src/util.c
+++ b/src/util.c
@@ -29,6 +29,7 @@
 #include <stdlib.h>
 #include <ncurses.h>
 #include <stdarg.h>
+#include <unistd.h>
 #ifdef HAVE_LOCALE_H
 #include <locale.h>
 #endif
@@ -371,12 +372,12 @@ char *getpath(struct dir *cur) {
 
   if(datl == 0) {
     datl = i;
-    dat = malloc(i);
+    dat = xmalloc(i);
   } else if(datl < i) {
     datl = i;
-    dat = realloc(dat, i);
+    dat = xrealloc(dat, i);
   }
-  list = malloc(c*sizeof(struct dir *));
+  list = xmalloc(c*sizeof(struct dir *));
 
   c = 0;
   for(d=cur; d!=NULL; d=d->parent)
@@ -413,3 +414,21 @@ void addparentstats(struct dir *d, int64_t size, int64_t asize, uint64_t mtime,
     d = d->parent;
   }
 }
+
+
+/* Apparently we can just resume drawing after endwin() and ncurses will pick
+ * up where it left. Probably not very portable...  */
+#define oom_msg "\nOut of memory, press enter to try again or Ctrl-C to give up.\n"
+#define wrap_oom(f) \
+  void *ptr;\
+  char buf[128];\
+  while((ptr = f) == NULL) {\
+    close_nc();\
+    write(2, oom_msg, sizeof(oom_msg));\
+    read(0, buf, sizeof(buf));\
+  }\
+  return ptr;
+
+void *xmalloc(size_t size) { wrap_oom(malloc(size)) }
+void *xcalloc(size_t n, size_t size) { wrap_oom(calloc(n, size)) }
+void *xrealloc(void *mem, size_t size) { wrap_oom(realloc(mem, size)) }
diff --git a/src/util.h b/src/util.h
index 50344e3..f3c6925 100644
--- a/src/util.h
+++ b/src/util.h
@@ -170,13 +170,13 @@ void addparentstats(struct dir *, int64_t, int64_t, uint64_t, int);
 #define nstack_init(_s) do {\
     (_s)->size = 10;\
     (_s)->top = 0;\
-    (_s)->list = malloc(10*sizeof(*(_s)->list));\
+    (_s)->list = xmalloc(10*sizeof(*(_s)->list));\
   } while(0)
 
 #define nstack_push(_s, _v) do {\
     if((_s)->size <= (_s)->top) {\
       (_s)->size *= 2;\
-      (_s)->list = realloc((_s)->list, (_s)->size*sizeof(*(_s)->list));\
+      (_s)->list = xrealloc((_s)->list, (_s)->size*sizeof(*(_s)->list));\
     }\
     (_s)->list[(_s)->top++] = _v;\
   } while(0)
@@ -186,5 +186,10 @@ void addparentstats(struct dir *, int64_t, int64_t, uint64_t, int);
 #define nstack_free(_s) free((_s)->list)
 
 
+/* Malloc wrappers that exit on OOM */
+void *xmalloc();
+void *xcalloc(size_t, size_t);
+void *xrealloc(void *, size_t);
+
 #endif
 
-- 
GitLab