Merge pull request #37 from iovisor/rand

add rand builtin

README.md

@@ -219,6 +219,7 @@ Variables:
 - `retval` - Return value from function being traced
 - `func` - Name of the function currently being traced
 - `curtask` - Current task_struct as a u64.
+- `rand` - Random number of type u32.
 
 Functions:
 - `hist(int n)` - Produce a log2 histogram of values of `n`

src/ast/codegen_llvm.cpp

@@ -74,6 +74,10 @@ void CodegenLLVM::visit(Builtin &builtin)
   {
     expr_ = b_.CreateGetCurrentTask();
   }
+  else if (builtin.ident == "rand")
+  {
+    expr_ = b_.CreateGetRandom();
+  }
   else if (builtin.ident == "comm")
   {
     AllocaInst *buf = b_.CreateAllocaBPF(builtin.type, "comm");

src/ast/irbuilderbpf.cpp

@@ -320,6 +320,19 @@ CallInst *IRBuilderBPF::CreateGetCurrentTask()
   return CreateCall(getcurtask_func, {}, "get_cur_task");
 }
 
+CallInst *IRBuilderBPF::CreateGetRandom()
+{
+  // u64 bpf_get_prandom_u32(void)
+  // Return: random
+  FunctionType *getrandom_func_type = FunctionType::get(getInt64Ty(), false);
+  PointerType *getrandom_func_ptr_type = PointerType::get(getrandom_func_type, 0);
+  Constant *getrandom_func = ConstantExpr::getCast(
+      Instruction::IntToPtr,
+      getInt64(BPF_FUNC_get_prandom_u32),
+      getrandom_func_ptr_type);
+  return CreateCall(getrandom_func, {}, "get_random");
+}
+
 CallInst *IRBuilderBPF::CreateGetStackId(Value *ctx, bool ustack)
 {
   Value *map_ptr = CreateBpfPseudoCall(bpftrace_.stackid_map_->mapfd_);

src/ast/irbuilderbpf.h

@@ -37,6 +37,7 @@ public:
   CallInst   *CreateGetUidGid();
   CallInst   *CreateGetCpuId();
   CallInst   *CreateGetCurrentTask();
+  CallInst   *CreateGetRandom();
   CallInst   *CreateGetStackId(Value *ctx, bool ustack);
   CallInst   *CreateGetJoinMap(Value *ctx);
   void        CreateGetCurrentComm(AllocaInst *buf, size_t size);

src/ast/semantic_analyser.cpp

@@ -33,6 +33,7 @@ void SemanticAnalyser::visit(Builtin &builtin)
       builtin.ident == "gid" ||
       builtin.ident == "cpu" ||
       builtin.ident == "curtask" ||
+      builtin.ident == "rand" ||
       builtin.ident == "retval") {
     builtin.type = SizedType(Type::integer, 8);
   }

src/lexer.l

@@ -45,7 +45,7 @@ header <(\\.|[_\-\./a-zA-Z0-9])*>
 <COMMENT>"EOF"          driver.error(loc, std::string("end of file during comment"));
 <COMMENT>.|"\n"         ;
 
-pid|tid|uid|gid|nsecs|cpu|comm|stack|ustack|arg[0-9]|retval|func|name|curtask {
+pid|tid|uid|gid|nsecs|cpu|comm|stack|ustack|arg[0-9]|retval|func|name|curtask|rand {
                           return Parser::make_BUILTIN(yytext, loc); }
 {ident}                 { return Parser::make_IDENT(yytext, loc); }
 {path}                  { return Parser::make_PATH(yytext, loc); }

tests/codegen.cpp

@@ -525,6 +525,42 @@ attributes #1 = { argmemonly nounwind }
 )EXPECTED");
 }
 
+TEST(codegen, builtin_rand)
+{
+  test("kprobe:f { @x = rand }",
+
+R"EXPECTED(; Function Attrs: nounwind
+declare i64 @llvm.bpf.pseudo(i64, i64) #0
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+define i64 @"kprobe:f"(i8* nocapture readnone) local_unnamed_addr section "s_kprobe:f" {
+entry:
+  %"@x_val" = alloca i64, align 8
+  %"@x_key" = alloca i64, align 8
+  %get_random = tail call i64 inttoptr (i64 7 to i64 ()*)()
+  %1 = bitcast i64* %"@x_key" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
+  store i64 0, i64* %"@x_key", align 8
+  %2 = bitcast i64* %"@x_val" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %2)
+  store i64 %get_random, i64* %"@x_val", align 8
+  %pseudo = tail call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %update_elem = call i64 inttoptr (i64 2 to i64 (i8*, i8*, i8*, i64)*)(i64 %pseudo, i64* nonnull %"@x_key", i64* nonnull %"@x_val", i64 0)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %2)
+  ret i64 0
+}
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nounwind }
+)EXPECTED");
+}
+
 TEST(codegen, builtin_comm)
 {
   test("kprobe:f { @x = comm }",

tests/parser.cpp

@@ -30,6 +30,7 @@ TEST(Parser, builtin_variables)
   test("kprobe:f { nsecs }", "Program\n kprobe:f\n  builtin: nsecs\n");
   test("kprobe:f { cpu }", "Program\n kprobe:f\n  builtin: cpu\n");
   test("kprobe:f { curtask }", "Program\n kprobe:f\n  builtin: curtask\n");
+  test("kprobe:f { rand }", "Program\n kprobe:f\n  builtin: rand\n");
   test("kprobe:f { comm }", "Program\n kprobe:f\n  builtin: comm\n");
   test("kprobe:f { stack }", "Program\n kprobe:f\n  builtin: stack\n");
   test("kprobe:f { ustack }", "Program\n kprobe:f\n  builtin: ustack\n");

tests/semantic_analyser.cpp

@@ -68,6 +68,7 @@ TEST(semantic_analyser, builtin_variables)
   test("kprobe:f { nsecs }", 0);
   test("kprobe:f { cpu }", 0);
   test("kprobe:f { curtask }", 0);
+  test("kprobe:f { rand }", 0);
   test("kprobe:f { comm }", 0);
   test("kprobe:f { stack }", 0);
   test("kprobe:f { ustack }", 0);