config.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"
)

type Config struct {
	requests         int
	concurrency      int
	timelimit        int
	executionTimeout time.Duration

	method              string
	bodyContent         []byte
	contentType         string
	headers             []string
	cookies             []string
	gzip                bool
	keepAlive           bool
	basicAuthentication string
	userAgent           string

	url  string
	host string
	port int
}

func LoadConfig() (config *Config, err error) {
	// setup command-line flags
	flag.IntVar(&Verbosity, "v", 0, "How much troubleshooting info to print")
	flag.IntVar(&GoMaxProcs, "G", runtime.NumCPU(), "Number of CPU")
	flag.BoolVar(&ContinueOnError, "r", false, "Don't exit when errors")

	request := flag.Int("n", 1, "Number of requests to perform")
	concurrency := flag.Int("c", 1, "Number of multiple requests to make")
	timelimit := flag.Int("t", 0, "Seconds to max. wait for responses")

	postFile := flag.String("p", "", "File containing data to POST. Remember also to set -T")
	putFile := flag.String("u", "", "File containing data to PUT. Remember also to set -T")
	headMethod := flag.Bool("i", false, "Use HEAD instead of GET")
	contentType := flag.String("T", "text/plain", "Content-type header for POSTing, eg. 'application/x-www-form-urlencoded' Default is 'text/plain'")

	var headers, cookies stringSet
	flag.Var(&headers, "H", "Add Arbitrary header line, eg. 'Accept-Encoding: gzip' Inserted after all normal header lines. (repeatable)")
	flag.Var(&cookies, "C", "Add cookie, eg. 'Apache=1234. (repeatable)")

	basicAuthentication := flag.String("A", "", "Add Basic WWW Authentication, the attributes are a colon separated username and password.")
	keepAlive := flag.Bool("k", false, "Use HTTP KeepAlive feature")
	gzip := flag.Bool("z", false, "Use HTTP Gzip feature")

	showHelp := flag.Bool("h", false, "Display usage information (this message)")

	flag.Usage = func() {
		fmt.Print("Usage: gb [options] http[s]://hostname[:port]/path\nOptions are:\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	if *showHelp {
		flag.Usage()
		os.Exit(0)
	}

	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(-1)
	}

	urlStr := strings.Trim(strings.Join(flag.Args(), ""), " ")
	isURL, _ := regexp.MatchString(`http.*?://.*`, urlStr)

	if !isURL {
		flag.Usage()
		os.Exit(-1)
	}

	// build configuration
	config = &Config{}
	config.requests = *request
	config.concurrency = *concurrency

	switch {
	case *postFile != "":
		config.method = "POST"
		if err = loadFile(config, *postFile); err != nil {
			return
		}
	case *putFile != "":
		config.method = "PUT"
		if err = loadFile(config, *putFile); err != nil {
			return
		}
	case *headMethod:
		config.method = "HEAD"
	default:
		config.method = "GET"
	}

	if *timelimit > 0 {
		config.timelimit = *timelimit
		if config.requests == 1 {
			config.requests = MaxRequests
		}
	}
	config.executionTimeout = MaxExecutionTimeout

	config.contentType = *contentType
	config.keepAlive = *keepAlive
	config.gzip = *gzip
	config.basicAuthentication = *basicAuthentication
	config.headers = []string(headers)
	config.cookies = []string(cookies)
	config.userAgent = "GoHttpBench/" + GBVersion

	URL, err := url.Parse(urlStr)
	if err != nil {
		return
	}
	config.host, config.port = extractHostAndPort(URL)
	config.url = urlStr

	if Verbosity > 1 {
		fmt.Printf("dump config: %#+v\n", config)
	}

	// validate configuration
	if config.requests < 1 || config.concurrency < 1 || config.timelimit < 0 || GoMaxProcs < 1 || Verbosity < 0 {
		err = errors.New("wrong number of arguments")
		return
	}

	if config.concurrency > config.requests {
		err = errors.New("Cannot use concurrency level greater than total number of requests")
		return
	}

	return

}

func loadFile(config *Config, filename string) error {
	bytes, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}
	config.bodyContent = bytes
	return nil
}

type stringSet []string

func (f *stringSet) String() string {
	return fmt.Sprint([]string(*f))
}

func (f *stringSet) Set(value string) error {
	*f = append(*f, value)
	return nil
}

func extractHostAndPort(url *url.URL) (host string, port int) {

	hostname := url.Host
	pos := strings.LastIndex(hostname, ":")
	if pos > 0 {
		portInt64, _ := strconv.Atoi(hostname[pos+1:])
		host = hostname[0:pos]
		port = int(portInt64)
	} else {
		host = hostname
		if url.Scheme == "http" {
			port = 80
		} else if url.Scheme == "https" {
			port = 443
		} else {
			panic("unsupported protocol schema:" + url.Scheme)
		}
	}

	return
}