bpf program fragment: Add support for checking possible return values.

This records the precise behavior of "return" instructions in the fragment,
rather than simply recording whether there was a "return" instruction at
all. In turn, this allows a caller to verify which return value is returned.

This is useful in an upcoming change to the BST struct, where the fragment
being recorded now has the potential to return (not just jump), and so it
becomes useful to check that the return values it can return are those we
expect.

PiperOrigin-RevId: 581422796

pkg/bpf/program_builder.go

@@ -19,6 +19,8 @@ import (
 	"math"
 	"sort"
 	"strings"
+
+	"gvisor.dev/gvisor/pkg/abi/linux"
 )
 
 const (
@@ -264,9 +266,13 @@ type FragmentOutcomes struct {
 	// fragment may jump to.
 	MayJumpToUnresolvedLabels map[string]struct{}
 
-	// MayReturn is true if executing the fragment may cause a return statement
-	// to be executed.
-	MayReturn bool
+	// MayReturnImmediate contains the set of possible immediate return values
+	// that the fragment may return.
+	MayReturnImmediate map[linux.BPFAction]struct{}
+
+	// MayReturnRegisterA is true if the fragment may return the value of
+	// register A.
+	MayReturnRegisterA bool
 }
 
 // String returns a list of possible human-readable outcomes.
@@ -275,21 +281,29 @@ func (o FragmentOutcomes) String() string {
 	if o.MayJumpToKnownOffsetBeyondFragment {
 		s = append(s, "may jump to known offset beyond fragment")
 	}
-	if o.MayJumpToUnresolvedLabels != nil {
-		sortedLabels := make([]string, 0, len(o.MayJumpToUnresolvedLabels))
-		for lbl := range o.MayJumpToUnresolvedLabels {
-			sortedLabels = append(sortedLabels, lbl)
-		}
-		sort.Strings(sortedLabels)
-		for _, lbl := range sortedLabels {
-			s = append(s, fmt.Sprintf("may jump to unresolved label %q", lbl))
-		}
+	sortedLabels := make([]string, 0, len(o.MayJumpToUnresolvedLabels))
+	for lbl := range o.MayJumpToUnresolvedLabels {
+		sortedLabels = append(sortedLabels, lbl)
+	}
+	sort.Strings(sortedLabels)
+	for _, lbl := range sortedLabels {
+		s = append(s, fmt.Sprintf("may jump to unresolved label %q", lbl))
 	}
 	if o.MayFallThrough {
 		s = append(s, "may fall through")
 	}
-	if o.MayReturn {
-		s = append(s, "may return")
+	sortedReturnValues := make([]uint32, 0, len(o.MayReturnImmediate))
+	for v := range o.MayReturnImmediate {
+		sortedReturnValues = append(sortedReturnValues, uint32(v))
+	}
+	sort.Slice(sortedReturnValues, func(i, j int) bool {
+		return sortedReturnValues[i] < sortedReturnValues[j]
+	})
+	for _, v := range sortedReturnValues {
+		s = append(s, fmt.Sprintf("may return '0x%x'", v))
+	}
+	if o.MayReturnRegisterA {
+		s = append(s, "may return register A")
 	}
 	if len(s) == 0 {
 		return "no outcomes (this should never happen)"
@@ -297,6 +311,11 @@ func (o FragmentOutcomes) String() string {
 	return strings.Join(s, ", ")
 }
 
+// MayReturn returns whether the fragment may return for any reason.
+func (o FragmentOutcomes) MayReturn() bool {
+	return len(o.MayReturnImmediate) > 0 || o.MayReturnRegisterA
+}
+
 // Outcomes returns the set of possible outcomes that executing this fragment
 // may result into.
 func (f ProgramFragment) Outcomes() FragmentOutcomes {
@@ -308,13 +327,19 @@ func (f ProgramFragment) Outcomes() FragmentOutcomes {
 	}
 	outcomes := FragmentOutcomes{
 		MayJumpToUnresolvedLabels: make(map[string]struct{}),
+		MayReturnImmediate:        make(map[linux.BPFAction]struct{}),
 	}
 	for pc := f.fromPC; pc < f.toPC; pc++ {
 		ins := f.b.instructions[pc]
 		isLastInstruction := pc == f.toPC-1
 		switch ins.OpCode & instructionClassMask {
 		case Ret:
-			outcomes.MayReturn = true
+			switch ins.OpCode {
+			case Ret | K:
+				outcomes.MayReturnImmediate[linux.BPFAction(ins.K)] = struct{}{}
+			case Ret | A:
+				outcomes.MayReturnRegisterA = true
+			}
 		case Jmp:
 			for _, offset := range ins.JumpOffsets() {
 				var foundLabelName string

pkg/bpf/program_builder_test.go

@@ -18,6 +18,8 @@ import (
 	"fmt"
 	"reflect"
 	"testing"
+
+	"gvisor.dev/gvisor/pkg/abi/linux"
 )
 
 func validate(p *ProgramBuilder, expected []Instruction) error {
@@ -190,6 +192,9 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 		if f.MayJumpToUnresolvedLabels == nil {
 			f.MayJumpToUnresolvedLabels = map[string]struct{}{}
 		}
+		if f.MayReturnImmediate == nil {
+			f.MayReturnImmediate = map[linux.BPFAction]struct{}{}
+		}
 		return f
 	}
 	for _, test := range []struct {
@@ -203,6 +208,9 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 		// by `build` alone.
 		wantLocal FragmentOutcomes
 
+		// Expected value of calling `MayReturn` on the local fragment.
+		wantLocalMayReturn bool
+
 		// Expected outcomes from recording the instructions added
 		// to the program since the test began.
 		wantOverall FragmentOutcomes
@@ -325,19 +333,44 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 			},
 		},
 		{
-			name: "add return",
+			name: "add immediate return",
 			build: func() {
 				p.AddStmt(Ret|K, 1337)
 			},
 			wantLocal: FragmentOutcomes{
-				MayReturn: true,
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
 			},
+			wantLocalMayReturn: true,
 			wantOverall: FragmentOutcomes{
 				MayJumpToUnresolvedLabels: map[string]struct{}{
 					"falselabel": struct{}{},
 				},
 				MayFallThrough: true, // From jump in previous test.
-				MayReturn:      true,
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
+			},
+		},
+		{
+			name: "add register A return",
+			build: func() {
+				p.AddStmt(Ret|A, 0)
+			},
+			wantLocal: FragmentOutcomes{
+				MayReturnRegisterA: true,
+			},
+			wantLocalMayReturn: true,
+			wantOverall: FragmentOutcomes{
+				MayJumpToUnresolvedLabels: map[string]struct{}{
+					"falselabel": struct{}{},
+				},
+				MayFallThrough: false, // Jump no longer pointing at end of fragment.
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
+				MayReturnRegisterA: true,
 			},
 		},
 		{
@@ -352,8 +385,11 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 				MayJumpToUnresolvedLabels: map[string]struct{}{
 					"falselabel": struct{}{},
 				},
-				MayReturn:      true,
-				MayFallThrough: true,
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
+				MayReturnRegisterA: true,
+				MayFallThrough:     true,
 			},
 		},
 		{
@@ -368,8 +404,11 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 				MayJumpToUnresolvedLabels: map[string]struct{}{
 					"falselabel": struct{}{},
 				},
-				MayReturn:      true,
-				MayFallThrough: true,
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
+				MayReturnRegisterA: true,
+				MayFallThrough:     true,
 			},
 		},
 		{
@@ -387,8 +426,11 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 				MayJumpToUnresolvedLabels: map[string]struct{}{
 					"falselabel": struct{}{},
 				},
-				MayReturn:      true,
-				MayFallThrough: true,
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
+				MayReturnRegisterA: true,
+				MayFallThrough:     true,
 			},
 		},
 		{
@@ -407,7 +449,10 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 				MayJumpToUnresolvedLabels: map[string]struct{}{
 					"falselabel": struct{}{},
 				},
-				MayReturn: true,
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
+				MayReturnRegisterA: true,
 			},
 		},
 		{
@@ -420,8 +465,11 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 			},
 			wantOverall: FragmentOutcomes{
 				MayJumpToKnownOffsetBeyondFragment: true,
-				MayReturn:                          true,
-				MayFallThrough:                     true,
+				MayReturnImmediate: map[linux.BPFAction]struct{}{
+					1337: struct{}{},
+				},
+				MayReturnRegisterA: true,
+				MayFallThrough:     true,
 			},
 		},
 	} {
@@ -429,9 +477,13 @@ func TestProgramBuilderOutcomes(t *testing.T) {
 			getLocalFragment := p.Record()
 			test.build()
 			localFragment := getLocalFragment()
-			if localOutcomes := localFragment.Outcomes(); !reflect.DeepEqual(fixup(localOutcomes), fixup(test.wantLocal)) {
+			localOutcomes := localFragment.Outcomes()
+			if !reflect.DeepEqual(fixup(localOutcomes), fixup(test.wantLocal)) {
 				t.Errorf("local fragment %v: got outcomes %v want %v", localFragment, localOutcomes, test.wantLocal)
 			}
+			if gotMayReturn := localOutcomes.MayReturn(); gotMayReturn != test.wantLocalMayReturn {
+				t.Errorf("local fragment MayReturn(): got %v want %v", gotMayReturn, test.wantLocalMayReturn)
+			}
 			overallFragment := getOverallFragment()
 			if overallOutcomes := overallFragment.Outcomes(); !reflect.DeepEqual(fixup(overallOutcomes), fixup(test.wantOverall)) {
 				t.Errorf("overall fragment %v: got outcomes %v want %v", overallFragment, overallOutcomes, test.wantOverall)

pkg/seccomp/seccomp.go

@@ -194,7 +194,7 @@ func (f syscallProgramFragment) MustHaveJumpedTo(labels ...label) {
 	if outcomes.MayFallThrough {
 		panic(fmt.Sprintf("fragment %v may fall through", fragment))
 	}
-	if outcomes.MayReturn {
+	if outcomes.MayReturn() {
 		panic(fmt.Sprintf("fragment %v may return", fragment))
 	}
 	if outcomes.MayJumpToKnownOffsetBeyondFragment {