diff --git a/.gitignore b/.gitignore index e34d93e..d8f5673 100644 --- a/.gitignore +++ b/.gitignore @@ -26,6 +26,7 @@ _testmain.go # tools and virtualenv .idea env/ +.vscode/ # test coverage reports overalls.coverprofile diff --git a/cmd/tcpgoon.go b/cmd/tcpgoon.go index c73238f..4a25286 100644 --- a/cmd/tcpgoon.go +++ b/cmd/tcpgoon.go @@ -11,6 +11,8 @@ import ( "github.com/dachad/tcpgoon/debugging" "github.com/dachad/tcpgoon/mtcpclient" "github.com/dachad/tcpgoon/tcpclient" + "github.com/dachad/tcpgoon/docker" + "github.com/spf13/cobra" ) @@ -24,6 +26,8 @@ type tcpgoonParams struct { debug bool reportingInterval int assumeyes bool + docker bool + containerID string } var params tcpgoonParams @@ -33,12 +37,12 @@ var runCmd = &cobra.Command{ Short: "Run tcpgoon test", Long: ``, PreRun: func(cmd *cobra.Command, args []string) { + enableDebuggingIfFlagSet(params) if err := validateRequiredArgs(¶ms, args); err != nil { cmd.Println(err) cmd.Println(cmd.UsageString()) os.Exit(1) } - enableDebuggingIfFlagSet(params) autorunValidation(params) }, Run: func(cmd *cobra.Command, args []string) { @@ -53,26 +57,65 @@ func init() { runCmd.Flags().BoolVarP(¶ms.debug, "debug", "d", false, "Print debugging information to the standard error") runCmd.Flags().IntVarP(¶ms.reportingInterval, "interval", "i", 1, "Interval, in seconds, between stats updates") runCmd.Flags().BoolVarP(¶ms.assumeyes, "assume-yes", "y", false, "Force execution without asking for confirmation") + runCmd.Flags().BoolVarP(¶ms.docker, "docker", "", false, "Running against a docker image. " + + "Use the image reference instead of the , and a custom mapped port if the first exposed one is not ok (optional") } func validateRequiredArgs(params *tcpgoonParams, args []string) error { - if len(args) != 2 { - return errors.New("Number of required parameters doesn't match") + var minNumberRequiredArgs int + var maxNumberRequiredArgs int + + if params.docker { + minNumberRequiredArgs = 1 + maxNumberRequiredArgs = 2 + } else { + minNumberRequiredArgs = 2 + maxNumberRequiredArgs = 2 } - params.target = args[0] - addrs, err := net.LookupIP(params.target) - if err != nil || len(addrs) == 0 { - return errors.New("Domain name not resolvable") + + if len(args) < minNumberRequiredArgs || len(args) > maxNumberRequiredArgs{ + return errors.New("Number of required parameters doesn't match") } - params.targetip = addrs[0].String() - fmt.Fprintln(debugging.DebugOut, "TCPGOON target: Hostname(", params.target, "), IP (", params.targetip, ")") + if params.docker { + dockerImage := args[0] + + fmt.Println("Starting Docker test for", dockerImage) + + var port int + if len(args) == 2 { + + portInt, err := strconv.Atoi(args[1]) + port = portInt - port, err := strconv.Atoi(args[1]) - if err != nil && port <= 0 { - return errors.New("Port argument is not a valid integer") + if err != nil && port <= 0 { + return errors.New("Port argument is not a valid integer") + } + } else { + port = 0 + } + + target, containerID := docker.DownloadAndRun(dockerImage, port) + params.containerID = containerID + params.target = target.IP + params.port = target.Port + + } else { + params.target = args[0] + addrs, err := net.LookupIP(params.target) + if err != nil || len(addrs) == 0 { + return errors.New("Domain name not resolvable") + } + + params.targetip = addrs[0].String() + fmt.Fprintln(debugging.DebugOut, "TCPGOON target: Hostname(", params.target, "), IP (", params.targetip, ")") + + port, err := strconv.Atoi(args[1]) + if err != nil && port <= 0 { + return errors.New("Port argument is not a valid integer") + } + params.port = port } - params.port = port return nil } @@ -86,11 +129,15 @@ func enableDebuggingIfFlagSet(params tcpgoonParams) { func autorunValidation(params tcpgoonParams) { if !(params.assumeyes || cmdutil.AskForUserConfirmation(params.target, params.port, params.numberConnections)) { fmt.Fprintln(debugging.DebugOut, "Execution not approved by the user") + if params.docker{ + docker.Stop(params.containerID) + } cmdutil.CloseAbruptly() } } func run(params tcpgoonParams) { + tcpclient.DefaultDialTimeoutInMs = params.connDialTimeout // TODO: we should decouple the caller from the mtcpclient package (too many structures being moved from @@ -98,6 +145,12 @@ func run(params tcpgoonParams) { // may help connStatusCh, connStatusTracker := mtcpclient.StartBackgroundReporting(params.numberConnections, params.reportingInterval) closureCh := mtcpclient.StartBackgroundClosureTrigger(*connStatusTracker) + if params.docker{ + go func() { + <-closureCh + docker.Stop(params.containerID) + }() + } mtcpclient.MultiTCPConnect(params.numberConnections, params.delay, params.target, params.port, connStatusCh, closureCh) fmt.Fprintln(debugging.DebugOut, "Tests execution completed") diff --git a/docker/docker.go b/docker/docker.go new file mode 100644 index 0000000..0099e8c --- /dev/null +++ b/docker/docker.go @@ -0,0 +1,163 @@ +package docker + +import ( + "context" + "fmt" + "io" + "math/rand" + "net" + + //"os" + "strconv" + + "github.com/docker/docker/api/types" + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/client" + "github.com/docker/go-connections/nat" + + "github.com/dachad/tcpgoon/debugging" +) + +type IPandPort struct { + IP string + Port int +} + +func lookForAvailableLocalPort() (availablePort int){ + minPort := 1025 + maxPort := 65535 + + for { + availablePort = rand.Intn(maxPort - minPort) + minPort + ln, err := net.Listen("tcp", ":" + strconv.Itoa(availablePort)) + ln.Close() + if err == nil { + break + } + } + return availablePort +} + +func DownloadAndRun(image string, port int) (targetinfo IPandPort, containerID string) { + ctx := context.Background() + cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation()) + if err != nil { + panic(err) + } + + fmt.Fprintln(debugging.DebugOut,"Pulling Docker image", image) + + out, err := cli.ImagePull(ctx, image, types.ImagePullOptions{}) + if err != nil { + panic(err) + } + io.Copy(debugging.DebugOut, out) + + defer out.Close() + + fmt.Fprintln(debugging.DebugOut, "Docker image", image, "downloaded") + + mappedPort := "" + mapperProto := "" + if port == 0 { + // if not specific mapping defined we look for exposed ports + // getting the exposed port before decide the mapping + + // TODO support mapping volumes, some containers required them + respCreate, err := cli.ContainerCreate(ctx, &container.Config{ + Image: image, + }, nil, nil, "") + if err != nil { + panic(err) + } + + if err := cli.ContainerStart(ctx, respCreate.ID, types.ContainerStartOptions{}); err != nil { + panic(err) + } + + respInspect := inspect(respCreate.ID) + natPorts := respInspect.NetworkSettings.NetworkSettingsBase.Ports + + for k, _ := range natPorts { + mappedPort = k.Port() + mapperProto = k.Proto() + break + } + // Stopping container used to get the port mapping + Stop(respCreate.ID) + } else { + mappedPort = strconv.Itoa(port) + mapperProto = "tcp" + } + + if mappedPort == "" || mapperProto != "tcp" { + panic("Not found any TCP port mapping") + } else { + fmt.Println("Internal Docker port binding:", mapperProto, mappedPort) + } + + // hack https://stackoverflow.com/questions/47395973/issue-while-using-docker-api-for-go-cannot-import-nat + + hostBinding := nat.PortBinding{ + HostIP: "0.0.0.0", + HostPort: strconv.Itoa(lookForAvailableLocalPort()), + } + containerPort, err := nat.NewPort(mapperProto, mappedPort) + if err != nil { + panic("Unable to get the port") + } + + portBinding := nat.PortMap{containerPort: []nat.PortBinding{hostBinding}} + resp, err := cli.ContainerCreate( + context.Background(), + &container.Config{ + Image: image, + }, + &container.HostConfig{ + PortBindings: portBinding, + }, nil, "") + if err != nil { + panic(err) + } + fmt.Fprintln(debugging.DebugOut, "Docker container built") + + if err := cli.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil { + panic(err) + } + fmt.Fprintln(debugging.DebugOut, "Started docker container", resp.ID) + + targetinfo.IP = "127.0.0.1" + targetinfo.Port, _ = strconv.Atoi(hostBinding.HostPort) + + return targetinfo, resp.ID +} + +func Stop(containerID string) { + ctx := context.Background() + cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation()) + if err != nil { + panic(err) + } + + fmt.Fprintln(debugging.DebugOut, "Stopping container", containerID, "...") + if err := cli.ContainerStop(ctx, containerID, nil); err != nil { + panic(err) + } + fmt.Fprintln(debugging.DebugOut, "Container stopped successfully") +} + +func inspect(containerID string) types.ContainerJSON { + ctx := context.Background() + cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation()) + if err != nil { + panic(err) + } + + fmt.Fprintln(debugging.DebugOut, "Getting container info", containerID, "...") + resp, err := cli.ContainerInspect(ctx, containerID) + if err != nil { + panic(err) + } + + return resp +}