pkg/aflow: provide reusable helpers to extract crash type

Lots of workflows may want to have special handling for particular bug types.
Provide common helpers that make it easy to act on bug type.

3 files changed, 36 insertions, 26 deletions

diff --git a/pkg/aflow/flow/assessment/moderation.go b/pkg/aflow/flow/assessment/moderation.go
index f5b7a4bfb..d85e9eb2a 100644
--- a/pkg/aflow/flow/assessment/moderation.go
+++ b/pkg/aflow/flow/assessment/moderation.go
@@ -4,13 +4,10 @@
 package assessmenet
 
 import (
-	"fmt"
-
 	"github.com/google/syzkaller/pkg/aflow"
 	"github.com/google/syzkaller/pkg/aflow/action/kernel"
 	"github.com/google/syzkaller/pkg/aflow/ai"
 	"github.com/google/syzkaller/pkg/aflow/tool/codesearcher"
-	"github.com/google/syzkaller/pkg/report/crash"
 )
 
 type moderationInputs struct {
@@ -34,7 +31,6 @@ func init() {
 		"assess if a bug report is consistent and actionable or not",
 		&aflow.Flow{
 			Root: aflow.Pipeline(
-				aflow.NewFuncAction("extract-crash-type", extractCrashType),
 				kernel.Checkout,
 				kernel.Build,
 				codesearcher.PrepareIndex,
@@ -62,7 +58,7 @@ report is actionable or not. Actionable means that it contains enough info to ro
 the underlying bug, and that the report is self-consistent and makes sense, rather than
 e.g. a one-off nonsensical crash induced by a previous memory corruption.
 
-{{if .IsUAF}}
+{{if titleIsUAF .BugTitle}}
 The bug report is about a use-after-free bug generated by KASAN tool.
 It should contain 3 stack traces: the bad memory access stack, the heap block allocation stack,
 and the heap block free stack. If the report does not contain 3 stacks, it's not actionable.
@@ -91,23 +87,3 @@ The bug report is:
 
 {{.CrashReport}}
 `
-
-type extractArgs struct {
-	BugTitle string
-}
-
-type extractResult struct {
-	IsUAF bool
-}
-
-func extractCrashType(ctx *aflow.Context, args extractArgs) (extractResult, error) {
-	var res extractResult
-	typ := crash.TitleToType(args.BugTitle)
-	switch {
-	case typ.IsUAF():
-		res.IsUAF = true
-	default:
-		return res, fmt.Errorf("unsupported bug type")
-	}
-	return res, nil
-}
diff --git a/pkg/aflow/template.go b/pkg/aflow/template.go
index cf8417bb9..d853abe46 100644
--- a/pkg/aflow/template.go
+++ b/pkg/aflow/template.go
@@ -8,9 +8,12 @@ import (
 	"fmt"
 	"io"
 	"reflect"
+	"slices"
 	"strings"
 	"text/template"
 	"text/template/parse"
+
+	"github.com/google/syzkaller/pkg/report/crash"
 )
 
 // formatTemplate formats template 'text' using the standard text/template logic.
@@ -101,5 +104,16 @@ func walkTemplate(n parse.Node, used map[string]bool, errp *error) {
 }
 
 func parseTemplate(prompt string) (*template.Template, error) {
-	return template.New("").Option("missingkey=error").Parse(prompt)
+	return template.New("").Option("missingkey=error").Funcs(templateFuncs).Parse(prompt)
+}
+
+var templateFuncs = template.FuncMap{
+	"titleIsUAF":     titleIs(crash.KASANUseAfterFreeRead, crash.KASANUseAfterFreeWrite),
+	"titleIsWarning": titleIs(crash.Warning),
+}
+
+func titleIs(types ...crash.Type) func(string) bool {
+	return func(title string) bool {
+		return slices.Contains(types, crash.TitleToType(title))
+	}
 }
diff --git a/pkg/aflow/template_test.go b/pkg/aflow/template_test.go
index 2de13a2c1..8e0e75804 100644
--- a/pkg/aflow/template_test.go
+++ b/pkg/aflow/template_test.go
@@ -4,6 +4,7 @@
 package aflow
 
 import (
+	"bytes"
 	"fmt"
 	"maps"
 	"reflect"
@@ -85,3 +86,22 @@ func TestTemplate(t *testing.T) {
 		})
 	}
 }
+
+func TestTemplateRender(t *testing.T) {
+	data := map[string]any{
+		"Title": "WARNING: something is wrong",
+	}
+	const text = `
+{{if titleIsUAF .Title}}It is UAF.{{end}}
+{{if titleIsWarning .Title}}It is WARNING.{{end}}
+`
+	const want = `
+
+It is WARNING.
+`
+	templ, err := parseTemplate(text)
+	require.NoError(t, err)
+	buf := new(bytes.Buffer)
+	require.NoError(t, templ.Execute(buf, data))
+	require.Equal(t, want, buf.String())
+}