From 31179bc75602cbe8f0421b44f19ff1b960039644 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Thu, 28 Dec 2023 21:31:00 +0100
Subject: prog: support conditional fields

pkg/compiler restructures conditional fields in structures into unions,
so we only have to implement the support for unions.

Semantics is as follows:
If a union has conditions, syzkaller picks the first field whose
condition matches. Since we require the last union field to have no
conditions, we can always construct an object.

Changes from this commit aim at ensuring that the selected union fields
always follow the rule above.
---
 prog/size.go | 64 ++++++++++++++++++++++++++++++++++++++----------------------
 1 file changed, 41 insertions(+), 23 deletions(-)

(limited to 'prog/size.go')

diff --git a/prog/size.go b/prog/size.go
index 70ec678b6..6463821e0 100644
--- a/prog/size.go
+++ b/prog/size.go
@@ -45,20 +45,34 @@ func (target *Target) assignArgSize(arg Arg, args []Arg, fields []Field, parents
 	}
 }
 
-func (target *Target) assignSizeStruct(dst *ConstArg, buf Arg, path []string, parentsMap map[Arg]Arg) {
+func (target *Target) assignSize(dst *ConstArg, pos Arg, path []string, args []Arg,
+	fields []Field, parentsMap map[Arg]Arg, overlayField int) {
+	found := target.findArg(pos, path, args, fields, parentsMap, overlayField)
+	if found != nil && !found.isAnyPtr {
+		dst.Val = target.computeSize(found.arg, found.offset, dst.Type().(*LenType))
+	}
+}
+
+type foundArg struct {
+	arg      Arg
+	offset   uint64
+	isAnyPtr bool
+}
+
+func (target *Target) findFieldStruct(buf Arg, path []string, parentsMap map[Arg]Arg) *foundArg {
 	switch arg := buf.(type) {
 	case *GroupArg:
 		typ := arg.Type().(*StructType)
-		target.assignSize(dst, buf, path, arg.Inner, typ.Fields, parentsMap, typ.OverlayField)
+		return target.findArg(buf, path, arg.Inner, typ.Fields, parentsMap, typ.OverlayField)
 	case *UnionArg:
-		target.assignSize(dst, buf, path, nil, nil, parentsMap, 0)
+		return target.findArg(buf, path, nil, nil, parentsMap, 0)
 	default:
-		panic(fmt.Sprintf("unexpected arg type %v", arg))
+		panic(fmt.Sprintf("unexpected arg type %#v", arg))
 	}
 }
 
-func (target *Target) assignSize(dst *ConstArg, pos Arg, path []string, args []Arg,
-	fields []Field, parentsMap map[Arg]Arg, overlayField int) {
+func (target *Target) findArg(pos Arg, path []string, args []Arg, fields []Field,
+	parentsMap map[Arg]Arg, overlayField int) *foundArg {
 	elem := path[0]
 	path = path[1:]
 	var offset uint64
@@ -66,6 +80,9 @@ func (target *Target) assignSize(dst *ConstArg, pos Arg, path []string, args []A
 		if i == overlayField {
 			offset = 0
 		}
+		if buf == nil {
+			continue
+		}
 		if elem != fields[i].Name {
 			offset += buf.Size()
 			continue
@@ -74,46 +91,43 @@ func (target *Target) assignSize(dst *ConstArg, pos Arg, path []string, args []A
 			// If path points into squashed argument, we don't have the target argument.
 			// In such case we simply leave size argument as is. It can't happen during generation,
 			// only during mutation and mutation can set size to random values, so it should be fine.
-			return
+			return &foundArg{buf, offset, true}
 		}
 		buf = InnerArg(buf)
 		if buf == nil {
-			dst.Val = 0 // target is an optional pointer
-			return
+			return &foundArg{nil, offset, false}
 		}
 		if len(path) != 0 {
-			target.assignSizeStruct(dst, buf, path, parentsMap)
-			return
+			return target.findFieldStruct(buf, path, parentsMap)
 		}
-		dst.Val = target.computeSize(buf, offset, dst.Type().(*LenType))
-		return
+		return &foundArg{buf, offset, false}
 	}
 	if elem == ParentRef {
 		buf := parentsMap[pos]
 		if len(path) != 0 {
-			target.assignSizeStruct(dst, buf, path, parentsMap)
-			return
+			return target.findFieldStruct(buf, path, parentsMap)
 		}
-		dst.Val = target.computeSize(buf, noOffset, dst.Type().(*LenType))
-		return
+		return &foundArg{buf, noOffset, false}
 	}
 	for buf := parentsMap[pos]; buf != nil; buf = parentsMap[buf] {
 		if elem != buf.Type().TemplateName() {
 			continue
 		}
 		if len(path) != 0 {
-			target.assignSizeStruct(dst, buf, path, parentsMap)
-			return
+			return target.findFieldStruct(buf, path, parentsMap)
 		}
-		dst.Val = target.computeSize(buf, noOffset, dst.Type().(*LenType))
-		return
+		return &foundArg{buf, noOffset, false}
 	}
 	var fieldNames []string
 	for _, field := range fields {
 		fieldNames = append(fieldNames, field.Name)
 	}
-	panic(fmt.Sprintf("len field %q references non existent field %q, pos=%q, argsMap: %v, path: %v",
-		dst.Type().Name(), elem, pos.Type().Name(), fieldNames, path))
+	posName := "nil"
+	if pos != nil {
+		posName = pos.Type().Name()
+	}
+	panic(fmt.Sprintf("path references non existent field %q, pos=%q, argsMap: %v, path: %v",
+		elem, posName, fieldNames, path))
 }
 
 const noOffset = ^uint64(0)
@@ -125,6 +139,10 @@ func (target *Target) computeSize(arg Arg, offset uint64, lenType *LenType) uint
 		}
 		return offset * 8 / lenType.BitSize
 	}
+	if arg == nil {
+		// For e.g. optional pointers.
+		return 0
+	}
 	bitSize := lenType.BitSize
 	if bitSize == 0 {
 		bitSize = 8
-- 
cgit mrf-deployment