From 19404bf509f8968b3270fd9940bc218197bb481d Mon Sep 17 00:00:00 2001
From: Jakub Jelen <jjelen@redhat.com>
Date: Tue, 27 Jun 2023 14:53:41 +0200
Subject: [PATCH] bench: Add missing allocations checks

Signed-off-by: Jakub Jelen <jjelen@redhat.com>
Reviewed-by: Andreas Schneider <asn@cryptomilk.org>
Reviewed-by: Norbert Pocs <npocs@redhat.com>
---
 tests/benchmarks/bench_raw.c  | 16 ++++++++++++++--
 tests/benchmarks/bench_sftp.c |  3 +++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/tests/benchmarks/bench_raw.c b/tests/benchmarks/bench_raw.c
index db1a057c..0fd1446d 100644
--- a/tests/benchmarks/bench_raw.c
+++ b/tests/benchmarks/bench_raw.c
@@ -49,6 +49,9 @@ static char *get_python_eater(unsigned long bytes){
   char *ptr;
   char buf[12];
 
+  if (eater == NULL) {
+    return NULL;
+  }
   memcpy(eater,python_eater,sizeof(python_eater));
   ptr=strstr(eater,"XXXXXXXXXX");
   if(!ptr){
@@ -116,7 +119,10 @@ int benchmarks_raw_up (ssh_session session, struct argument_s *args,
   unsigned long total=0;
 
   bytes = args->datasize * 1024 * 1024;
-  script =get_python_eater(bytes);
+  script = get_python_eater(bytes);
+  if (script == NULL) {
+    return -1;
+  }
   err=upload_script(session,"/tmp/eater.py",script);
   free(script);
   if(err<0)
@@ -205,6 +211,9 @@ static char *get_python_giver(unsigned long bytes){
   char *ptr;
   char buf[12];
 
+  if (giver == NULL) {
+    return NULL;
+  }
   memcpy(giver,python_giver,sizeof(python_giver));
   ptr=strstr(giver,"XXXXXXXXXX");
   if(!ptr){
@@ -236,7 +245,10 @@ int benchmarks_raw_down (ssh_session session, struct argument_s *args,
   unsigned long total=0;
 
   bytes = args->datasize * 1024 * 1024;
-  script =get_python_giver(bytes);
+  script = get_python_giver(bytes);
+  if (script == NULL) {
+    return -1;
+  }
   err=upload_script(session,"/tmp/giver.py",script);
   free(script);
   if(err<0)
diff --git a/tests/benchmarks/bench_sftp.c b/tests/benchmarks/bench_sftp.c
index 601ecec0..d5766879 100644
--- a/tests/benchmarks/bench_sftp.c
+++ b/tests/benchmarks/bench_sftp.c
@@ -180,6 +180,9 @@ int benchmarks_async_sftp_down (ssh_session session, struct argument_s *args,
   if(!file)
     goto error;
   ids = malloc(concurrent_downloads * sizeof(int));
+  if (ids == NULL) {
+    return -1;
+  }
   if(args->verbose>0)
     fprintf(stdout,"Starting download of %lu bytes now, using %d concurrent downloads\n",bytes,
         concurrent_downloads);