1 files changed, 61 insertions, 70 deletions

diff --git a/tools/testing/selftests/bpf/test_progs.c b/tools/testing/selftests/bpf/test_progs.c
index db00196c8315..12895d03d58b 100644
--- a/tools/testing/selftests/bpf/test_progs.c
+++ b/tools/testing/selftests/bpf/test_progs.c
@@ -40,14 +40,20 @@ static bool should_run(struct test_selector *sel, int num, const char *name)
 
 static void dump_test_log(const struct prog_test_def *test, bool failed)
 {
+        if (stdout == env.stdout)
+                return;
+
+        fflush(stdout); /* exports env.log_buf & env.log_cnt */
+
         if (env.verbose || test->force_log || failed) {
                 if (env.log_cnt) {
-                        fprintf(stdout, "%s", env.log_buf);
+                        fprintf(env.stdout, "%s", env.log_buf);
                         if (env.log_buf[env.log_cnt - 1] != '\n')
-                                fprintf(stdout, "\n");
+                                fprintf(env.stdout, "\n");
                 }
         }
-        env.log_cnt = 0;
+
+        fseeko(stdout, 0, SEEK_SET); /* rewind */
 }
 
 void test__end_subtest()
@@ -62,7 +68,7 @@ void test__end_subtest()
 
         dump_test_log(test, sub_error_cnt);
 
-        printf("#%d/%d %s:%s\n",
+        fprintf(env.stdout, "#%d/%d %s:%s\n",
                test->test_num, test->subtest_num,
                test->subtest_name, sub_error_cnt ? "FAIL" : "OK");
 }
@@ -79,7 +85,8 @@ bool test__start_subtest(const char *name)
         test->subtest_num++;
 
         if (!name || !name[0]) {
-                fprintf(stderr, "Subtest #%d didn't provide sub-test name!\n",
+                fprintf(env.stderr,
+                        "Subtest #%d didn't provide sub-test name!\n",
                         test->subtest_num);
                 return false;
         }
@@ -98,66 +105,6 @@ void test__force_log() {
         env.test->force_log = true;
 }
 
-void test__vprintf(const char *fmt, va_list args)
-{
-        size_t rem_sz;
-        int ret = 0;
-
-        if (env.verbose || (env.test && env.test->force_log)) {
-                vfprintf(stderr, fmt, args);
-                return;
-        }
-
-try_again:
-        rem_sz = env.log_cap - env.log_cnt;
-        if (rem_sz) {
-                va_list ap;
-
-                va_copy(ap, args);
-                /* we reserved extra byte for \0 at the end */
-                ret = vsnprintf(env.log_buf + env.log_cnt, rem_sz + 1, fmt, ap);
-                va_end(ap);
-
-                if (ret < 0) {
-                        env.log_buf[env.log_cnt] = '\0';
-                        fprintf(stderr, "failed to log w/ fmt '%s'\n", fmt);
-                        return;
-                }
-        }
-
-        if (!rem_sz || ret > rem_sz) {
-                size_t new_sz = env.log_cap * 3 / 2;
-                char *new_buf;
-
-                if (new_sz < 4096)
-                        new_sz = 4096;
-                if (new_sz < ret + env.log_cnt)
-                        new_sz = ret + env.log_cnt;
-
-                /* +1 for guaranteed space for terminating \0 */
-                new_buf = realloc(env.log_buf, new_sz + 1);
-                if (!new_buf) {
-                        fprintf(stderr, "failed to realloc log buffer: %d\n",
-                                errno);
-                        return;
-                }
-                env.log_buf = new_buf;
-                env.log_cap = new_sz;
-                goto try_again;
-        }
-
-        env.log_cnt += ret;
-}
-
-void test__printf(const char *fmt, ...)
-{
-        va_list args;
-
-        va_start(args, fmt);
-        test__vprintf(fmt, args);
-        va_end(args);
-}
-
 struct ipv4_packet pkt_v4 = {
         .eth.h_proto = __bpf_constant_htons(ETH_P_IP),
         .iph.ihl = 5,
@@ -331,7 +278,7 @@ enum ARG_KEYS {
         ARG_VERIFIER_STATS = 's',
         ARG_VERBOSE = 'v',
 };
-        
+
 static const struct argp_option opts[] = {
         { "num", ARG_TEST_NUM, "NUM", 0,
           "Run test number NUM only " },
@@ -349,7 +296,7 @@ static int libbpf_print_fn(enum libbpf_print_level level,
 {
         if (!env.very_verbose && level == LIBBPF_DEBUG)
                 return 0;
-        test__vprintf(format, args);
+        vprintf(format, args);
         return 0;
 }
 
@@ -477,6 +424,48 @@ static error_t parse_arg(int key, char *arg, struct argp_state *state)
         return 0;
 }
 
+static void stdio_hijack(void)
+{
+#ifdef __GLIBC__
+        env.stdout = stdout;
+        env.stderr = stderr;
+
+        if (env.verbose) {
+                /* nothing to do, output to stdout by default */
+                return;
+        }
+
+        /* stdout and stderr -> buffer */
+        fflush(stdout);
+
+        stdout = open_memstream(&env.log_buf, &env.log_cnt);
+        if (!stdout) {
+                stdout = env.stdout;
+                perror("open_memstream");
+                return;
+        }
+
+        stderr = stdout;
+#endif
+}
+
+static void stdio_restore(void)
+{
+#ifdef __GLIBC__
+        if (stdout == env.stdout)
+                return;
+
+        fclose(stdout);
+        free(env.log_buf);
+
+        env.log_buf = NULL;
+        env.log_cnt = 0;
+
+        stdout = env.stdout;
+        stderr = env.stderr;
+#endif
+}
+
 int main(int argc, char **argv)
 {
         static const struct argp argp = {
@@ -496,6 +485,7 @@ int main(int argc, char **argv)
 
         env.jit_enabled = is_jit_enabled();
 
+        stdio_hijack();
         for (i = 0; i < prog_test_cnt; i++) {
                 struct prog_test_def *test = &prog_test_defs[i];
                 int old_pass_cnt = pass_cnt;
@@ -523,13 +513,14 @@ int main(int argc, char **argv)
 
                 dump_test_log(test, test->error_cnt);
 
-                printf("#%d %s:%s\n", test->test_num, test->test_name,
-                       test->error_cnt ? "FAIL" : "OK");
+                fprintf(env.stdout, "#%d %s:%s\n",
+                        test->test_num, test->test_name,
+                        test->error_cnt ? "FAIL" : "OK");
         }
+        stdio_restore();
         printf("Summary: %d/%d PASSED, %d FAILED\n",
                env.succ_cnt, env.sub_succ_cnt, env.fail_cnt);
 
-        free(env.log_buf);
         free(env.test_selector.num_set);
         free(env.subtest_selector.num_set);