From 977091ec495120b79476a0d319bb109542205dee Mon Sep 17 00:00:00 2001
From: Brenden Blanco <bblanco@plumgrid.com>
Date: Thu, 5 May 2016 11:37:59 -0700
Subject: [PATCH] Add check for number of arguments

There are two problems here:
 1. bcc is leaving the extra function parameters in the prototype, when
 these really just serve as an annotation on the method being traced.
 They aren't really passed into the function. LLVM lowering phase later
 errors out on such functions.
 2. bcc doesn't limit the size of the argument list, when currently it
 just supports up to the number of arguments held in registers (6).

Fix 1. by rewriting the arguments out of the prototype and into the
preamble where they are currently being initialized.
Fix 2. by enforcing the limit and returning a more meaningful error
message.

Added a test case.

Extra fix included for string type on python3.

Fixes: #497
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>
---
 src/cc/frontends/clang/b_frontend_action.cc | 20 +++++++++++++++++++-
 tests/python/test_clang.py                  | 10 ++++++++++
 tests/python/test_stackid.py                |  2 +-
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
index da73c4b..709daff 100644
--- a/src/cc/frontends/clang/b_frontend_action.cc
+++ b/src/cc/frontends/clang/b_frontend_action.cc
@@ -32,6 +32,7 @@
 
 namespace ebpf {
 
+constexpr int MAX_CALLING_CONV_REGS = 6;
 const char *calling_conv_regs_x86[] = {
   "di", "si", "dx", "cx", "r8", "r9"
 };
@@ -255,6 +256,11 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
   if (D->isExternallyVisible() && D->hasBody()) {
     string attr = string("__attribute__((section(\"") + BPF_FN_PREFIX + D->getName().str() + "\")))\n";
     rewriter_.InsertText(D->getLocStart(), attr);
+    if (D->param_size() > MAX_CALLING_CONV_REGS + 1) {
+      error(D->getParamDecl(MAX_CALLING_CONV_REGS + 1)->getLocStart(),
+            "too many arguments, bcc only supports in-register parameters");
+      return false;
+    }
     // remember the arg names of the current function...first one is the ctx
     fn_args_.clear();
     string preamble = "{";
@@ -265,12 +271,24 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
       }
       fn_args_.push_back(arg);
       if (fn_args_.size() > 1) {
+        // Move the args into a preamble section where the same params are
+        // declared and initialized from pt_regs.
+        // Todo: this init should be done only when the program requests it.
+        string text = rewriter_.getRewrittenText(
+            SourceRange(arg->getLocStart(), arg->getLocEnd()));
         arg->addAttr(UnavailableAttr::CreateImplicit(C, "ptregs"));
         size_t d = fn_args_.size() - 2;
         const char *reg = calling_conv_regs[d];
-        preamble += arg->getName().str() + " = " + fn_args_[0]->getName().str() + "->" + string(reg) + ";";
+        preamble += " " + text + " = " + fn_args_[0]->getName().str() + "->" +
+                    string(reg) + ";";
       }
     }
+    if (D->param_size() > 1) {
+      rewriter_.ReplaceText(
+          SourceRange(D->getParamDecl(0)->getLocEnd(),
+                      D->getParamDecl(D->getNumParams() - 1)->getLocEnd()),
+          fn_args_[0]->getName());
+    }
     // for each trace argument, convert the variable from ptregs to something on stack
     if (CompoundStmt *S = dyn_cast<CompoundStmt>(D->getBody()))
       rewriter_.ReplaceText(S->getLBracLoc(), 1, preamble);
diff --git a/tests/python/test_clang.py b/tests/python/test_clang.py
index 214ffaa..f4e897f 100755
--- a/tests/python/test_clang.py
+++ b/tests/python/test_clang.py
@@ -328,5 +328,15 @@ BPF_TABLE("hash", struct bpf_tunnel_key, int, t1, 1);
         t1 = b["t1"]
         print(t1.Key().remote_ipv4)
 
+    def test_too_many_args(self):
+        text = """
+#include <uapi/linux/ptrace.h>
+int many(struct pt_regs *ctx, int a, int b, int c, int d, int e, int f, int g) {
+    return 0;
+}
+"""
+        with self.assertRaises(Exception):
+            b = BPF(text=text)
+
 if __name__ == "__main__":
     main()
diff --git a/tests/python/test_stackid.py b/tests/python/test_stackid.py
index 4d258d7..481adbd 100755
--- a/tests/python/test_stackid.py
+++ b/tests/python/test_stackid.py
@@ -47,7 +47,7 @@ int kprobe__htab_map_lookup_elem(struct pt_regs *ctx, struct bpf_map *map, u64 *
         stackid = stack_entries[k]
         self.assertIsNotNone(stackid)
         stack = stack_traces[stackid].ip
-        self.assertEqual(b.ksym(stack[0]), b"htab_map_lookup_elem")
+        self.assertEqual(b.ksym(stack[0]), "htab_map_lookup_elem")
 
 
 if __name__ == "__main__":
-- 
2.7.4