loadr/main.go

package main

import (
	"flag"
	"fmt"
	"loadr/lib"
	"math"
	"os"
)

// Version number of the loadr tool
var version string = "1.0.2"

// parseCommandLine processes command line arguments and returns configuration parameters
// Returns:
//   - requestsPerSecond: target rate of requests
//   - maxRequests: total number of requests to send
//   - url: target endpoint
//   - patterns: array of request patterns
//   - jsonFilePath: path to JSON file containing request body
//   - bearerToken: authorization token
func parseCommandLine() (float64, int, string, []lib.RequestPattern, string, string) {
	requestsPerSecond := flag.Float64("rate", 10, "Number of requests per second")
	maxRequests := flag.Int("max", 50, "Maximum number of requests to send (0 for unlimited)")
	url := flag.String("url", "https://example.com", "The URL to make requests to")
	pattern := flag.String("pattern", "", `Request pattern (e.g., "5p" for 5 POSTs, "1p5g" for sequential, "20%p80%g" for probabilistic)`)
	jsonFilePath := flag.String("json", "", "Path to the JSON file with request data")
	bearerToken := flag.String("token", "", "Bearer token for authorization")
	versionFlag := flag.Bool("version", false, "Print the version and exit")
	versionFlagShort := flag.Bool("v", false, "Print the version and exit")

	flag.Parse()

	if *versionFlag || *versionFlagShort {
		fmt.Println("Version:", version)
		os.Exit(0)
	}

	patterns, err := parsePattern(*pattern)
	if err != nil {
		fmt.Printf("Warning: %v. Using default GET pattern.\n", err)
		patterns = []lib.RequestPattern{{Verb: "GET", Sequence: 1}}
	}

	return *requestsPerSecond, *maxRequests, *url, patterns, *jsonFilePath, *bearerToken
}

// parsePattern interprets the pattern string and converts it to RequestPattern structs
// Pattern formats:
//   - Sequential: "5p" (5 POSTs), "1p5g" (1 POST then 5 GETs)
//   - Probabilistic: "20%p80%g" (20% POSTs, 80% GETs)
//
// Returns error if pattern is invalid or percentages don't sum to 100
func parsePattern(pattern string) ([]lib.RequestPattern, error) {
	if pattern == "" {
		return []lib.RequestPattern{{Verb: "GET", Sequence: 1}}, nil
	}

	var patterns []lib.RequestPattern
	var current int
	isProbabilistic := false

	// Parse pattern string character by character
	for i := 0; i < len(pattern); i++ {
		c := pattern[i]
		switch {
		case c >= '0' && c <= '9':
			current = current*10 + int(c-'0')
		case c == '%':
			isProbabilistic = true
		case c == 'p' || c == 'P' || c == 'g' || c == 'G':
			verb := "GET"
			if c == 'p' || c == 'P' {
				verb = "POST"
			}

			if isProbabilistic {
				patterns = append(patterns, lib.RequestPattern{
					Verb:       verb,
					Percentage: float64(current),
				})
			} else {
				if current == 0 {
					current = 1 // Default to 1 if no number specified
				}
				patterns = append(patterns, lib.RequestPattern{
					Verb:     verb,
					Sequence: current,
				})
			}
			current = 0
		default:
			return nil, fmt.Errorf("invalid pattern character: %c", c)
		}
	}

	if len(patterns) == 0 {
		return nil, fmt.Errorf("no valid patterns found in: %s", pattern)
	}

	// For probabilistic patterns, ensure percentages sum to 100
	if isProbabilistic {
		total := 0.0
		for _, p := range patterns {
			total += p.Percentage
		}
		if math.Abs(total-100.0) > 0.001 {
			return nil, fmt.Errorf("percentages must sum to 100, got: %.1f", total)
		}
	}

	return patterns, nil
}

// readJSONFile loads and returns the contents of a JSON file if specified
// Returns nil if filePath is empty
func readJSONFile(filePath string) ([]byte, error) {
	if filePath == "" {
		return nil, nil
	}
	return os.ReadFile(filePath)
}

func main() {
	requestsPerSecond, maxRequests, url, patterns, jsonFilePath, bearerToken := parseCommandLine()

	if maxRequests <= 0 {
		fmt.Println("Error: max must be an integer greater than 0")
		return
	}

	jsonData, err := readJSONFile(jsonFilePath)
	if err != nil {
		fmt.Println("Error reading JSON file:", err)
		return
	}

	lib.SendRequests(url, patterns, maxRequests, requestsPerSecond, bearerToken, jsonData)
}