Make printk safer and use less stack

3 files changed, 35 insertions, 33 deletions

diff --git a/pk/console.c b/pk/console.c
index a15cb8a..69bb9e9 100644
--- a/pk/console.c
+++ b/pk/console.c
@@ -6,10 +6,11 @@
 #include <stdbool.h>
 #include <string.h>
 
-static void vsprintk(char* out, const char* s, va_list vl)
+static int vsnprintf(char* out, size_t n, const char* s, va_list vl)
 {
   bool format = false;
   bool longarg = false;
+  size_t pos = 0;
   for( ; *s; s++)
   {
     if(format)
@@ -21,15 +22,14 @@ static void vsprintk(char* out, const char* s, va_list vl)
           break;
         case 'p':
           longarg = true;
-          *out++ = '0';
-          *out++ = 'x';
+          if (++pos < n) out[pos-1] = '0';
+          if (++pos < n) out[pos-1] = 'x';
         case 'x':
         {
-          long n = longarg ? va_arg(vl,long) : va_arg(vl,int);
-          for(int i = 2*(longarg ? sizeof(long) : sizeof(int))-1; i >= 0; i--)
-          {
-            int d = (n >> (4*i)) & 0xF;
-            *out++ = (d < 10 ? '0'+d : 'a'+d-10);
+          long num = longarg ? va_arg(vl, long) : va_arg(vl, int);
+          for(int i = 2*(longarg ? sizeof(long) : sizeof(int))-1; i >= 0; i--) {
+            int d = (num >> (4*i)) & 0xF;
+            if (++pos < n) out[pos-1] = (d < 10 ? '0'+d : 'a'+d-10);
           }
           longarg = false;
           format = false;
@@ -37,20 +37,19 @@ static void vsprintk(char* out, const char* s, va_list vl)
         }
         case 'd':
         {
-          long n = longarg ? va_arg(vl,long) : va_arg(vl,int);
-          if(n < 0)
-          {
-            n = -n;
-            *out++ = '-';
+          long num = longarg ? va_arg(vl, long) : va_arg(vl, int);
+          if (num < 0) {
+            num = -num;
+            if (++pos < n) out[pos-1] = '-';
           }
           long digits = 1;
-          for(long nn = n ; nn /= 10; digits++);
-          for(int i = digits-1; i >= 0; i--)
-          {
-            out[i] = '0' + n%10;
-            n /= 10;
+          for (long nn = num; nn /= 10; digits++)
+            ;
+          for (int i = digits-1; i >= 0; i--) {
+            if (pos + i + 1 < n) out[pos + i] = '0' + (num % 10);
+            num /= 10;
           }
-          out += digits;
+          pos += digits;
           longarg = false;
           format = false;
           break;
@@ -59,14 +58,14 @@ static void vsprintk(char* out, const char* s, va_list vl)
         {
           const char* s2 = va_arg(vl,const char*);
           while(*s2)
-            *out++ = *s2++;
+            if (++pos < n) out[pos-1] = *s2++;
           longarg = false;
           format = false;
           break;
         }
         case 'c':
         {
-          *out++ = (char)va_arg(vl,int);
+          if (++pos < n) out[pos-1] = (char)va_arg(vl,int);
           longarg = false;
           format = false;
           break;
@@ -78,16 +77,20 @@ static void vsprintk(char* out, const char* s, va_list vl)
     else if(*s == '%')
       format = true;
     else
-      *out++ = *s;
+      if (++pos < n) out[pos-1] = *s;
   }
-  *out++ = '\0';
+  if (pos < n)
+    out[pos] = 0;
+  else if (n)
+    out[n-1] = 0;
+  return pos;
 }
 
 static void vprintk(const char* s, va_list vl)
 {
-  char out[1024]; // XXX
-  vsprintk(out, s, vl);
-  file_write(stderr, out, strlen(out));
+  char out[256]; // XXX
+  int res = vsnprintf(out, sizeof(out), s, vl);
+  file_write(stderr, out, res < sizeof(out) ? res : sizeof(out));
 }
 
 void printk(const char* s, ...)
@@ -100,14 +103,13 @@ void printk(const char* s, ...)
   va_end(vl);
 }
 
-void sprintk(char* out, const char* s, ...)
+int snprintf(char* out, size_t n, const char* s, ...)
 {
   va_list vl;
-  va_start(vl,s);
-
-  vsprintk(out,s,vl);
-
+  va_start(vl, s);
+  int res = vsnprintf(out, n, s, vl);
   va_end(vl);
+  return res;
 }
 
 void dump_tf(trapframe_t* tf)
diff --git a/pk/mtrap.c b/pk/mtrap.c
index 27936fe..1d461cf 100644
--- a/pk/mtrap.c
+++ b/pk/mtrap.c
@@ -126,7 +126,7 @@ static uintptr_t mcall_console_putchar(uint8_t ch)
 }
 
 #define printm(str, ...) ({ \
-  char buf[1024], *p = buf; sprintk(buf, str, __VA_ARGS__); \
+  char buf[128], *p = buf; snprintf(buf, sizeof(buf), str, __VA_ARGS__); \
   while (*p) mcall_console_putchar(*p++); })
 
 static uintptr_t mcall_dev_req(sbi_device_message *m)
diff --git a/pk/pk.h b/pk/pk.h
index 06d68ee..d21a890 100644
--- a/pk/pk.h
+++ b/pk/pk.h
@@ -48,7 +48,7 @@ extern uint32_t num_harts;
 
 struct mainvars* parse_args(struct mainvars*);
 void printk(const char* s, ...);
-void sprintk(char* out, const char* s, ...);
+int snprintf(char* out, size_t n, const char* s, ...);
 void init_tf(trapframe_t*, long pc, long sp, int user64);
 void pop_tf(trapframe_t*) __attribute__((noreturn));
 void dump_tf(trapframe_t*);