fix: support slices of enums and fix NfTablesConntrackStateMatch

We already have the code which supports custom enums, so let's extend it to support custom enums in slices and fix the NfTablesConntrackStateMatch proto definition. Signed-off-by: Dmitriy Matrenichev <[email protected]>
siderolabs · Nov 29, 2023 · 0b111ec · 0b111ec
1 parent 9a85217
commit 0b111ec
Show file tree

Hide file tree

Showing 11 changed files with 521 additions and 517 deletions.
diff --git a/api/api.descriptors b/api/api.descriptors
diff --git a/api/resource/definitions/network/network.proto b/api/resource/definitions/network/network.proto
@@ -193,7 +193,7 @@ message NfTablesClampMSS {
 
 // NfTablesConntrackStateMatch describes the match on the connection tracking state.
 message NfTablesConntrackStateMatch {
-  repeated uint32 states = 1;
+  repeated talos.resource.definitions.enums.NethelpersConntrackState states = 1;
 }
 
 // NfTablesIfNameMatch describes the match on the interface name.

diff --git a/hack/structprotogen/proto/proto.go b/hack/structprotogen/proto/proto.go
@@ -270,7 +270,12 @@ func PrepareProtoData(pkgsTypes slices.Sorted[*types.Type], constants consts.Con
 			if fieldTyp, ok := types.MatchTypeData[types.Slice](fieldTypeData); ok {
 				var importName, typeName string
 
+				block, isEnum := constants.Get(fieldTyp.Pkg, fieldTyp.Name)
+
 				switch {
+				case isEnum:
+					importName = "resource/definitions/enums/enums.proto"
+					typeName = "repeated talos.resource.definitions.enums." + block.ProtoMessageName()
 				case fieldTyp.Pkg == "" && fieldTyp.Name == "byte" && fieldTyp.Is2DSlice: //nolint:goconst
 					typeName = "repeated bytes"
 				case fieldTyp.Pkg == "" && fieldTyp.Name == "byte":
@@ -345,7 +350,7 @@ func PrepareProtoData(pkgsTypes slices.Sorted[*types.Type], constants consts.Con
 func mustFormatTypeName(fieldTypePkg string, fieldType string, declPkg string) (string, string) {
 	importPath, name := formatTypeName(fieldTypePkg, fieldType, declPkg)
 	if name == "" {
-		panic(fmt.Errorf("unknown type %s.%s", fieldType, fieldTypePkg))
+		panic(fmt.Errorf("unknown type %s.%s", fieldTypePkg, fieldType))
 	}
 
 	return importPath, name

diff --git a/internal/app/machined/pkg/adapters/network/nftables_rule.go b/internal/app/machined/pkg/adapters/network/nftables_rule.go
@@ -311,7 +311,7 @@ func (a nftablesRule) Compile() (*NfTablesCompiled, error) {
 					SourceRegister: 1,
 					DestRegister:   1,
 					Len:            4,
-					Mask:           binaryutil.NativeEndian.PutUint32(match.States[0]),
+					Mask:           binaryutil.NativeEndian.PutUint32(uint32(match.States[0])),
 					Xor:            []byte{0x0, 0x0, 0x0, 0x0},
 				},
 				// [ cmp neq reg 1 0x00000000 ]
@@ -325,7 +325,7 @@ func (a nftablesRule) Compile() (*NfTablesCompiled, error) {
 			result.Sets = append(result.Sets,
 				NfTablesSet{
 					Kind:            SetKindConntrackState,
-					ConntrackStates: xslices.Map(match.States, func(s uint32) nethelpers.ConntrackState { return nethelpers.ConntrackState(s) }),
+					ConntrackStates: match.States,
 				})
 
 			rulePre = append(rulePre,

diff --git a/internal/app/machined/pkg/adapters/network/nftables_rule_test.go b/internal/app/machined/pkg/adapters/network/nftables_rule_test.go
@@ -645,8 +645,8 @@ func TestNfTablesRuleCompile(t *testing.T) { //nolint:tparallel
 			name: "ct state",
 			spec: networkres.NfTablesRule{
 				MatchConntrackState: &networkres.NfTablesConntrackStateMatch{
-					States: []uint32{
-						uint32(nethelpers.ConntrackStateInvalid),
+					States: []nethelpers.ConntrackState{
+						nethelpers.ConntrackStateInvalid,
 					},
 				},
 			},
@@ -675,9 +675,9 @@ func TestNfTablesRuleCompile(t *testing.T) { //nolint:tparallel
 			name: "ct states",
 			spec: networkres.NfTablesRule{
 				MatchConntrackState: &networkres.NfTablesConntrackStateMatch{
-					States: []uint32{
-						uint32(nethelpers.ConntrackStateRelated),
-						uint32(nethelpers.ConntrackStateEstablished),
+					States: []nethelpers.ConntrackState{
+						nethelpers.ConntrackStateRelated,
+						nethelpers.ConntrackStateEstablished,
 					},
 				},
 			},

diff --git a/internal/app/machined/pkg/controllers/network/nftables_chain_test.go b/internal/app/machined/pkg/controllers/network/nftables_chain_test.go
@@ -169,17 +169,17 @@ func (s *NfTablesChainSuite) TestConntrackCounter() {
 	chain.TypedSpec().Rules = []network.NfTablesRule{
 		{
 			MatchConntrackState: &network.NfTablesConntrackStateMatch{
-				States: []uint32{
-					uint32(nethelpers.ConntrackStateEstablished),
-					uint32(nethelpers.ConntrackStateRelated),
+				States: []nethelpers.ConntrackState{
+					nethelpers.ConntrackStateEstablished,
+					nethelpers.ConntrackStateRelated,
 				},
 			},
 			Verdict: pointer.To(nethelpers.VerdictAccept),
 		},
 		{
 			MatchConntrackState: &network.NfTablesConntrackStateMatch{
-				States: []uint32{
-					uint32(nethelpers.ConntrackStateInvalid),
+				States: []nethelpers.ConntrackState{
+					nethelpers.ConntrackStateInvalid,
 				},
 			},
 			AnonCounter: true,

diff --git a/pkg/machinery/api/resource/definitions/network/network.pb.go b/pkg/machinery/api/resource/definitions/network/network.pb.go
diff --git a/pkg/machinery/api/resource/definitions/network/network_vtproto.pb.go b/pkg/machinery/api/resource/definitions/network/network_vtproto.pb.go
diff --git a/pkg/machinery/resources/network/deep_copy.generated.go b/pkg/machinery/resources/network/deep_copy.generated.go
diff --git a/pkg/machinery/resources/network/nftables_chain.go b/pkg/machinery/resources/network/nftables_chain.go
@@ -133,8 +133,7 @@ type NfTablesLimitMatch struct {
 //
 //gotagsrewrite:gen
 type NfTablesConntrackStateMatch struct {
-	// TODO: should be []nethelpers.ConntrackState, but structprotogen needs to be fixed to support it.
-	States []uint32 `yaml:"states" protobuf:"1"`
+	States []nethelpers.ConntrackState `yaml:"states" protobuf:"1"`
 }
 
 // NewNfTablesChain initializes a NfTablesChain resource.

diff --git a/website/content/v1.6/reference/api.md b/website/content/v1.6/reference/api.md
@@ -3095,7 +3095,7 @@ NfTablesConntrackStateMatch describes the match on the connection tracking state
 
 | Field | Type | Label | Description |
 | ----- | ---- | ----- | ----------- |
-| states | [uint32](#uint32) | repeated |  |
+| states | [talos.resource.definitions.enums.NethelpersConntrackState](#talos.resource.definitions.enums.NethelpersConntrackState) | repeated |  |