network attack: support bandwidth limit (#91)

* support limit network bandwidth

Signed-off-by: xiang <xiang13225080@163.com>

cmd/attack/network.go

@@ -48,6 +48,7 @@ func NewNetworkAttackCommand(uid *string) *cobra.Command {
 		NetworkPartitionCommand(dep, options),
 		NetworkDNSCommand(dep, options),
 		NewNetworkPortOccupiedCommand(dep, options),
+		NewNetworkBandwidthCommand(dep, options),
 	)
 
 	return cmd
@@ -215,6 +216,30 @@ func NetworkDNSCommand(dep fx.Option, options *core.NetworkCommand) *cobra.Comma
 	return cmd
 }
 
+func NewNetworkBandwidthCommand(dep fx.Option, options *core.NetworkCommand) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "bandwidth",
+		Short: "limit network bandwidth",
+
+		Run: func(*cobra.Command, []string) {
+			options.Action = core.NetworkBandwidthAction
+			options.CompleteDefaults()
+			utils.FxNewAppWithoutLog(dep, fx.Invoke(commonNetworkAttackFunc)).Run()
+		},
+	}
+
+	cmd.Flags().StringVarP(&options.Rate, "rate", "r", "", "the speed knob, allows bps, kbps, mbps, gbps, tbps unit. bps means bytes per second")
+	cmd.Flags().Uint32VarP(&options.Limit, "limit", "l", 0, "the number of bytes that can be queued waiting for tokens to become available")
+	cmd.Flags().Uint32VarP(&options.Buffer, "buffer", "b", 0, "the maximum amount of bytes that tokens can be available for instantaneously")
+	cmd.Flags().Uint64VarP(options.Peakrate, "peakrate", "", 0, "the maximum depletion rate of the bucket")
+	cmd.Flags().Uint32VarP(options.Minburst, "minburst", "m", 0, "specifies the size of the peakrate bucket")
+	cmd.Flags().StringVarP(&options.Device, "device", "d", "", "the network interface to impact")
+	cmd.Flags().StringVarP(&options.IPAddress, "ip", "i", "", "only impact egress traffic to these IP addresses")
+	cmd.Flags().StringVarP(&options.Hostname, "hostname", "H", "", "only impact traffic to these hostnames")
+
+	return cmd
+}
+
 func commonNetworkAttackFunc(options *core.NetworkCommand, chaos *chaosd.Server) {
 	if err := options.Validate(); err != nil {
 		utils.ExitWithError(utils.ExitBadArgs, err)
@@ -234,7 +259,7 @@ func NewNetworkPortOccupiedCommand(dep fx.Option, options *core.NetworkCommand)
 		Short: "attack network port",
 
 		Run: func(cmd *cobra.Command, args []string) {
-			options.Action = core.NetworkPortOccupied
+			options.Action = core.NetworkPortOccupiedAction
 			options.CompleteDefaults()
 			utils.FxNewAppWithoutLog(dep, fx.Invoke(commonNetworkAttackFunc)).Run()
 		},

pkg/core/network.go

@@ -20,7 +20,9 @@ import (
 	"strings"
 	"time"
 
+	"github.com/chaos-mesh/chaos-mesh/api/v1alpha1"
 	"github.com/chaos-mesh/chaos-mesh/pkg/chaosdaemon/pb"
+	"github.com/chaos-mesh/chaos-mesh/pkg/netem"
 	"github.com/pingcap/errors"
 
 	"github.com/chaos-mesh/chaosd/pkg/utils"
@@ -49,6 +51,7 @@ type NetworkCommand struct {
 	DNSIp     string
 	DNSHost   string
 
+	*BandwidthSpec `json:",inline"`
 	// only the packet which match the tcp flag can be accepted, others will be dropped.
 	// only set when the IPProtocol is tcp, used for partition.
 	AcceptTCPFlags string
@@ -57,13 +60,14 @@ type NetworkCommand struct {
 var _ AttackConfig = &NetworkCommand{}
 
 const (
-	NetworkDelayAction     = "delay"
-	NetworkLossAction      = "loss"
-	NetworkCorruptAction   = "corrupt"
-	NetworkDuplicateAction = "duplicate"
-	NetworkDNSAction       = "dns"
-	NetworkPartitionAction = "partition"
-	NetworkPortOccupied    = "occupied"
+	NetworkDelayAction        = "delay"
+	NetworkLossAction         = "loss"
+	NetworkCorruptAction      = "corrupt"
+	NetworkDuplicateAction    = "duplicate"
+	NetworkDNSAction          = "dns"
+	NetworkPartitionAction    = "partition"
+	NetworkBandwidthAction    = "bandwidth"
+	NetworkPortOccupiedAction = "occupied"
 )
 
 func (n *NetworkCommand) Validate() error {
@@ -79,8 +83,10 @@ func (n *NetworkCommand) Validate() error {
 		return n.validNetworkDNS()
 	case NetworkPartitionAction:
 		return n.validNetworkPartition()
-	case NetworkPortOccupied:
+	case NetworkPortOccupiedAction:
 		return n.validNetworkOccupied()
+	case NetworkBandwidthAction:
+		return n.validNetworkBandwidth()
 	default:
 		return errors.Errorf("network action %s not supported", n.Action)
 	}
@@ -116,6 +122,14 @@ func (n *NetworkCommand) validNetworkDelay() error {
 	return checkProtocolAndPorts(n.IPProtocol, n.SourcePort, n.EgressPort)
 }
 
+func (n *NetworkCommand) validNetworkBandwidth() error {
+	if len(n.Rate) == 0 || n.Limit == 0 || n.Buffer == 0 {
+		return errors.Errorf("rate, limit and buffer both are required when action is bandwidth")
+	}
+
+	return nil
+}
+
 func (n *NetworkCommand) validNetworkCommon() error {
 	if len(n.Percent) == 0 {
 		return errors.New("percent is required")
@@ -325,6 +339,26 @@ func (n *NetworkCommand) ToDuplicateNetem() (*pb.Netem, error) {
 }
 
 func (n *NetworkCommand) ToTC(ipset string) (*pb.Tc, error) {
+	if n.Action == NetworkBandwidthAction {
+		tbf, err := netem.FromBandwidth(&v1alpha1.BandwidthSpec{
+			Rate:     n.Rate,
+			Limit:    n.Limit,
+			Buffer:   n.Buffer,
+			Peakrate: n.Peakrate,
+			Minburst: n.Minburst,
+		})
+
+		if err != nil {
+			return nil, err
+		}
+
+		return &pb.Tc{
+			Type:  pb.Tc_BANDWIDTH,
+			Tbf:   tbf,
+			Ipset: ipset,
+		}, nil
+	}
+
 	tc := &pb.Tc{
 		Type:       pb.Tc_NETEM,
 		Ipset:      ipset,
@@ -405,7 +439,7 @@ func (n *NetworkCommand) NeedApplyIptables() bool {
 
 func (n *NetworkCommand) NeedApplyTC() bool {
 	switch n.Action {
-	case NetworkDelayAction, NetworkLossAction, NetworkCorruptAction, NetworkDuplicateAction:
+	case NetworkDelayAction, NetworkLossAction, NetworkCorruptAction, NetworkDuplicateAction, NetworkBandwidthAction:
 		return true
 	default:
 		return false
@@ -468,5 +502,9 @@ func NewNetworkCommand() *NetworkCommand {
 		CommonAttackConfig: CommonAttackConfig{
 			Kind: NetworkAttack,
 		},
+		BandwidthSpec: &BandwidthSpec{
+			Peakrate: new(uint64),
+			Minburst: new(uint32),
+		},
 	}
 }

pkg/server/chaosd/network.go

@@ -63,10 +63,10 @@ func (networkAttack) Attack(options core.AttackConfig, env Environment) (err err
 			}
 		}
 
-	case core.NetworkPortOccupied:
+	case core.NetworkPortOccupiedAction:
 		return env.Chaos.applyPortOccupied(attack)
 
-	case core.NetworkDelayAction, core.NetworkLossAction, core.NetworkCorruptAction, core.NetworkDuplicateAction, core.NetworkPartitionAction:
+	case core.NetworkDelayAction, core.NetworkLossAction, core.NetworkCorruptAction, core.NetworkDuplicateAction, core.NetworkBandwidthAction:
 		if attack.NeedApplyIPSet() {
 			ipsetName, err = env.Chaos.applyIPSet(attack, env.AttackUid)
 			if err != nil {
@@ -91,7 +91,7 @@ func (networkAttack) Attack(options core.AttackConfig, env Environment) (err err
 }
 
 func (s *Server) applyIPSet(attack *core.NetworkCommand, uid string) (string, error) {
-	ipset, err := attack.ToIPSet(fmt.Sprintf("chaos-%s", uid[:16]))
+	ipset, err := attack.ToIPSet(fmt.Sprintf("chaos-%.16s", uid))
 	if err != nil {
 		return "", errors.WithStack(err)
 	}
@@ -192,6 +192,14 @@ func (s *Server) applyTC(attack *core.NetworkCommand, ipset string, uid string)
 			Duplicate:   attack.Percent,
 			Correlation: attack.Correlation,
 		}
+	case core.NetworkBandwidthAction:
+		tc.Bandwidth = &core.BandwidthSpec{
+			Rate:     attack.Rate,
+			Limit:    attack.Limit,
+			Buffer:   attack.Buffer,
+			Peakrate: attack.Peakrate,
+			Minburst: attack.Minburst,
+		}
 	default:
 		return errors.Errorf("network %s attack not supported", attack.Action)
 	}
@@ -320,9 +328,9 @@ func (networkAttack) Recover(exp core.Experiment, env Environment) error {
 			}
 		}
 		return env.Chaos.recoverDNSServer(attack)
-	case core.NetworkPortOccupied:
+	case core.NetworkPortOccupiedAction:
 		return env.Chaos.recoverPortOccupied(attack, env.AttackUid)
-	case core.NetworkDelayAction, core.NetworkLossAction, core.NetworkCorruptAction, core.NetworkDuplicateAction, core.NetworkPartitionAction:
+	case core.NetworkDelayAction, core.NetworkLossAction, core.NetworkCorruptAction, core.NetworkDuplicateAction, core.NetworkPartitionAction, core.NetworkBandwidthAction:
 		if attack.NeedApplyIPSet() {
 			if err := env.Chaos.recoverIPSet(env.AttackUid); err != nil {
 				return errors.WithStack(err)