Fix crash using some per-cpu maps as map keys

'count', 'sum', 'avg', 'min', 'max' are all castable
per-cpu map aggregation types so they should be valid
as map keys when printing.

Also add semantic analyser checks for invalid per-cpu
map aggregation types.

Issue: bpftrace#3591

Author: Jordan Rome

.github/include/aot_skip.txt

@@ -91,6 +91,7 @@ aot.other.no symbolize enum after arithmetic
 aot.other.no symbolize enum after arithmetic mixed
 aot.other.per_cpu_map_arithmetic
 aot.other.per_cpu_map_avg_if
+aot.other.per_cpu_map_as_map_key
 # Enum metadata is not being serialized
 aot.other.symbolize enum in map key
 # Enum metadata is not being serialized

CHANGELOG.md

@@ -112,6 +112,8 @@ and this project adheres to
   - [#3542](https://github.com/bpftrace/bpftrace/pull/3542)
 - Fix field access and offsetof for strings that are builtin types
   - [#3565](https://github.com/bpftrace/bpftrace/pull/3565)
+- Fix crash when using castable per-cpu map types as map keys
+  - [#3594](https://github.com/bpftrace/bpftrace/pull/3594)
 #### Security
 #### Docs
 - Remove mention of unsupported character literals

src/ast/passes/semantic_analyser.cpp

@@ -1672,6 +1672,10 @@ void SemanticAnalyser::validate_map_key(const SizedType &key,
     LOG(ERROR, loc, err_) << "context cannot be used as a map key";
   }
 
+  if (key.IsHistTy() || key.IsLhistTy() || key.IsStatsTy()) {
+    LOG(ERROR, loc, err_) << key << " cannot be used as a map key";
+  }
+
   if (is_final_pass() && key.IsNoneTy()) {
     LOG(ERROR, loc, err_) << "Invalid expression for assignment: ";
   }

src/output.cpp

@@ -452,6 +452,11 @@ std::string Output::map_key_str(BPFtrace &bpftrace,
     case Type::array:
     case Type::mac_address:
     case Type::record:
+    case Type::count_t:
+    case Type::avg_t:
+    case Type::max_t:
+    case Type::min_t:
+    case Type::sum_t:
       return value_to_str(bpftrace, arg, data, false, 1, true);
     case Type::tuple: {
       std::vector<std::string> elems;
@@ -466,17 +471,12 @@ std::string Output::map_key_str(BPFtrace &bpftrace,
     }
     case Type::cgroup_path_t:
     case Type::strerror_t:
-    case Type::avg_t:
-    case Type::count_t:
     case Type::hist_t:
     case Type::lhist_t:
-    case Type::max_t:
-    case Type::min_t:
     case Type::none:
     case Type::reference:
     case Type::stack_mode:
     case Type::stats_t:
-    case Type::sum_t:
     case Type::timestamp_mode:
     case Type::voidtype:
       LOG(BUG) << "Invalid mapkey argument type: " << arg;

tests/runtime/other

@@ -275,6 +275,10 @@ NAME per_cpu_map_cast
 PROG BEGIN { @a = count(); @b = sum(10); printf("%d-%d\n", (uint64)@a, (int64)@b); exit();}
 EXPECT 1-10
 
+NAME per_cpu_map_as_map_key
+PROG BEGIN {@a = count(); @b = sum(5); @c = min(1); @d = max(10); @e = avg(7); @[@a, @b, @c, @d, @e] = 1; exit(); }
+EXPECT @[1, 5, 1, 10, 7]: 1
+
 NAME dry run empty output
 RUN {{BPFTRACE}} --dry-run -e 'BEGIN { printf("hello\n"); @ = 0; }'
 EXPECT_NONE hello

tests/semantic_analyser.cpp

@@ -692,6 +692,9 @@ kprobe:f { $x = hist(1); }
 )");
   test_error("kprobe:f { @x[hist(1)] = 1; }", R"(
 stdin:1:12-22: ERROR: hist() should be directly assigned to a map
+kprobe:f { @x[hist(1)] = 1; }
+           ~~~~~~~~~~
+stdin:1:12-22: ERROR: hist_t cannot be used as a map key
 kprobe:f { @x[hist(1)] = 1; }
            ~~~~~~~~~~
 )");
@@ -764,6 +767,9 @@ stdin:1:12-21: ERROR: lhist() should be directly assigned to a map
 kprobe:f { @[lhist()] = 1; }
            ~~~~~~~~~
 stdin:1:12-21: ERROR: lhist() requires 4 arguments (0 provided)
+kprobe:f { @[lhist()] = 1; }
+           ~~~~~~~~~
+stdin:1:12-21: ERROR: lhist_t cannot be used as a map key
 kprobe:f { @[lhist()] = 1; }
            ~~~~~~~~~
 )");
@@ -2556,6 +2562,33 @@ stdin:4:9-30: ERROR: Argument mismatch for @x: trying to access with arguments:
 )");
 }
 
+TEST(semantic_analyser, per_cpu_map_as_map_key)
+{
+  test("BEGIN { @x = count(); @y[@x] = 1; }");
+  test("BEGIN { @x = sum(10); @y[@x] = 1; }");
+  test("BEGIN { @x = min(1); @y[@x] = 1; }");
+  test("BEGIN { @x = max(1); @y[@x] = 1; }");
+  test("BEGIN { @x = avg(1); @y[@x] = 1; }");
+
+  test_error("BEGIN { @x = hist(10); @y[@x] = 1; }", R"(
+stdin:1:24-29: ERROR: hist_t cannot be used as a map key
+BEGIN { @x = hist(10); @y[@x] = 1; }
+                       ~~~~~
+)");
+
+  test_error("BEGIN { @x = lhist(10, 0, 10, 1); @y[@x] = 1; }", R"(
+stdin:1:35-40: ERROR: lhist_t cannot be used as a map key
+BEGIN { @x = lhist(10, 0, 10, 1); @y[@x] = 1; }
+                                  ~~~~~
+)");
+
+  test_error("BEGIN { @x = stats(10); @y[@x] = 1; }", R"(
+stdin:1:25-30: ERROR: stats_t cannot be used as a map key
+BEGIN { @x = stats(10); @y[@x] = 1; }
+                        ~~~~~
+)");
+}
+
 TEST(semantic_analyser, probe_short_name)
 {
   test("t:a:b { args }");