Add some memory error handling.

2 files changed, 14 insertions, 3 deletions

diff --git a/loader2/nfsinstall.c b/loader2/nfsinstall.c
index 9ee5f093c..132465b98 100644
--- a/loader2/nfsinstall.c
+++ b/loader2/nfsinstall.c
@@ -218,6 +218,11 @@ char * mountNfsImage(struct installMethod * method,
 
             rc = asprintf(&buf, "%.*s/RHupdates", (int) (strrchr(fullPath, '/')-fullPath),
                           fullPath);
+            if (rc == -1) {
+                logMessage(CRITICAL, "Couldn't create RHupdates path, aborting");
+                abort();
+            }
+
             logMessage(INFO, "mounting nfs path %s for updates", buf);
 
             if (!doPwMount(buf, "/tmp/update-disk", "nfs", mountOpts)) {
diff --git a/loader2/urlinstall.c b/loader2/urlinstall.c
index 5da206d6a..ad06d3a6c 100644
--- a/loader2/urlinstall.c
+++ b/loader2/urlinstall.c
@@ -102,14 +102,18 @@ static void copyErrorFn (char *msg) {
 }
 
 static int loadUrlImages(struct iurlinfo * ui) {
-    char *buf, *path, *dest;
+    char *buf, *path, *dest, *slash;
     int rc;
 
     /* Figure out the path where updates.img and product.img files are
      * kept.  Since ui->prefix points to a stage2 image file, we just need
      * to trim off the file name and look in the same directory.
      */
-    path = strndup(ui->prefix, strrchr(ui->prefix, '/') - ui->prefix);
+    if ((slash = strrchr(ui->prefix, '/')) == NULL)
+        return 0;
+
+    if ((path = strndup(ui->prefix, slash - ui->prefix)) == NULL)
+        path = ui->prefix;
 
     /* grab the updates.img before stage2.img so that we minimize our
      * ramdisk usage */
@@ -410,7 +414,9 @@ void setKickstartUrl(struct loaderData_s * loaderData, int argc,
         return;
     }
 
-    loaderData->stage2Data = calloc(sizeof(struct urlInstallData *), 1);
+    if ((loaderData->stage2Data = calloc(sizeof(struct urlInstallData *), 1)) == NULL)
+        return;
+
     ((struct urlInstallData *)loaderData->stage2Data)->url = url;
 
     logMessage(INFO, "results of url ks, url %s", url);