posix runtime: add malloc checks

2 files changed, 20 insertions, 7 deletions

diff --git a/runtime/POSIX/fd_init.c b/runtime/POSIX/fd_init.c
index 8845fc9c..a8d557e7 100644
--- a/runtime/POSIX/fd_init.c
+++ b/runtime/POSIX/fd_init.c
@@ -14,12 +14,9 @@
 #include "klee/klee.h"
 
 #include <assert.h>
-#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 #include <sys/stat.h>
-#include <sys/syscall.h>
-#include <unistd.h>
 
 exe_file_system_t __exe_fs;
 
@@ -46,6 +43,9 @@ exe_sym_env_t __exe_env = {
 static void __create_new_dfile(exe_disk_file_t *dfile, unsigned size, 
                                const char *name, struct stat64 *defaults) {
   struct stat64 *s = malloc(sizeof(*s));
+  if (!s)
+    klee_report_error(__FILE__, __LINE__, "out of memory in klee_init_env", "user.err");
+
   const char *sp;
   char sname[64];
   for (sp=name; *sp; ++sp)
@@ -56,6 +56,8 @@ static void __create_new_dfile(exe_disk_file_t *dfile, unsigned size,
 
   dfile->size = size;
   dfile->contents = malloc(dfile->size);
+  if (!dfile->contents)
+    klee_report_error(__FILE__, __LINE__, "out of memory in klee_init_env", "user.err");
   klee_make_symbolic(dfile->contents, dfile->size, name);
   
   klee_make_symbolic(s, sizeof(*s), sname);
@@ -118,6 +120,9 @@ void klee_init_fds(unsigned n_files, unsigned file_length,
 
   __exe_fs.n_sym_files = n_files;
   __exe_fs.sym_files = malloc(sizeof(*__exe_fs.sym_files) * n_files);
+  if (n_files && !__exe_fs.sym_files)
+    klee_report_error(__FILE__, __LINE__, "out of memory in klee_init_env", "user.err");
+
   for (k=0; k < n_files; k++) {
     name[0] = 'A' + k;
     __create_new_dfile(&__exe_fs.sym_files[k], file_length, name, &s);
@@ -126,6 +131,8 @@ void klee_init_fds(unsigned n_files, unsigned file_length,
   /* setting symbolic stdin */
   if (stdin_length) {
     __exe_fs.sym_stdin = malloc(sizeof(*__exe_fs.sym_stdin));
+    if (!__exe_fs.sym_stdin)
+      klee_report_error(__FILE__, __LINE__, "out of memory in klee_init_env", "user.err");
     __create_new_dfile(__exe_fs.sym_stdin, stdin_length, "stdin", &s);
     __exe_env.fds[0].dfile = __exe_fs.sym_stdin;
   }
@@ -138,6 +145,9 @@ void klee_init_fds(unsigned n_files, unsigned file_length,
     __exe_fs.close_fail = malloc(sizeof(*__exe_fs.close_fail));
     __exe_fs.ftruncate_fail = malloc(sizeof(*__exe_fs.ftruncate_fail));
     __exe_fs.getcwd_fail = malloc(sizeof(*__exe_fs.getcwd_fail));
+    if (!(__exe_fs.read_fail && __exe_fs.write_fail && __exe_fs.close_fail
+          && __exe_fs.ftruncate_fail && __exe_fs.getcwd_fail))
+      klee_report_error(__FILE__, __LINE__, "out of memory in klee_init_env", "user.err");
 
     klee_make_symbolic(__exe_fs.read_fail, sizeof(*__exe_fs.read_fail), "read_fail");
     klee_make_symbolic(__exe_fs.write_fail, sizeof(*__exe_fs.write_fail), "write_fail");
@@ -149,6 +159,8 @@ void klee_init_fds(unsigned n_files, unsigned file_length,
   /* setting symbolic stdout */
   if (sym_stdout_flag) {
     __exe_fs.sym_stdout = malloc(sizeof(*__exe_fs.sym_stdout));
+    if (!__exe_fs.sym_stdout)
+      klee_report_error(__FILE__, __LINE__, "out of memory in klee_init_env", "user.err");
     __create_new_dfile(__exe_fs.sym_stdout, 1024, "stdout", &s);
     __exe_env.fds[1].dfile = __exe_fs.sym_stdout;
     __exe_fs.stdout_writes = 0;
diff --git a/runtime/POSIX/klee_init_env.c b/runtime/POSIX/klee_init_env.c
index f45ddf3c..aaee4c4e 100644
--- a/runtime/POSIX/klee_init_env.c
+++ b/runtime/POSIX/klee_init_env.c
@@ -16,9 +16,6 @@
 #include <stdlib.h>
 #include <string.h>
 #include <assert.h>
-#include <errno.h>
-#include <sys/syscall.h>
-#include <unistd.h>
 
 static void __emit_error(const char *msg) {
   klee_report_error(__FILE__, __LINE__, msg, "user.err");
@@ -62,12 +59,14 @@ static int __streq(const char *a, const char *b) {
 static char *__get_sym_str(int numChars, char *name) {
   int i;
   char *s = malloc(numChars+1);
+  if (!s)
+    __emit_error("out of memory in klee_init_env");
   klee_mark_global(s);
   klee_make_symbolic(s, numChars+1, name);
 
   for (i=0; i<numChars; i++)
     klee_posix_prefer_cex(s, __isprint(s[i]));
-  
+
   s[numChars] = '\0';
   return s;
 }
@@ -223,6 +222,8 @@ usage: (klee_init_env) [options] [program arguments]\n\
   }
 
   final_argv = (char **)malloc((new_argc + 1) * sizeof(*final_argv));
+  if (!final_argv)
+    __emit_error("out of memory in klee_init_env");
   klee_mark_global(final_argv);
   memcpy(final_argv, new_argv, new_argc * sizeof(*final_argv));
   final_argv[new_argc] = 0;