struct_ops: fix to enforce specify a btf.Int as Key

Tests are updated to use Key instead of KeySize/ValueSize literals

Signed-off-by: shun159 <[email protected]>

Author: shun159

collection_test.go

@@ -791,10 +791,9 @@ func TestStructOpsMapSpecSimpleLoadAndAssign(t *testing.T) {
 				Name:       "testmod_ops",
 				Type:       StructOpsMap,
 				Flags:      sys.BPF_F_LINK,
-				KeySize:    4,
-				ValueSize:  448,
-				MaxEntries: 1,
+				Key:        &btf.Int{Size: 4},
 				Value:      &btf.Struct{Name: "bpf_struct_ops_bpf_testmod_ops"},
+				MaxEntries: 1,
 				Contents: []MapKV{
 					{
 						Key:   uint32(0),

map.go

@@ -554,42 +554,61 @@ func (spec *MapSpec) createMap(inner *sys.FD, c *btf.Cache) (_ *Map, err error)
 		defer handle.Close()
 
 		if spec.Type == StructOpsMap {
+			if handle == nil {
+				return nil, fmt.Errorf("struct_ops: BTF handle is not resolved")
+			}
+
 			if spec.Value == nil {
 				return nil, fmt.Errorf("struct_ops map: missing value type information")
 			}
 
-			valueType, ok := btf.As[*btf.Struct](spec.Value)
+			if spec.Key == nil {
+				return nil, fmt.Errorf("struct_ops map: Key must be specified")
+			}
+
+			keyType, ok := btf.As[*btf.Int](spec.Key)
 			if !ok {
-				return nil, fmt.Errorf("value must be Struct type")
+				return nil, fmt.Errorf("key must be Int type")
 			}
 
-			if spec.KeySize != 4 {
-				return nil, fmt.Errorf("struct_ops: KeySize must be 4")
+			// The key size of StructOps is always 4
+			if keyType.Size != 4 {
+				return nil, fmt.Errorf("key size must be 4")
+			}
+
+			valueType, ok := btf.As[*btf.Struct](spec.Value)
+			if !ok {
+				return nil, fmt.Errorf("value must be Struct type")
 			}
 
 			// struct_ops: resolve value type ("bpf_struct_ops_<name>") and
 			// record kernel-specific BTF IDs / FDs needed for map creation.
-			target := btf.Type((*btf.Struct)(nil))
-			s, module, err := findTargetInKernel(valueType.Name, &target, c)
+			vTarget := btf.Type((*btf.Struct)(nil))
+			s, module, err := findTargetInKernel(valueType.Name, &vTarget, c)
 			if err != nil {
 				return nil, fmt.Errorf("lookup value type %q: %w", valueType.Name, err)
 			}
 			defer module.Close()
 
-			vType := target.(*btf.Struct)
-
+			vType := vTarget.(*btf.Struct)
 			btfValueTypeId, err := s.TypeID(vType)
 			if err != nil {
 				return nil, fmt.Errorf("lookup type_id: %w", err)
 			}
 
-			attr.ValueSize = spec.ValueSize
+			attr.BtfFd = uint32(handle.FD())
 			attr.BtfVmlinuxValueTypeId = btfValueTypeId
-
-			if handle == nil {
-				return nil, fmt.Errorf("struct_ops: BTF handle is not resolved")
+			attr.KeySize, spec.KeySize = keyType.Size, keyType.Size
+
+			if spec.ValueSize == 0 {
+				attr.ValueSize, spec.ValueSize = vType.Size, vType.Size
+			} else {
+				// If spec.ValueSize is non-zero but does not match vtype.size, it should fail.
+				if spec.ValueSize != vType.Size {
+					return nil, fmt.Errorf("struct_ops: ValueSize (%d) mismatches vtype size (%d)", spec.ValueSize, vType.Size)
+				}
+				attr.ValueSize = spec.ValueSize
 			}
-			attr.BtfFd = uint32(handle.FD())
 
 			if module != nil {
 				// BPF_F_VTYPE_BTF_OBJ_FD is required if the type comes from a module

struct_ops_test.go

@@ -15,10 +15,9 @@ func TestCreateStructOpsMapSpecSimple(t *testing.T) {
 		Name:       "testmod_ops",
 		Type:       StructOpsMap,
 		Flags:      sys.BPF_F_LINK,
-		KeySize:    4,
-		ValueSize:  448,
-		MaxEntries: 1,
+		Key:        &btf.Int{Size: 4},
 		Value:      &btf.Struct{Name: "bpf_struct_ops_bpf_testmod_ops"},
+		MaxEntries: 1,
 		Contents: []MapKV{
 			{
 				Key:   uint32(0),