Simulate port is already allocated (#62) (#67)

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

Author: luanqi521-s

cmd/attack/network.go

@@ -45,6 +45,7 @@ func NewNetworkAttackCommand() *cobra.Command {
 		NewNetworkCorruptCommand(dep, options),
 		NetworkDuplicateCommand(dep, options),
 		NetworkDNSCommand(dep, options),
+		NewNetworkPortOccupiedCommand(dep, options),
 	)
 
 	return cmd
@@ -201,3 +202,19 @@ func commonNetworkAttackFunc(options *core.NetworkCommand, chaos *chaosd.Server)
 
 	utils.NormalExit(fmt.Sprintf("Attack network successfully, uid: %s", uid))
 }
+
+func NewNetworkPortOccupiedCommand(dep fx.Option, options *core.NetworkCommand) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "port",
+		Short: "attack network port",
+
+		Run: func(cmd *cobra.Command, args []string) {
+			options.Action = core.NetworkPortOccupied
+			options.CompleteDefaults()
+			utils.FxNewAppWithoutLog(dep, fx.Invoke(commonNetworkAttackFunc)).Run()
+		},
+	}
+
+	cmd.Flags().StringVarP(&options.Port, "port", "p", "", "this specified port is to occupied")
+	return cmd
+}

cmd/main

@@ -43,6 +43,8 @@ type NetworkCommand struct {
 
 	// used for DNS attack
 	DNSServer string
+	Port      string
+	PortPid   int32
 	DNSIp     string
 	DNSHost   string
 }
@@ -55,6 +57,7 @@ const (
 	NetworkCorruptAction   = "corrupt"
 	NetworkDuplicateAction = "duplicate"
 	NetworkDNSAction       = "dns"
+	NetworkPortOccupied    = "occupied"
 )
 
 func (n *NetworkCommand) Validate() error {
@@ -68,6 +71,8 @@ func (n *NetworkCommand) Validate() error {
 		return n.validNetworkCommon()
 	case NetworkDNSAction:
 		return n.validNetworkDNS()
+	case NetworkPortOccupied:
+		return n.validNetworkOccupied()
 	default:
 		return errors.Errorf("network action %s not supported", n.Action)
 	}
@@ -143,6 +148,13 @@ func (n *NetworkCommand) validNetworkDNS() error {
 	return nil
 }
 
+func (n *NetworkCommand) validNetworkOccupied() error {
+	if len(n.Port) == 0 {
+		return errors.New("port is required")
+	}
+	return nil
+}
+
 func (n *NetworkCommand) CompleteDefaults() {
 	switch n.Action {
 	case NetworkDelayAction:

pkg/core/network.go

@@ -23,6 +23,10 @@ import (
 	"os/exec"
 	"regexp"
 	"strings"
+	"syscall"
+
+	"github.com/chaos-mesh/chaos-mesh/pkg/bpm"
+	"github.com/shirou/gopsutil/process"
 
 	"go.uber.org/zap"
 
@@ -59,6 +63,9 @@ func (networkAttack) Attack(options core.AttackConfig, env Environment) (err err
 			}
 		}
 
+	case core.NetworkPortOccupied:
+		return env.Chaos.applyPortOccupied(attack)
+
 	case core.NetworkDelayAction, core.NetworkLossAction, core.NetworkCorruptAction, core.NetworkDuplicateAction:
 		if attack.NeedApplyIPSet() {
 			ipsetName, err = env.Chaos.applyIPSet(attack, env.AttackUid)
@@ -316,7 +323,8 @@ func (networkAttack) Recover(exp core.Experiment, env Environment) error {
 			}
 		}
 		return env.Chaos.recoverDNSServer(attack)
-
+	case core.NetworkPortOccupied:
+		return env.Chaos.recoverPortOccupied(attack, env.AttackUid)
 	case core.NetworkDelayAction, core.NetworkLossAction, core.NetworkCorruptAction, core.NetworkDuplicateAction:
 		if attack.NeedApplyIPSet() {
 			if err := env.Chaos.recoverIPSet(env.AttackUid); err != nil {
@@ -411,6 +419,83 @@ func (s *Server) recoverDNSServer(attack *core.NetworkCommand) error {
 	return nil
 }
 
+func (s *Server) applyPortOccupied(attack *core.NetworkCommand) error {
+
+	if len(attack.Port) == 0 {
+		return nil
+	}
+
+	flag, err := checkPortIsListened(attack.Port)
+	if err != nil {
+		if flag {
+			return errors.Errorf("port %s has been occupied", attack.Port)
+		}
+		return errors.WithStack(err)
+	}
+
+	if flag {
+		return errors.Errorf("port %s has been occupied", attack.Port)
+	}
+
+	args := fmt.Sprintf("-p=%s", attack.Port)
+	cmd := bpm.DefaultProcessBuilder("PortOccupyTool", args).Build()
+
+	cmd.Cmd.SysProcAttr = &syscall.SysProcAttr{}
+
+	backgroundProcessManager := bpm.NewBackgroundProcessManager()
+	err = backgroundProcessManager.StartProcess(cmd)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+
+	attack.PortPid = int32(cmd.Process.Pid)
+
+	return nil
+}
+
+func checkPortIsListened(port string) (bool, error) {
+	checkStatement := fmt.Sprintf("lsof -i:%s | awk '{print $2}' | grep -v PID", port)
+	cmd := exec.Command("sh", "-c", checkStatement)
+
+	stdout, err := cmd.CombinedOutput()
+	if err != nil {
+		if err.Error() == "exit status 1" && string(stdout) == "" {
+			return false, nil
+		}
+		log.Error(cmd.String()+string(stdout), zap.Error(err))
+		return true, errors.WithStack(err)
+	}
+
+	if string(stdout) == "" {
+		return false, nil
+	}
+	return true, nil
+}
+
+func (s *Server) recoverPortOccupied(attack *core.NetworkCommand, uid string) error {
+
+	proc, err := process.NewProcess(attack.PortPid)
+	if err != nil {
+		return err
+	}
+
+	procName, err := proc.Name()
+	if err != nil {
+		return err
+	}
+
+	if !strings.Contains(procName, "PortOccupyTool") {
+		log.Warn("the process is not PortOccupyTool, maybe it is killed by manual")
+		return nil
+	}
+
+	if err := proc.Kill(); err != nil {
+		log.Error("the port occupy process kill failed", zap.Error(err))
+		return err
+	}
+	return nil
+}
+
 func (s *Server) recoverEtcHosts(attack *core.NetworkCommand, uid string) error {
 	cmd := "mv /etc/hosts.chaosd." + uid + " /etc/hosts"
 	recoverCmd := exec.Command("/bin/bash", "-c", cmd) // #nosec

pkg/server/chaosd/network.go

@@ -0,0 +1,65 @@
+// Copyright 2020 Chaos Mesh Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/log"
+	"github.com/spf13/cobra"
+	"go.uber.org/zap"
+)
+
+func PortOccupyTool() error {
+	var Port string
+	var rootCmd = &cobra.Command{
+		Use:   "http server start",
+		Short: "http server start",
+		Run: func(cmd *cobra.Command, args []string) {
+			startHttp(Port)
+		},
+	}
+
+	rootCmd.Flags().StringVarP(&Port, "port", "p", "", "port to occupy")
+	if err := rootCmd.Execute(); err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}
+
+func generateHttpPort(port string) string {
+	s := fmt.Sprintf(":%s", port)
+	return s
+}
+
+func startHttp(porttoOccupy string) {
+
+	s := generateHttpPort(porttoOccupy)
+
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(w, "Hello cmd!")
+	})
+	if err := http.ListenAndServe(s, nil); err != nil {
+		log.Error("ListenAndServe", zap.Error(err))
+	}
+}
+
+func main() {
+	if err := PortOccupyTool(); err != nil {
+		log.Error("PortOccupyTool run fail", zap.Error(err))
+	}
+}