mirror of
https://github.com/aykhans/sarin.git
synced 2026-02-28 14:59:14 +00:00
v1.0.0: here we go again
This commit is contained in:
336
internal/sarin/request.go
Normal file
336
internal/sarin/request.go
Normal file
@@ -0,0 +1,336 @@
|
||||
package sarin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"maps"
|
||||
"math/rand/v2"
|
||||
"net/url"
|
||||
"strings"
|
||||
"text/template"
|
||||
|
||||
"github.com/joho/godotenv"
|
||||
"github.com/valyala/fasthttp"
|
||||
"go.aykhans.me/sarin/internal/types"
|
||||
utilsSlice "go.aykhans.me/utils/slice"
|
||||
)
|
||||
|
||||
type RequestGenerator func(*fasthttp.Request) error
|
||||
|
||||
type RequestGeneratorWithData func(*fasthttp.Request, any) error
|
||||
|
||||
type valuesData struct {
|
||||
Values map[string]string
|
||||
}
|
||||
|
||||
// NewRequestGenerator creates a new RequestGenerator function that generates HTTP requests
|
||||
// with the specified configuration. The returned RequestGenerator is NOT safe for concurrent
|
||||
// use by multiple goroutines.
|
||||
func NewRequestGenerator(
|
||||
methods []string,
|
||||
requestURL *url.URL,
|
||||
params types.Params,
|
||||
headers types.Headers,
|
||||
cookies types.Cookies,
|
||||
bodies []string,
|
||||
values []string,
|
||||
) (RequestGenerator, bool) {
|
||||
randSource := NewDefaultRandSource()
|
||||
//nolint:gosec // G404: Using non-cryptographic rand for load testing, not security
|
||||
localRand := rand.New(randSource)
|
||||
templateFuncMap := NewDefaultTemplateFuncMap(randSource)
|
||||
|
||||
methodGenerator, isMethodGeneratorDynamic := NewMethodGeneratorFunc(localRand, methods, templateFuncMap)
|
||||
paramsGenerator, isParamsGeneratorDynamic := NewParamsGeneratorFunc(localRand, params, templateFuncMap)
|
||||
headersGenerator, isHeadersGeneratorDynamic := NewHeadersGeneratorFunc(localRand, headers, templateFuncMap)
|
||||
cookiesGenerator, isCookiesGeneratorDynamic := NewCookiesGeneratorFunc(localRand, cookies, templateFuncMap)
|
||||
|
||||
bodyTemplateFuncMapData := &BodyTemplateFuncMapData{}
|
||||
bodyTemplateFuncMap := NewDefaultBodyTemplateFuncMap(randSource, bodyTemplateFuncMapData)
|
||||
bodyGenerator, isBodyGeneratorDynamic := NewBodyGeneratorFunc(localRand, bodies, bodyTemplateFuncMap)
|
||||
|
||||
valuesGenerator := NewValuesGeneratorFunc(values, templateFuncMap)
|
||||
|
||||
return func(req *fasthttp.Request) error {
|
||||
req.SetRequestURI(requestURL.Path)
|
||||
req.Header.SetHost(requestURL.Host)
|
||||
|
||||
data, err := valuesGenerator()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := methodGenerator(req, data); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
bodyTemplateFuncMapData.ClearFormDataContenType()
|
||||
if err := bodyGenerator(req, data); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := headersGenerator(req, data); err != nil {
|
||||
return err
|
||||
}
|
||||
if bodyTemplateFuncMapData.GetFormDataContenType() != "" {
|
||||
req.Header.Add("Content-Type", bodyTemplateFuncMapData.GetFormDataContenType())
|
||||
}
|
||||
|
||||
if err := paramsGenerator(req, data); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := cookiesGenerator(req, data); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if requestURL.Scheme == "https" {
|
||||
req.URI().SetScheme("https")
|
||||
}
|
||||
return nil
|
||||
}, isMethodGeneratorDynamic ||
|
||||
isParamsGeneratorDynamic ||
|
||||
isHeadersGeneratorDynamic ||
|
||||
isCookiesGeneratorDynamic ||
|
||||
isBodyGeneratorDynamic
|
||||
}
|
||||
|
||||
func NewMethodGeneratorFunc(localRand *rand.Rand, methods []string, templateFunctions template.FuncMap) (RequestGeneratorWithData, bool) {
|
||||
methodGenerator, isDynamic := buildStringSliceGenerator(localRand, methods, templateFunctions)
|
||||
|
||||
var (
|
||||
method string
|
||||
err error
|
||||
)
|
||||
return func(req *fasthttp.Request, data any) error {
|
||||
method, err = methodGenerator()(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
req.Header.SetMethod(method)
|
||||
return nil
|
||||
}, isDynamic
|
||||
}
|
||||
|
||||
func NewBodyGeneratorFunc(localRand *rand.Rand, bodies []string, templateFunctions template.FuncMap) (RequestGeneratorWithData, bool) {
|
||||
bodyGenerator, isDynamic := buildStringSliceGenerator(localRand, bodies, templateFunctions)
|
||||
|
||||
var (
|
||||
body string
|
||||
err error
|
||||
)
|
||||
return func(req *fasthttp.Request, data any) error {
|
||||
body, err = bodyGenerator()(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
req.SetBody([]byte(body))
|
||||
return nil
|
||||
}, isDynamic
|
||||
}
|
||||
|
||||
func NewParamsGeneratorFunc(localRand *rand.Rand, params types.Params, templateFunctions template.FuncMap) (RequestGeneratorWithData, bool) {
|
||||
generators, isDynamic := buildKeyValueGenerators(localRand, params, templateFunctions)
|
||||
|
||||
var (
|
||||
key, value string
|
||||
err error
|
||||
)
|
||||
return func(req *fasthttp.Request, data any) error {
|
||||
for _, gen := range generators {
|
||||
key, err = gen.Key(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
value, err = gen.Value()(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
req.URI().QueryArgs().Add(key, value)
|
||||
}
|
||||
return nil
|
||||
}, isDynamic
|
||||
}
|
||||
|
||||
func NewHeadersGeneratorFunc(localRand *rand.Rand, headers types.Headers, templateFunctions template.FuncMap) (RequestGeneratorWithData, bool) {
|
||||
generators, isDynamic := buildKeyValueGenerators(localRand, headers, templateFunctions)
|
||||
|
||||
var (
|
||||
key, value string
|
||||
err error
|
||||
)
|
||||
return func(req *fasthttp.Request, data any) error {
|
||||
for _, gen := range generators {
|
||||
key, err = gen.Key(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
value, err = gen.Value()(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
return nil
|
||||
}, isDynamic
|
||||
}
|
||||
|
||||
func NewCookiesGeneratorFunc(localRand *rand.Rand, cookies types.Cookies, templateFunctions template.FuncMap) (RequestGeneratorWithData, bool) {
|
||||
generators, isDynamic := buildKeyValueGenerators(localRand, cookies, templateFunctions)
|
||||
|
||||
var (
|
||||
key, value string
|
||||
err error
|
||||
)
|
||||
if len(generators) > 0 {
|
||||
return func(req *fasthttp.Request, data any) error {
|
||||
cookieStrings := make([]string, 0, len(generators))
|
||||
for _, gen := range generators {
|
||||
key, err = gen.Key(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
value, err = gen.Value()(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
cookieStrings = append(cookieStrings, key+"="+value)
|
||||
}
|
||||
req.Header.Add("Cookie", strings.Join(cookieStrings, "; "))
|
||||
return nil
|
||||
}, isDynamic
|
||||
}
|
||||
|
||||
return func(req *fasthttp.Request, data any) error {
|
||||
return nil
|
||||
}, isDynamic
|
||||
}
|
||||
|
||||
func NewValuesGeneratorFunc(values []string, templateFunctions template.FuncMap) func() (valuesData, error) {
|
||||
generators := make([]func(_ any) (string, error), len(values))
|
||||
|
||||
for i, v := range values {
|
||||
generators[i], _ = createTemplateFunc(v, templateFunctions)
|
||||
}
|
||||
|
||||
var (
|
||||
rendered string
|
||||
data map[string]string
|
||||
err error
|
||||
)
|
||||
return func() (valuesData, error) {
|
||||
result := make(map[string]string)
|
||||
for _, generator := range generators {
|
||||
rendered, err = generator(nil)
|
||||
if err != nil {
|
||||
return valuesData{}, fmt.Errorf("values rendering: %w", err)
|
||||
}
|
||||
|
||||
data, err = godotenv.Unmarshal(rendered)
|
||||
if err != nil {
|
||||
return valuesData{}, fmt.Errorf("values rendering: %w", err)
|
||||
}
|
||||
|
||||
maps.Copy(result, data)
|
||||
}
|
||||
|
||||
return valuesData{Values: result}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func createTemplateFunc(value string, templateFunctions template.FuncMap) (func(data any) (string, error), bool) {
|
||||
tmpl, err := template.New("").Funcs(templateFunctions).Parse(value)
|
||||
if err == nil && hasTemplateActions(tmpl) {
|
||||
var err error
|
||||
return func(data any) (string, error) {
|
||||
var buf bytes.Buffer
|
||||
if err = tmpl.Execute(&buf, data); err != nil {
|
||||
return "", fmt.Errorf("template rendering: %w", err)
|
||||
}
|
||||
return buf.String(), nil
|
||||
}, true
|
||||
}
|
||||
return func(_ any) (string, error) { return value, nil }, false
|
||||
}
|
||||
|
||||
type keyValueGenerator struct {
|
||||
Key func(data any) (string, error)
|
||||
Value func() func(data any) (string, error)
|
||||
}
|
||||
|
||||
type keyValueItem interface {
|
||||
types.Param | types.Header | types.Cookie
|
||||
}
|
||||
|
||||
func buildKeyValueGenerators[T keyValueItem](
|
||||
localRand *rand.Rand,
|
||||
items []T,
|
||||
templateFunctions template.FuncMap,
|
||||
) ([]keyValueGenerator, bool) {
|
||||
isDynamic := false
|
||||
generators := make([]keyValueGenerator, len(items))
|
||||
|
||||
for generatorIndex, item := range items {
|
||||
// Convert to KeyValue to access fields
|
||||
keyValue := types.KeyValue[string, []string](item)
|
||||
|
||||
// Generate key function
|
||||
keyFunc, keyIsDynamic := createTemplateFunc(keyValue.Key, templateFunctions)
|
||||
if keyIsDynamic {
|
||||
isDynamic = true
|
||||
}
|
||||
|
||||
// Generate value functions
|
||||
valueFuncs := make([]func(data any) (string, error), len(keyValue.Value))
|
||||
for j, v := range keyValue.Value {
|
||||
valueFunc, valueIsDynamic := createTemplateFunc(v, templateFunctions)
|
||||
if valueIsDynamic {
|
||||
isDynamic = true
|
||||
}
|
||||
valueFuncs[j] = valueFunc
|
||||
}
|
||||
|
||||
generators[generatorIndex] = keyValueGenerator{
|
||||
Key: keyFunc,
|
||||
Value: utilsSlice.RandomCycle(localRand, valueFuncs...),
|
||||
}
|
||||
|
||||
if len(keyValue.Value) > 1 {
|
||||
isDynamic = true
|
||||
}
|
||||
}
|
||||
|
||||
return generators, isDynamic
|
||||
}
|
||||
|
||||
func buildStringSliceGenerator(
|
||||
localRand *rand.Rand,
|
||||
values []string,
|
||||
templateFunctions template.FuncMap,
|
||||
) (func() func(data any) (string, error), bool) {
|
||||
// Return a function that returns an empty string generator if values is empty
|
||||
if len(values) == 0 {
|
||||
emptyFunc := func(_ any) (string, error) { return "", nil }
|
||||
return func() func(_ any) (string, error) { return emptyFunc }, false
|
||||
}
|
||||
|
||||
isDynamic := len(values) > 1
|
||||
valueFuncs := make([]func(data any) (string, error), len(values))
|
||||
|
||||
for i, value := range values {
|
||||
valueFunc, valueIsDynamic := createTemplateFunc(value, templateFunctions)
|
||||
if valueIsDynamic {
|
||||
isDynamic = true
|
||||
}
|
||||
valueFuncs[i] = valueFunc
|
||||
}
|
||||
|
||||
return utilsSlice.RandomCycle(localRand, valueFuncs...), isDynamic
|
||||
}
|
||||
Reference in New Issue
Block a user