Merge branch 'master' of https://github.com/vanhauser-thc/AFLplusplus

1 files changed, 18 insertions, 17 deletions

diff --git a/test/test-compcov.c b/test/test-compcov.c
index bbad3aed..978d5551 100644
--- a/test/test-compcov.c
+++ b/test/test-compcov.c
@@ -4,25 +4,26 @@
 #include <string.h>
 
 int main(int argc, char** argv) {
+  char *input = argv[1], *buf, buffer[20];
 
-  char *buf;
-
-  if (argc > 1) {
-  
-    if (strcmp(argv[1], "LIBTOKENCAP") == 0)
-      printf("your string was libtokencap\n");
-    else if (strcmp(argv[1], "BUGMENOT") == 0)
-      printf("your string was bugmenot\n");
-    else if (strcmp(argv[1], "BUFFEROVERFLOW") == 0) {
-      buf = malloc(16);
-      strcpy(buf, "TEST");
-      strcat(buf, argv[1]);
-      printf("This will only crash with libdislocator: %s\n", buf);
-      return 0;
-    } else
-      printf("I do not know your string\n");
-  
+  if (argc < 2) {
+    ssize_t ret = read(0, buffer, sizeof(buffer) - 1);
+    buffer[ret] = 0;
+    input = buffer;
   }
+  
+  if (strcmp(input, "LIBTOKENCAP") == 0)
+    printf("your string was libtokencap\n");
+  else if (strcmp(input, "BUGMENOT") == 0)
+    printf("your string was bugmenot\n");
+  else if (strcmp(input, "BUFFEROVERFLOW") == 0) {
+    buf = malloc(16);
+    strcpy(buf, "TEST");
+    strcat(buf, input);
+    printf("This will only crash with libdislocator: %s\n", buf);
+    return 0;
+  } else
+    printf("I do not know your string\n");
 
   return 0;