From b190e1b74b422b1488c1735ac3557e0e89377dcf Mon Sep 17 00:00:00 2001
From: Peter Collingbourne <peter@pcc.me.uk>
Date: Thu, 30 Oct 2014 13:23:01 +0000
Subject: [PATCH] Update __dfsw_s{,n}printf custom functions for new calling
 convention.

Differential Revision: http://reviews.llvm.org/D6029

llvm-svn: 220907
---
 compiler-rt/lib/dfsan/dfsan_custom.cc | 37 ++++++++++++---------------
 compiler-rt/test/dfsan/custom.cc      |  5 ++++
 2 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/compiler-rt/lib/dfsan/dfsan_custom.cc b/compiler-rt/lib/dfsan/dfsan_custom.cc
index ffd8b000bd9c..55f7170df455 100644
--- a/compiler-rt/lib/dfsan/dfsan_custom.cc
+++ b/compiler-rt/lib/dfsan/dfsan_custom.cc
@@ -908,7 +908,8 @@ static int format_chunk(char *str, size_t off, bool has_size, size_t size,
 // out which bytes of the output string depends on which argument and thus to
 // propagate labels more precisely.
 static int format_buffer(char *str, bool has_size, size_t size,
-                         const char *format, va_list ap) {
+                         const char *format, dfsan_label *va_labels,
+                         dfsan_label *ret_label, va_list ap) {
   InternalMmapVector<Chunk> chunks(8);
   size_t off = 0;
 
@@ -1055,16 +1056,8 @@ static int format_buffer(char *str, bool has_size, size_t size,
     off += status;
   }
 
-  // Consume the labels of the output buffer, (optional) size, and format
-  // string.
-  //
   // TODO(martignlo): Decide how to combine labels (e.g., whether to ignore or
   // not the label of the format string).
-  va_arg(ap, dfsan_label_va);
-  if (has_size) {
-    va_arg(ap, dfsan_label_va);
-  }
-  va_arg(ap, dfsan_label_va);
 
   // Label each output chunk according to the label supplied as argument to the
   // function. We need to go through all the chunks and arguments even if the
@@ -1077,17 +1070,17 @@ static int format_buffer(char *str, bool has_size, size_t size,
         dfsan_set_label(0, (void*) chunk.ptr, chunk.size);
         break;
       case Chunk::IGNORED:
-        va_arg(ap, dfsan_label_va);
+        va_labels++;
         dfsan_set_label(0, (void*) chunk.ptr, chunk.size);
         break;
       case Chunk::NUMERIC: {
-        dfsan_label label = va_arg(ap, dfsan_label_va);
+        dfsan_label label = *va_labels++;
         dfsan_set_label(label, (void*) chunk.ptr, chunk.size);
         break;
       }
       case Chunk::STRING: {
         // Consume the label of the pointer to the string
-        va_arg(ap, dfsan_label_va);
+        va_labels++;
         internal_memcpy(shadow_for((void *) chunk.ptr),
                         shadow_for((void *) chunk.arg),
                         sizeof(dfsan_label) * (strlen(chunk.arg)));
@@ -1096,27 +1089,31 @@ static int format_buffer(char *str, bool has_size, size_t size,
     }
   }
 
-  dfsan_label *ret_label_ptr = va_arg(ap, dfsan_label *);
-  *ret_label_ptr = 0;
+  *ret_label = 0;
 
   // Number of bytes written in total.
   return off;
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
-int __dfsw_sprintf(char *str, const char *format, ...) {
+int __dfsw_sprintf(char *str, const char *format, dfsan_label str_label,
+                   dfsan_label format_label, dfsan_label *va_labels,
+                   dfsan_label *ret_label, ...) {
   va_list ap;
-  va_start(ap, format);
-  int ret = format_buffer(str, false, 0, format, ap);
+  va_start(ap, ret_label);
+  int ret = format_buffer(str, false, 0, format, va_labels, ret_label, ap);
   va_end(ap);
   return ret;
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
-int __dfsw_snprintf(char *str, size_t size, const char *format, ...) {
+int __dfsw_snprintf(char *str, size_t size, const char *format,
+                    dfsan_label str_label, dfsan_label size_label,
+                    dfsan_label format_label, dfsan_label *va_labels,
+                    dfsan_label *ret_label, ...) {
   va_list ap;
-  va_start(ap, format);
-  int ret = format_buffer(str, true, size, format, ap);
+  va_start(ap, ret_label);
+  int ret = format_buffer(str, true, size, format, va_labels, ret_label, ap);
   va_end(ap);
   return ret;
 }
diff --git a/compiler-rt/test/dfsan/custom.cc b/compiler-rt/test/dfsan/custom.cc
index 8a49c32a0c76..d7bb3e3073fb 100644
--- a/compiler-rt/test/dfsan/custom.cc
+++ b/compiler-rt/test/dfsan/custom.cc
@@ -813,6 +813,11 @@ void test_sprintf() {
   assert(strcmp(buf, "Hello world!") == 0);
   ASSERT_READ_LABEL(buf, sizeof(buf), 0);
 
+  // Test for extra arguments.
+  assert(sprintf(buf, "Hello world!", 42, "hello") == 12);
+  assert(strcmp(buf, "Hello world!") == 0);
+  ASSERT_READ_LABEL(buf, sizeof(buf), 0);
+
   // Test formatting & label propagation (multiple conversion specifiers): %s,
   // %d, %n, %f, and %%.
   const char* s = "world";