Codegen: Don't emit lifetime.end for variables used as map keys

Previously a lifetime.end was always inserted for expressions used as map keys. This is not correct if the expression is a variable which is used again later in the program. $blah should not be marked as dead after line 4, it should persist at least until after line 5: 1 BEGIN 2 { 3 $blah = "abc"; 4 @x[$blah] = 1; 5 @y[$blah] = 1; 6 }
bpftrace · May 26, 2024 · d44a6bf · d44a6bf
1 parent 51f6882
commit d44a6bf
Show file tree

Hide file tree

Showing 9 changed files with 17 additions and 23 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -61,6 +61,8 @@ and this project adheres to
   - [#3024](https://github.com/bpftrace/bpftrace/pull/3024)
 - Fix error in dereferencing kernel double pointers
   - [#3024](https://github.com/bpftrace/bpftrace/pull/3024)
+- Fix variable corruption when used as map key
+  - [#3195](https://github.com/bpftrace/bpftrace/pull/3195)
 #### Docs
 #### Tools
 

diff --git a/src/ast/passes/codegen_llvm.cpp b/src/ast/passes/codegen_llvm.cpp
@@ -2777,6 +2777,7 @@ std::tuple<Value *, CodegenLLVM::ScopedExprDeleter> CodegenLLVM::getMapKey(
     Map &map)
 {
   Value *key;
+  bool alloca_created_here = true;
   if (map.vargs) {
     // A single value as a map key (e.g., @[comm] = 0;)
     if (map.vargs->size() == 1) {
@@ -2793,6 +2794,10 @@ std::tuple<Value *, CodegenLLVM::ScopedExprDeleter> CodegenLLVM::getMapKey(
           key = expr_;
           // Call-ee freed
           scoped_del.disarm();
+
+          // We don't have enough visibility into where key comes from to safely
+          // end its lifetime. It could be a variable, for example.
+          alloca_created_here = false;
         }
       } else {
         key = b_.CreateAllocaBPF(expr->type, map.ident + "_key");
@@ -2815,8 +2820,8 @@ std::tuple<Value *, CodegenLLVM::ScopedExprDeleter> CodegenLLVM::getMapKey(
     b_.CreateStore(b_.getInt64(0), key);
   }
 
-  auto key_deleter = [this, key]() {
-    if (dyn_cast<AllocaInst>(key))
+  auto key_deleter = [this, key, alloca_created_here]() {
+    if (alloca_created_here)
       b_.CreateLifetimeEnd(key);
   };
   return std::make_tuple(key, ScopedExprDeleter(std::move(key_deleter)));
@@ -2832,6 +2837,9 @@ AllocaInst *CodegenLLVM::getMultiMapKey(Map &map,
   for (auto *extra_key : extra_keys) {
     size += module_->getDataLayout().getTypeAllocSize(extra_key->getType());
   }
+
+  // If key ever changes to not be allocated here, be sure to update getMapKey()
+  // as well to take the new lifetime semantics into account.
   AllocaInst *key = b_.CreateAllocaBPF(size, map.ident + "_key");
   auto *key_type = ArrayType::get(b_.getInt8Ty(), size);
 

diff --git a/tests/codegen/llvm/call_ntop_key.ll b/tests/codegen/llvm/call_ntop_key.ll
@@ -55,8 +55,6 @@ lookup_failure:                                   ; preds = %entry
 lookup_merge:                                     ; preds = %lookup_failure, %lookup_success
   %11 = bitcast i64* %lookup_elem_val to i8*
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* %11)
-  %12 = bitcast %inet_t* %inet to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %12)
   ret i64 0
 }
 

diff --git a/tests/codegen/llvm/call_usym_key.ll b/tests/codegen/llvm/call_usym_key.ll
@@ -56,8 +56,6 @@ lookup_failure:                                   ; preds = %entry
 lookup_merge:                                     ; preds = %lookup_failure, %lookup_success
   %11 = bitcast i64* %lookup_elem_val to i8*
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* %11)
-  %12 = bitcast %usym_t* %usym to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %12)
   ret i64 0
 }
 

diff --git a/tests/codegen/llvm/for_map_strings.ll b/tests/codegen/llvm/for_map_strings.ll
@@ -29,10 +29,8 @@ entry:
   call void @llvm.lifetime.start.p0i8(i64 -1, i8* %2)
   store [4 x i8] c"abc\00", [4 x i8]* %str1, align 1
   %update_elem = call i64 inttoptr (i64 2 to i64 (%"struct map_t"*, [4 x i8]*, [4 x i8]*, i64)*)(%"struct map_t"* @AT_map, [4 x i8]* %str1, [4 x i8]* %str, i64 0)
-  %3 = bitcast [4 x i8]* %str1 to i8*
+  %3 = bitcast [4 x i8]* %str to i8*
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* %3)
-  %4 = bitcast [4 x i8]* %str to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %4)
   %for_each_map_elem = call i64 inttoptr (i64 164 to i64 (%"struct map_t"*, i64 (i8*, i8*, i8*, i8*)*, i8*, i64)*)(%"struct map_t"* @AT_map, i64 (i8*, i8*, i8*, i8*)* @map_for_each_cb, i8* null, i64 0)
   ret i64 0
 }

diff --git a/tests/codegen/llvm/literal_strncmp.ll b/tests/codegen/llvm/literal_strncmp.ll
@@ -100,8 +100,6 @@ lookup_failure:                                   ; preds = %pred_true
 lookup_merge:                                     ; preds = %lookup_failure, %lookup_success
   %19 = bitcast i64* %lookup_elem_val to i8*
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* %19)
-  %20 = bitcast [16 x i8]* %comm5 to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %20)
   ret i64 0
 }
 

diff --git a/tests/codegen/llvm/string_equal_comparison.ll b/tests/codegen/llvm/string_equal_comparison.ll
@@ -130,8 +130,6 @@ lookup_failure:                                   ; preds = %pred_true
 lookup_merge:                                     ; preds = %lookup_failure, %lookup_success
   %25 = bitcast i64* %lookup_elem_val to i8*
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* %25)
-  %26 = bitcast [16 x i8]* %comm17 to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %26)
   ret i64 0
 }
 

diff --git a/tests/codegen/llvm/string_not_equal_comparison.ll b/tests/codegen/llvm/string_not_equal_comparison.ll
@@ -130,8 +130,6 @@ lookup_failure:                                   ; preds = %pred_true
 lookup_merge:                                     ; preds = %lookup_failure, %lookup_success
   %25 = bitcast i64* %lookup_elem_val to i8*
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* %25)
-  %26 = bitcast [16 x i8]* %comm17 to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %26)
   ret i64 0
 }
 

diff --git a/tests/codegen/llvm/variable_map_key_lifetime.ll b/tests/codegen/llvm/variable_map_key_lifetime.ll
@@ -39,16 +39,12 @@ entry:
   %update_elem = call i64 inttoptr (i64 2 to i64 (%"struct map_t"*, [4 x i8]*, i64*, i64)*)(%"struct map_t"* @AT_x, [4 x i8]* %"$myvar", i64* %"@x_val", i64 0)
   %8 = bitcast i64* %"@x_val" to i8*
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* %8)
-  %9 = bitcast [4 x i8]* %"$myvar" to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %9)
-  %10 = bitcast i64* %"@x_val1" to i8*
-  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %10)
+  %9 = bitcast i64* %"@x_val1" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* %9)
   store i64 1, i64* %"@x_val1", align 8
   %update_elem2 = call i64 inttoptr (i64 2 to i64 (%"struct map_t"*, [4 x i8]*, i64*, i64)*)(%"struct map_t"* @AT_x, [4 x i8]* %"$myvar", i64* %"@x_val1", i64 0)
-  %11 = bitcast i64* %"@x_val1" to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %11)
-  %12 = bitcast [4 x i8]* %"$myvar" to i8*
-  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %12)
+  %10 = bitcast i64* %"@x_val1" to i8*
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* %10)
   ret i64 0
 }