Merge pull request #4 from ozontech-forks/fix_hostname_params

fix hostname params and add full-disable params (chaos-mesh#226)
ozontech-forks · Jan 30, 2023 · e09976c · e09976c
2 parents 14f3ff6 + 9497688
commit e09976c
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git a/cmd/attack/network.go b/cmd/attack/network.go
@@ -85,6 +85,7 @@ func NewNetworkDelayCommand(dep fx.Option, options *core.NetworkCommand) *cobra.
 	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.")
+	cmd.Flags().BoolVar(&options.FullDisable, "full-disable", false, "network on device full disable flag")
 
 	return cmd
 }
@@ -114,7 +115,7 @@ func NewNetworkLossCommand(dep fx.Option, options *core.NetworkCommand) *cobra.C
 	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().BoolVar(&options.FullDisable, "full-disable", false, "full net on device disable flag")
+	cmd.Flags().BoolVar(&options.FullDisable, "full-disable", false, "network on device full disable flag")
 
 	return cmd
 }
@@ -144,7 +145,7 @@ func NewNetworkCorruptCommand(dep fx.Option, options *core.NetworkCommand) *cobr
 	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().BoolVar(&options.FullDisable, "full-disable", false, "network on device full disable flag")
 	return cmd
 }
 
@@ -173,7 +174,7 @@ func NetworkDuplicateCommand(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().BoolVar(&options.FullDisable, "full-disable", false, "network on device full disable flag")
 	return cmd
 }
 

diff --git a/pkg/core/network.go b/pkg/core/network.go
@@ -143,6 +143,10 @@ func (n *NetworkCommand) validNetworkDelay() error {
 		return errors.Errorf("protocol should be 'tcp' when set accept-tcp-flags")
 	}
 
+	if err := checkNetworkLimitParams(n.Hostname, n.IPAddress, n.FullDisable); err != nil {
+		return err
+	}
+
 	return checkProtocolAndPorts(n.IPProtocol, n.SourcePort, n.EgressPort)
 }
 
@@ -151,6 +155,10 @@ func (n *NetworkCommand) validNetworkBandwidth() error {
 		return errors.Errorf("rate, limit and buffer both are required when action is bandwidth")
 	}
 
+	if err := checkNetworkLimitParams(n.Hostname, n.IPAddress, n.FullDisable); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -175,14 +183,8 @@ func (n *NetworkCommand) validNetworkCommon() error {
 		return errors.Errorf("ip addressed %s not valid", n.IPAddress)
 	}
 
-	if len(n.Hostname) == 0 && len(n.IPAddress) == 0 && !n.FullDisable {
-		return errors.New("hostname or ip address is required")
-	}
-
-	if n.FullDisable {
-		if len(n.Hostname) > 0 || len(n.IPAddress) > 0 {
-			return errors.New("the host and address are set, but the flag full-disable is enabled")
-		}
+	if err := checkNetworkLimitParams(n.Hostname, n.IPAddress, n.FullDisable); err != nil {
+		return err
 	}
 
 	return checkProtocolAndPorts(n.IPProtocol, n.SourcePort, n.EgressPort)
@@ -344,6 +346,20 @@ func checkProtocolAndPorts(p string, sports string, dports string) error {
 	return nil
 }
 
+func checkNetworkLimitParams(hostname string, ipaddress string, fullDisable bool) error {
+	if len(hostname) == 0 && len(ipaddress) == 0 && !fullDisable {
+		return errors.New("hostname or ip address is required")
+	}
+
+	if fullDisable {
+		if len(hostname) > 0 || len(ipaddress) > 0 {
+			return errors.New("the host and address are set, but the flag full-disable is enabled")
+		}
+	}
+
+	return nil
+}
+
 func (n NetworkCommand) RecoverData() string {
 	data, _ := json.Marshal(n)