示例#1
0
文件: attack.go 项目: hpbuniat/vegeta
// attack validates the attack arguments, sets up the
// required resources, launches the attack and writes the results
func attack(rate uint64, duration time.Duration, targetsf, ordering,
	output string, redirects int, timeout time.Duration, unique bool,
	hdr http.Header) error {

	if rate == 0 {
		return fmt.Errorf(errRatePrefix + "can't be zero")
	}

	if duration == 0 {
		return fmt.Errorf(errDurationPrefix + "can't be zero")
	}

	in, err := file(targetsf, false)
	if err != nil {
		return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err)
	}
	defer in.Close()
	targets, err := vegeta.NewTargetsFrom(in, unique)
	if err != nil {
		return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err)
	}
	targets.SetHeader(hdr)

	switch ordering {
	case "random":
		targets.Shuffle(time.Now().UnixNano())
	case "sequential":
		break
	default:
		return fmt.Errorf(errOrderingPrefix+"`%s` is invalid", ordering)
	}

	out, err := file(output, true)
	if err != nil {
		return fmt.Errorf(errOutputFilePrefix+"(%s): %s", output, err)
	}
	defer out.Close()

	vegeta.DefaultAttacker.SetRedirects(redirects)

	if timeout > 0 {
		vegeta.DefaultAttacker.SetTimeout(timeout)
	}

	log.Printf("Vegeta is attacking %d targets in %s order for %s...\n",
		len(targets), ordering, duration)
	results := vegeta.Attack(targets, rate, duration)
	log.Println("Done!")
	log.Printf("Writing results to '%s'...", output)
	if err := results.Encode(out); err != nil {
		return err
	}
	return nil
}
示例#2
0
// attack validates the attack arguments, sets up the
// required resources, launches the attack and writes the results
func attack(opts *attackOpts) error {
	if opts.rate == 0 {
		return fmt.Errorf(errRatePrefix + "can't be zero")
	}

	if opts.duration == 0 {
		return fmt.Errorf(errDurationPrefix + "can't be zero")
	}

	in, err := file(opts.targetsf, false)
	if err != nil {
		return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", opts.targetsf, err)
	}
	defer in.Close()

	var body []byte
	if opts.bodyf != "" {
		bodyr, err := file(opts.bodyf, false)
		if err != nil {
			return fmt.Errorf(errBodyFilePrefix+"(%s): %s", opts.bodyf, err)
		}
		defer bodyr.Close()

		if body, err = ioutil.ReadAll(bodyr); err != nil {
			return fmt.Errorf(errBodyFilePrefix+"(%s): %s", opts.bodyf, err)
		}
	}

	targets, err := vegeta.NewTargetsFrom(in, body, opts.headers.Header)
	if err != nil {
		return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", opts.targetsf, err)
	}

	switch opts.ordering {
	case "random":
		targets.Shuffle(time.Now().UnixNano())
	case "sequential":
		break
	default:
		return fmt.Errorf(errOrderingPrefix+"`%s` is invalid", opts.ordering)
	}

	out, err := file(opts.outputf, true)
	if err != nil {
		return fmt.Errorf(errOutputFilePrefix+"(%s): %s", opts.outputf, err)
	}
	defer out.Close()

	atk := vegeta.NewAttacker(opts.redirects, opts.timeout, *opts.laddr.IPAddr)

	log.Printf(
		"Vegeta is attacking %d targets in %s order for %s...\n",
		len(targets),
		opts.ordering,
		opts.duration,
	)
	results := atk.Attack(targets, opts.rate, opts.duration)

	log.Printf("Done! Writing results to '%s'...", opts.outputf)
	return results.Encode(out)
}