Extend s{,n}printf custom wrappers to support '*' in the format specifiers

[pcc]

lib/dfsan/dfsan_custom.cc

971

I feel a little uncomfortable about introducing more of these macros. Can we perhaps introduce a struct to hold the state used by these functions and make these macros and the functions into methods of the struct?

[pcc]

LGTM with my suggestions.

lib/dfsan/dfsan_custom.cc

1013

If you keep track of the start of the chunk of the format string instead of its size, you will have only one thing to increment here and in the other loop.

1015

You could modify this switch statement to consume characters instead:

switch (*++formatter.fmt) {

Once you make the matching change to build_format_string, this will save you from having to confusingly increment and decrement formatter.fmt above.

lib/dfsan/dfsan_custom.cc

if (has_size) {
     r = vsprintf(str + off, chunk_format, ap);
   }
   va_end(ap);
 
   free(chunk_format);
   return r;
 }
 
+// Returns the size of a formatted chunk.
+static size_t size_of_chunk(char *str, size_t off, bool has_size, size_t size,
+                            int status) {
+  if (status < 0) {
+    return 0;
+  }
+
+  size_t chunk_size = status;
+
+  // A return value of {v,}snprintf of size or more means that the output was
+  // truncated.
+  if (has_size) {
+    if (off < size) {
+      if (chunk_size >= size - off) {
+        chunk_size -= size - off;
+      }
+    } else {
+      chunk_size = 0;
+    }
+  }
+
+  return chunk_size;
+}
+
 // Formats the input and propagates the input labels to the output. The output
 // is stored in 'str'. If 'has_size' is true, 'size' bounds the number of
 // output bytes. 'format' and 'ap' are the format string and the list of
 // arguments for formatting. Returns the return value vsnprintf would return.
 //
 // The function tokenizes the format string in chunks representing either a
 // constant string or a single format directive (e.g., '%.3f') and formats each
 // chunk independently into the output string. This approach allows to figure
 // out which bytes of the output string depends on which argument and thus to
 // propagate labels more precisely.
+//
+// WARNING: This implementation does not support conversion specifiers with
+// positional arguments.
 static int format_buffer(char *str, bool has_size, size_t size,
                          const char *format, dfsan_label *va_labels,
                          dfsan_label *ret_label, va_list ap) {
-  InternalMmapVector<Chunk> chunks(8);
   size_t off = 0;
 
   while (*format) {
-    chunks.push_back(Chunk());
-    Chunk& chunk = chunks.back();
-    chunk.ptr = str + off;
-    chunk.arg = nullptr;
-
     int status = 0;
 
     if (*format != '%') {
       // Ordinary character. Consume all the characters until a '%' or the end
       // of the string.
       size_t format_size = 0;
       for (; *format && *format != '%'; ++format, ++format_size) {}
       status = format_chunk(str, off, has_size, size, format - format_size,
                             format_size);
-      chunk.label_type = Chunk::NONE;
+      dfsan_set_label(0, str + off,
+                      size_of_chunk(str, off, has_size, size, status));
     } else {
       // Conversion directive. Consume all the characters until a conversion
       // specifier or the end of the string.
       bool end_format = false;
-#define FORMAT_CHUNK(t)                                                  \
-      format_chunk(str, off, has_size, size, format - format_size,  \
+      // Whether the field width is given in the next argument (i.e., the
+      // conversion specifiers contains a '*').
+      bool has_width_arg = false;
+      int chunk_width = 0;
+#define FORMAT_CHUNK(t)                                                        \
+      has_width_arg ?                                                          \
+          format_chunk(str, off, has_size, size, format - format_size,         \
+                       format_size + 1, chunk_width, va_arg(ap, t)) :          \
+          format_chunk(str, off, has_size, size, format - format_size,         \
                        format_size + 1, va_arg(ap, t))
-
+#define LABEL_CHUNK(l)                                                         \
+      dfsan_set_label(l, str + off,                                            \
+                  size_of_chunk(str, off, has_size, size, status));
       for (size_t format_size = 1; *++format && !end_format; ++format_size) {
         switch (*format) {
           case 'd':
           case 'i':
           case 'o':
           case 'u':
           case 'x':
           case 'X':
#define LABEL_CHUNK(l)                                                         \
                 status = FORMAT_CHUNK(size_t);
                 break;
               case 't':
                 status = FORMAT_CHUNK(size_t);
                 break;
               default:
                 status = FORMAT_CHUNK(int);
             }
-            chunk.label_type = Chunk::NUMERIC;
+            LABEL_CHUNK(*va_labels++);
             end_format = true;
             break;
 
           case 'a':
           case 'A':
           case 'e':
           case 'E':
           case 'f':
           case 'F':
           case 'g':
           case 'G':
             if (*(format - 1) == 'L') {
               status = FORMAT_CHUNK(long double);
             } else {
               status = FORMAT_CHUNK(double);
             }
-            chunk.label_type = Chunk::NUMERIC;
+            LABEL_CHUNK(*va_labels++);
             end_format = true;
             break;
 
           case 'c':
             status = FORMAT_CHUNK(int);
-            chunk.label_type = Chunk::NUMERIC;
+            LABEL_CHUNK(*va_labels++);
             end_format = true;
             break;
 
-          case 's':
-            chunk.arg = va_arg(ap, char *);
+          case 's': {
+            char *chunk_arg = va_arg(ap, char *);
+            if (has_width_arg) {
               status =
                   format_chunk(str, off, has_size, size,
                                format - format_size, format_size + 1,
-                             chunk.arg);
-            chunk.label_type = Chunk::STRING;
+                               chunk_width, chunk_arg);
+            } else {
+              status =
+                  format_chunk(str, off, has_size, size,
+                               format - format_size, format_size + 1,
+                               chunk_arg);
+            }
+            va_labels++;
+            internal_memcpy(
+                shadow_for(str + off), shadow_for(chunk_arg),
+                sizeof(dfsan_label) *
+                    (size_of_chunk(str, off, has_size, size, status)));
             end_format = true;
             break;
+          }
 
           case 'p':
             status = FORMAT_CHUNK(void *);
-            chunk.label_type = Chunk::NUMERIC;
+            LABEL_CHUNK(*va_labels++);
             end_format = true;
             break;
 
-          case 'n':
-            *(va_arg(ap, int *)) = (int)off;
-            chunk.label_type = Chunk::IGNORED;
+          case 'n': {
+            int* chunk_ptr = va_arg(ap, int *);
+            *chunk_ptr = (int)off;
+            va_labels++;
+            dfsan_set_label(0, chunk_ptr, sizeof(chunk_ptr));
             end_format = true;
             break;
+          }
 
           case '%':
             status = format_chunk(str, off, has_size, size,
                                   format - format_size, format_size + 1);
-            chunk.label_type = Chunk::NONE;
+            LABEL_CHUNK(0);
             end_format = true;
             break;
 
+          case '*':
+            has_width_arg = true;
+            chunk_width = va_arg(ap, int);
+            va_labels++;
+            break;
+
           default:
             break;
         }
       }
+#undef LABEL_CHUNK
 #undef FORMAT_CHUNK
     }
 
     if (status < 0) {
       return status;
     }
 
-    // A return value of {v,}snprintf of size or more means that the output was
-    // truncated.
-    if (has_size) {
-      if (off < size) {
-        size_t ustatus = (size_t) status;
-        chunk.size = ustatus >= (size - off) ?
-            ustatus - (size - off) : ustatus;
-      } else {
-        chunk.size = 0;
-      }
-    } else {
-      chunk.size = status;
-    }
     off += status;
   }
 
-  // TODO(martignlo): Decide how to combine labels (e.g., whether to ignore or
-  // not the label of the format string).
-
-  // 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
-  // string was only partially printed ({v,}snprintf case).
-  for (size_t i = 0; i < chunks.size(); ++i) {
-    const Chunk& chunk = chunks[i];
-    void *chunk_ptr = const_cast<char *>(chunk.ptr);
-
-    switch (chunk.label_type) {
-      case Chunk::NONE:
-        dfsan_set_label(0, chunk_ptr, chunk.size);
-        break;
-      case Chunk::IGNORED:
-        va_labels++;
-        dfsan_set_label(0, chunk_ptr, chunk.size);
-        break;
-      case Chunk::NUMERIC: {
-        dfsan_label label = *va_labels++;
-        dfsan_set_label(label, chunk_ptr, chunk.size);
-        break;
-      }
-      case Chunk::STRING: {
-        // Consume the label of the pointer to the string
-        va_labels++;
-        internal_memcpy(shadow_for(chunk_ptr),
-                        shadow_for(chunk.arg),
-                        sizeof(dfsan_label) * (strlen(chunk.arg)));
-        break;
-      }
-    }
-  }
-
   *ret_label = 0;
 
   // Number of bytes written in total.
   return off;
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
 int __dfsw_sprintf(char *str, const char *format, dfsan_label str_label,

test/dfsan/custom.cc

void test_sprintf() {
   test_sprintf_chunk("1.234560e-01", "%e", 0.123456);
   test_sprintf_chunk("1.234560E-01", "%E", 0.123456);
   test_sprintf_chunk("0.1234567891234560", "%.16Lf",
                      (long double) 0.123456789123456);
 
   test_sprintf_chunk("z", "%c", 'z');
 
   // %n, %s, %d, %f, and %% already tested
+
+  // Test formatting with width passed as an argument.
+  r = sprintf(buf, "hi %*d my %*s friend %.*f", 3, 1, 6, "dear", 4, 3.14159265359);
+  assert(r == 30);
+  assert(strcmp(buf, "hi   1 my   dear friend 3.1416") == 0);
 }
 
 void test_snprintf() {
   char buf[2048];
   memset(buf, 'a', sizeof(buf));
   dfsan_set_label(0, buf, sizeof(buf));
   const char* s = "world";
   int y = 2014;