Add selectors builder

selectors.go

@@ -2,6 +2,7 @@ package manager
 
 import (
 	"fmt"
+	"regexp"
 	"strings"
 )
 
@@ -199,3 +200,94 @@ func (be *BestEffort) EditProbeIdentificationPair(old ProbeIdentificationPair, n
 		selector.EditProbeIdentificationPair(old, new)
 	}
 }
+
+type ProbesSelectorBuilder struct {
+	selectors []ProbesSelector
+	currSelectorSlice *[]ProbesSelector
+}
+
+func NewProbesSelectorBuilder() *ProbesSelectorBuilder {
+	b := &ProbesSelectorBuilder{
+		selectors: make([]ProbesSelector, 0),
+	}
+
+	b.currSelectorSlice = &b.selectors
+	return b
+}
+
+type ProbeSelectorLocation int
+const (
+	ProbeSelectorLocationRoot ProbeSelectorLocation = iota
+	ProbeSelectorLocationNested
+)
+
+func (b *ProbesSelectorBuilder) AllOf(location ProbeSelectorLocation) *ProbesSelectorBuilder {
+	allof := &AllOf{}
+	if location == ProbeSelectorLocationRoot {
+		b.selectors = append(b.selectors, allof)
+	} else {
+		*b.currSelectorSlice = append(*b.currSelectorSlice, allof)
+	}
+	b.currSelectorSlice = &allof.Selectors
+
+	return b
+}
+
+func (b *ProbesSelectorBuilder) OneOf(location ProbeSelectorLocation) *ProbesSelectorBuilder {
+	oneof := &OneOf{}
+	if location == ProbeSelectorLocationRoot {
+		b.selectors = append(b.selectors, oneof)
+	} else {
+		*b.currSelectorSlice = append(*b.currSelectorSlice, oneof)
+	}
+	b.currSelectorSlice = &oneof.Selectors
+
+	return b
+}
+
+func (b *ProbesSelectorBuilder) BestEffort(location ProbeSelectorLocation) *ProbesSelectorBuilder {
+	be := &BestEffort{}
+	if location == ProbeSelectorLocationRoot {
+		b.selectors = append(b.selectors, be)
+	} else {
+		*b.currSelectorSlice = append(*b.currSelectorSlice, be)
+	}
+	b.currSelectorSlice = &be.Selectors
+
+	return b
+}
+
+type ProbeIdOptions int
+
+const (
+	ProbeIdAddRetprobe ProbeIdOptions = iota
+)
+
+var addRetRegex = regexp.MustCompile(`^(k|u)probe`)
+
+func (b *ProbesSelectorBuilder) ProbeID(funcName string, opts ...ProbeIdOptions) *ProbesSelectorBuilder {
+	pi := &ProbeIdentificationPair{
+		EBPFFuncName: funcName,
+	}
+
+	*b.currSelectorSlice = append(*b.currSelectorSlice, &ProbeSelector{ProbeIdentificationPair: *pi})
+
+	for _, opt := range opts {
+		switch opt {
+		case ProbeIdAddRetprobe:
+			modified := addRetRegex.ReplaceAllString(funcName, "${1}retprobe")
+
+			if modified == funcName {
+				panic("ProbeIdAddRetprobe can only be used with kprobe or uprobe")
+			}
+
+			b.ProbeID(modified)
+		}
+	}
+
+	return b
+}
+
+func (b *ProbesSelectorBuilder) Build() []ProbesSelector {
+	return b.selectors
+}

selectors_test.go

@@ -0,0 +1,56 @@
+package manager
+
+import "testing"
+
+func TestProbesSelectorBuilder(t *testing.T) {
+	selectors := (NewProbesSelectorBuilder().
+		AllOf(ProbeSelectorLocationRoot).
+			ProbeID("uprobe__func1").
+			ProbeID("uprobe__func2", ProbeIdAddRetprobe).
+		BestEffort(ProbeSelectorLocationRoot).
+			ProbeID("kprobe__func3").
+			ProbeID("kprobe__func4", ProbeIdAddRetprobe).
+			OneOf(ProbeSelectorLocationNested).
+				ProbeID("kprobe__func5").
+				ProbeID("kprobe__func6")).Build()
+
+	if len(selectors) != 2 {
+		t.Fatalf("expected 2 selectors, got %d", len(selectors))
+	}
+
+	allof, ok := selectors[0].(*AllOf)
+	if !ok {
+		t.Fatalf("expected AllOf, got %T", selectors[0])
+	}
+
+	// Testing with the String() method is easier than iterating, casting and comparing each elemet
+	// and results are the same
+	expectedString := "AllOf {UID: EBPFFuncName:uprobe__func1}, {UID: EBPFFuncName:uprobe__func2}, {UID: EBPFFuncName:uretprobe__func2}"
+	if allof.String() != expectedString {
+		t.Fatalf("expected: %s, got: %s", expectedString, allof.String())
+	}
+
+	be, ok := selectors[1].(*BestEffort)
+	if !ok {
+		t.Fatalf("expected BestEffort, got %T", selectors[0])
+	}
+
+	if len(be.Selectors) != 4 {
+		t.Fatalf("expected 4 selectors, got %d", len(be.Selectors))
+	}
+
+	oneof, ok := be.Selectors[3].(*OneOf)
+	if !ok {
+		t.Fatalf("expected OneOf, got %T", be.Selectors[3])
+	}
+
+	expectedString = "OneOf {UID: EBPFFuncName:kprobe__func5}, {UID: EBPFFuncName:kprobe__func6}"
+	if oneof.String() != expectedString {
+		t.Fatalf("expected: %s, got: %s", expectedString, oneof.String())
+	}
+
+	expectedString = "BestEffort {UID: EBPFFuncName:kprobe__func3}, {UID: EBPFFuncName:kprobe__func4}, {UID: EBPFFuncName:kretprobe__func4}, {UID: EBPFFuncName:kprobe__func5}, {UID: EBPFFuncName:kprobe__func6}"
+	if be.String() != expectedString {
+		t.Fatalf("expected: %s, got: %s", expectedString, be.String())
+	}
+}