[client] Mark netbird data plane traffic to identify interface traffic correctly (#3623)

client/firewall/nftables/router_linux.go

@@ -100,6 +100,10 @@ func (r *router) init(workTable *nftables.Table) error {
 		return fmt.Errorf("create containers: %w", err)
 	}
 
+	if err := r.setupDataPlaneMark(); err != nil {
+		log.Errorf("failed to set up data plane mark: %v", err)
+	}
+
 	return nil
 }
 
@@ -196,15 +200,21 @@ func (r *router) createContainers() error {
 		Type:     nftables.ChainTypeNAT,
 	})
 
-	// Chain is created by acl manager
-	// TODO: move creation to a common place
-	r.chains[chainNamePrerouting] = &nftables.Chain{
-		Name:     chainNamePrerouting,
+	r.chains[chainNameManglePostrouting] = r.conn.AddChain(&nftables.Chain{
+		Name:     chainNameManglePostrouting,
+		Table:    r.workTable,
+		Hooknum:  nftables.ChainHookPostrouting,
+		Priority: nftables.ChainPriorityMangle,
+		Type:     nftables.ChainTypeFilter,
+	})
+
+	r.chains[chainNameManglePrerouting] = r.conn.AddChain(&nftables.Chain{
+		Name:     chainNameManglePrerouting,
 		Table:    r.workTable,
 		Hooknum:  nftables.ChainHookPrerouting,
 		Priority: nftables.ChainPriorityMangle,
 		Type:     nftables.ChainTypeFilter,
-	}
+	})
 
 	// Add the single NAT rule that matches on mark
 	if err := r.addPostroutingRules(); err != nil {
@@ -220,7 +230,83 @@ func (r *router) createContainers() error {
 	}
 
 	if err := r.conn.Flush(); err != nil {
-		return fmt.Errorf("nftables: unable to initialize table: %v", err)
+		return fmt.Errorf("initialize tables: %v", err)
+	}
+
+	return nil
+}
+
+// setupDataPlaneMark configures the fwmark for the data plane
+func (r *router) setupDataPlaneMark() error {
+	if r.chains[chainNameManglePrerouting] == nil || r.chains[chainNameManglePostrouting] == nil {
+		return errors.New("no mangle chains found")
+	}
+
+	ctNew := getCtNewExprs()
+	preExprs := []expr.Any{
+		&expr.Meta{
+			Key:      expr.MetaKeyIIFNAME,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpEq,
+			Register: 1,
+			Data:     ifname(r.wgIface.Name()),
+		},
+	}
+	preExprs = append(preExprs, ctNew...)
+	preExprs = append(preExprs,
+		&expr.Immediate{
+			Register: 1,
+			Data:     binaryutil.NativeEndian.PutUint32(nbnet.DataPlaneMarkIn),
+		},
+		&expr.Ct{
+			Key:            expr.CtKeyMARK,
+			Register:       1,
+			SourceRegister: true,
+		},
+	)
+
+	preNftRule := &nftables.Rule{
+		Table: r.workTable,
+		Chain: r.chains[chainNameManglePrerouting],
+		Exprs: preExprs,
+	}
+	r.conn.AddRule(preNftRule)
+
+	postExprs := []expr.Any{
+		&expr.Meta{
+			Key:      expr.MetaKeyOIFNAME,
+			Register: 1,
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpEq,
+			Register: 1,
+			Data:     ifname(r.wgIface.Name()),
+		},
+	}
+	postExprs = append(postExprs, ctNew...)
+	postExprs = append(postExprs,
+		&expr.Immediate{
+			Register: 1,
+			Data:     binaryutil.NativeEndian.PutUint32(nbnet.DataPlaneMarkOut),
+		},
+		&expr.Ct{
+			Key:            expr.CtKeyMARK,
+			Register:       1,
+			SourceRegister: true,
+		},
+	)
+
+	postNftRule := &nftables.Rule{
+		Table: r.workTable,
+		Chain: r.chains[chainNameManglePostrouting],
+		Exprs: postExprs,
+	}
+	r.conn.AddRule(postNftRule)
+
+	if err := r.conn.Flush(); err != nil {
+		return fmt.Errorf("flush: %w", err)
 	}
 
 	return nil
@@ -516,26 +602,10 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 		op = expr.CmpOpNeq
 	}
 
-	exprs := []expr.Any{
-		// We only care about NEW connections to mark them and later identify them in the postrouting chain for masquerading.
-		// Masquerading will take care of the conntrack state, which means we won't need to mark established connections.
-		&expr.Ct{
-			Key:      expr.CtKeySTATE,
-			Register: 1,
-		},
-		&expr.Bitwise{
-			SourceRegister: 1,
-			DestRegister:   1,
-			Len:            4,
-			Mask:           binaryutil.NativeEndian.PutUint32(expr.CtStateBitNEW),
-			Xor:            binaryutil.NativeEndian.PutUint32(0),
-		},
-		&expr.Cmp{
-			Op:       expr.CmpOpNeq,
-			Register: 1,
-			Data:     []byte{0, 0, 0, 0},
-		},
-
+	// We only care about NEW connections to mark them and later identify them in the postrouting chain for masquerading.
+	// Masquerading will take care of the conntrack state, which means we won't need to mark established connections.
+	exprs := getCtNewExprs()
+	exprs = append(exprs,
 		// interface matching
 		&expr.Meta{
 			Key:      expr.MetaKeyIIFNAME,
@@ -546,7 +616,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 			Register: 1,
 			Data:     ifname(r.wgIface.Name()),
 		},
-	}
+	)
 
 	exprs = append(exprs, sourceExp...)
 	exprs = append(exprs, destExp...)
@@ -578,7 +648,7 @@ func (r *router) addNatRule(pair firewall.RouterPair) error {
 
 	r.rules[ruleKey] = r.conn.AddRule(&nftables.Rule{
 		Table:    r.workTable,
-		Chain:    r.chains[chainNamePrerouting],
+		Chain:    r.chains[chainNameManglePrerouting],
 		Exprs:    exprs,
 		UserData: []byte(ruleKey),
 	})
@@ -1324,3 +1394,24 @@ func applyPort(port *firewall.Port, isSource bool) []expr.Any {
 
 	return exprs
 }
+
+func getCtNewExprs() []expr.Any {
+	return []expr.Any{
+		&expr.Ct{
+			Key:      expr.CtKeySTATE,
+			Register: 1,
+		},
+		&expr.Bitwise{
+			SourceRegister: 1,
+			DestRegister:   1,
+			Len:            4,
+			Mask:           binaryutil.NativeEndian.PutUint32(expr.CtStateBitNEW),
+			Xor:            binaryutil.NativeEndian.PutUint32(0),
+		},
+		&expr.Cmp{
+			Op:       expr.CmpOpNeq,
+			Register: 1,
+			Data:     []byte{0, 0, 0, 0},
+		},
+	}
+}