From 8f5e7824d6ed7f038a7aa1066d96ac39dfa5024a Mon Sep 17 00:00:00 2001
From: Jiri Olsa <jolsa@kernel.org>
Date: Mon, 3 Jul 2023 20:31:09 +0000
Subject: [PATCH] tetraong: Add InRefMap operator

Adding InRefMap operator that allows to refference list of values,
which are normally defined directly for InMap operator.

With InRefMap operator we can spcify special 'ref:<NAME>' value
that refference another map of values, like:

   - matchArgs:
     - index: 0
       operator: "InRefMap"
       values:
       - "ref:killer"

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 pkg/selectors/kernel.go          | 51 ++++++++++++++++++++++++++------
 pkg/selectors/selectors.go       |  7 +++--
 pkg/sensors/tracing/selectors.go | 19 +++++++++++-
 3 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/pkg/selectors/kernel.go b/pkg/selectors/kernel.go
index 1e2d0ed88c6..093bb842f36 100644
--- a/pkg/selectors/kernel.go
+++ b/pkg/selectors/kernel.go
@@ -212,17 +212,19 @@ const (
 	SelectorOpPrefix  = 8
 	SelectorOpPostfix = 9
 	// Map ops
-	SelectorInMap    = 10
-	SelectorNotInMap = 11
+	SelectorInMap       = 10
+	SelectorNotInMap    = 11
+	SelectorInRefMap    = 12
+	SelectorNotInRefMap = 13
 
-	SelectorOpMASK = 12
+	SelectorOpMASK = 14
 
 	// socket ops
-	SelectorOpSaddr    = 13
-	SelectorOpDaddr    = 14
-	SelectorOpSport    = 15
-	SelectorOpDport    = 16
-	SelectorOpProtocol = 17
+	SelectorOpSaddr    = 15
+	SelectorOpDaddr    = 16
+	SelectorOpSport    = 17
+	SelectorOpDport    = 18
+	SelectorOpProtocol = 19
 )
 
 func SelectorOp(op string) (uint32, error) {
@@ -247,6 +249,10 @@ func SelectorOp(op string) (uint32, error) {
 		return SelectorInMap, nil
 	case "NotInMap":
 		return SelectorNotInMap, nil
+	case "InRefMap":
+		return SelectorInRefMap, nil
+	case "NotInRefMap":
+		return SelectorNotInRefMap, nil
 	case "mask", "Mask":
 		return SelectorOpMASK, nil
 	case "saddr", "Saddr", "SAddr":
@@ -345,6 +351,22 @@ func argSelectorType(arg *v1alpha1.ArgSelector, sig []v1alpha1.KProbeArg) (uint3
 	return 0, fmt.Errorf("argFilter for unknown index")
 }
 
+func setupRefMap(k *KernelSelectorState, values []string, ty uint32) error {
+	if len(values) != 1 {
+		return fmt.Errorf("RefMap needs single value")
+	}
+	if values[0] != "ref:killer" {
+		return fmt.Errorf("RefMap supports only killer map")
+	}
+
+	mid, m := k.newValueMap()
+	// mark as the killer map, it will be filled in later
+	m.IsKiller = true
+	// write the map id into the selector
+	WriteSelectorUint32(k, mid)
+	return nil
+}
+
 func writeMatchValuesInMap(k *KernelSelectorState, values []string, ty uint32) error {
 	mid, m := k.newValueMap()
 	for _, v := range values {
@@ -479,7 +501,13 @@ func ParseMatchArg(k *KernelSelectorState, arg *v1alpha1.ArgSelector, sig []v1al
 	if err != nil {
 		return fmt.Errorf("matcharg error: %w", err)
 	}
-	WriteSelectorUint32(k, op)
+
+	if op == SelectorInRefMap {
+		WriteSelectorUint32(k, SelectorInMap)
+	} else {
+		WriteSelectorUint32(k, op)
+	}
+
 	moff := AdvanceSelectorLength(k)
 	ty, err := argSelectorType(arg, sig)
 	if err != nil {
@@ -487,6 +515,11 @@ func ParseMatchArg(k *KernelSelectorState, arg *v1alpha1.ArgSelector, sig []v1al
 	}
 	WriteSelectorUint32(k, ty)
 	switch op {
+	case SelectorInRefMap:
+		err := setupRefMap(k, arg.Values, ty)
+		if err != nil {
+			return fmt.Errorf("setupRefMap error: %w", err)
+		}
 	case SelectorInMap, SelectorNotInMap:
 		err := writeMatchValuesInMap(k, arg.Values, ty)
 		if err != nil {
diff --git a/pkg/selectors/selectors.go b/pkg/selectors/selectors.go
index 588bcc47cfd..98dde6ed4e7 100644
--- a/pkg/selectors/selectors.go
+++ b/pkg/selectors/selectors.go
@@ -27,7 +27,8 @@ func (k *MatchBinariesMappings) GetBinSelNamesMap() map[uint32]uint32 {
 }
 
 type ValueMap struct {
-	Data map[[8]byte]struct{}
+	Data     map[[8]byte]struct{}
+	IsKiller bool
 }
 
 type KernelSelectorState struct {
@@ -145,10 +146,10 @@ func ArgSelectorValue(v string) ([]byte, uint32) {
 	return b, uint32(len(b))
 }
 
-func (k *KernelSelectorState) newValueMap() (uint32, ValueMap) {
+func (k *KernelSelectorState) newValueMap() (uint32, *ValueMap) {
 	mapid := len(k.valueMaps)
 	vm := ValueMap{}
 	vm.Data = make(map[[8]byte]struct{})
 	k.valueMaps = append(k.valueMaps, vm)
-	return uint32(mapid), k.valueMaps[mapid]
+	return uint32(mapid), &k.valueMaps[mapid]
 }
diff --git a/pkg/sensors/tracing/selectors.go b/pkg/sensors/tracing/selectors.go
index e06b4b487de..71795408c70 100644
--- a/pkg/sensors/tracing/selectors.go
+++ b/pkg/sensors/tracing/selectors.go
@@ -4,6 +4,7 @@
 package tracing
 
 import (
+	"encoding/binary"
 	"fmt"
 	"path/filepath"
 
@@ -102,7 +103,23 @@ func populateArgFilterMaps(
 	outerMap *ebpf.Map,
 ) error {
 	for i, vm := range k.ValueMaps() {
-		err := populateArgFilterMap(pinPathPrefix, outerMap, uint32(i), vm.Data)
+		data := vm.Data
+
+		if vm.IsKiller {
+			values, err := KillerMapValues()
+			if err != nil {
+				return err
+			}
+
+			for _, v := range values {
+				var val [8]byte
+
+				binary.LittleEndian.PutUint64(val[:], uint64(v))
+				data[val] = struct{}{}
+			}
+		}
+
+		err := populateArgFilterMap(pinPathPrefix, outerMap, uint32(i), data)
 		if err != nil {
 			return err
 		}