Add flag accept-tcp-flag to network delay (#195)

* add flag  to network delay

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

* edit some details

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

* update test

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

* change the logic of adding chains

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

* edit some details

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

Co-authored-by: root <[email protected]>

Author: FingerLeader

cmd/attack/network.go

@@ -84,6 +84,7 @@ func NewNetworkDelayCommand(dep fx.Option, options *core.NetworkCommand) *cobra.
 	cmd.Flags().StringVarP(&options.Hostname, "hostname", "H", "", "only impact traffic to these hostnames")
 	cmd.Flags().StringVarP(&options.IPProtocol, "protocol", "p", "",
 		"only impact traffic using this IP protocol, supported: tcp, udp, icmp, all")
+	cmd.Flags().StringVarP(&options.AcceptTCPFlags, "accept-tcp-flags", "", "", "only the packet which match the tcp flag can be accepted, others will be dropped. only set when the protocol is tcp.")
 
 	return cmd
 }

pkg/core/network.go

@@ -137,6 +137,10 @@ func (n *NetworkCommand) validNetworkDelay() error {
 		return errors.Errorf("ip addressed %s not valid", n.IPAddress)
 	}
 
+	if len(n.AcceptTCPFlags) > 0 && n.IPProtocol != "tcp" {
+		return errors.Errorf("protocol should be 'tcp' when set accept-tcp-flags")
+	}
+
 	return checkProtocolAndPorts(n.IPProtocol, n.SourcePort, n.EgressPort)
 }
 
@@ -519,24 +523,20 @@ func (n *NetworkCommand) NeedApplyTC() bool {
 	}
 }
 
-func (n *NetworkCommand) PartitionChain(ipset string) ([]*pb.Chain, error) {
-	if n.Action != NetworkPartitionAction {
-		return nil, nil
-	}
-
+func (n *NetworkCommand) AdditionalChain(ipset string) ([]*pb.Chain, error) {
 	chains := make([]*pb.Chain, 0, 2)
 	var toChains, fromChains []*pb.Chain
 	var err error
 
 	if n.Direction == "to" || n.Direction == "both" {
-		toChains, err = n.getPartitionChain(ipset, "to")
+		toChains, err = n.getAdditionalChain(ipset, "to")
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	if n.Direction == "from" || n.Direction == "both" {
-		fromChains, err = n.getPartitionChain(ipset, "from")
+		fromChains, err = n.getAdditionalChain(ipset, "from")
 		if err != nil {
 			return nil, err
 		}
@@ -548,7 +548,7 @@ func (n *NetworkCommand) PartitionChain(ipset string) ([]*pb.Chain, error) {
 	return chains, nil
 }
 
-func (n *NetworkCommand) getPartitionChain(ipset, direction string) ([]*pb.Chain, error) {
+func (n *NetworkCommand) getAdditionalChain(ipset, direction string) ([]*pb.Chain, error) {
 	var directionStr string
 	var directionChain pb.Chain_Direction
 	if direction == "to" {
@@ -573,14 +573,15 @@ func (n *NetworkCommand) getPartitionChain(ipset, direction string) ([]*pb.Chain
 		})
 	}
 
-	chains = append(chains, &pb.Chain{
-		Name:      fmt.Sprintf("%s/1", directionStr),
-		Ipsets:    []string{ipset},
-		Direction: directionChain,
-		Protocol:  n.IPProtocol,
-		Target:    "DROP",
-	})
-
+	if n.Action == NetworkPartitionAction {
+		chains = append(chains, &pb.Chain{
+			Name:      fmt.Sprintf("%s/1", directionStr),
+			Ipsets:    []string{ipset},
+			Direction: directionChain,
+			Protocol:  n.IPProtocol,
+			Target:    "DROP",
+		})
+	}
 	return chains, nil
 }
 
@@ -596,6 +597,13 @@ func (n *NetworkCommand) NeedApplyDNSServer() bool {
 	return len(n.DNSServer) > 0
 }
 
+func (n *NetworkCommand) NeedAdditionalChains() bool {
+	if n.Action != NetworkPartitionAction || (n.Action == NetworkDelayAction && len(n.AcceptTCPFlags) != 0) {
+		return true
+	}
+	return false
+}
+
 func NewNetworkCommand() *NetworkCommand {
 	return &NetworkCommand{
 		CommonAttackConfig: CommonAttackConfig{

pkg/core/network_test.go

@@ -130,7 +130,7 @@ func TestPatitionChain(t *testing.T) {
 			},
 		}
 		for _, tc := range testCases {
-			chains, err := tc.cmd.PartitionChain("test")
+			chains, err := tc.cmd.AdditionalChain("test")
 			if err != nil {
 				t.Errorf("failed to partition chain: %v", err)
 			}

pkg/server/chaosd/network.go

@@ -140,11 +140,14 @@ func (s *Server) applyIptables(attack *core.NetworkCommand, ipset, uid string) e
 		return perrors.WithStack(err)
 	}
 	chains := core.IptablesRuleList(iptables).ToChains()
-	newChains, err := attack.PartitionChain(ipset)
-	if err != nil {
-		return perrors.WithStack(err)
+	// Presently, only partition and delay with `accept-tcp-flags` need to add additional chains
+	if attack.NeedAdditionalChains() {
+		newChains, err := attack.AdditionalChain(ipset)
+		if err != nil {
+			return perrors.WithStack(err)
+		}
+		chains = append(chains, newChains...)
 	}
-	chains = append(chains, newChains...)
 
 	if _, err := s.svr.SetIptablesChains(context.Background(), &pb.IptablesChainsRequest{
 		Chains:  chains,