dodo/requests/request.go

package requests

import (
	"bytes"
	"context"
	"math/rand"
	"net/url"
	"strings"
	"text/template"
	"time"

	"github.com/aykhans/dodo/config"
	"github.com/aykhans/dodo/types"
	"github.com/aykhans/dodo/utils"
	"github.com/brianvoe/gofakeit/v7"
	"github.com/valyala/fasthttp"
)

type RequestGeneratorFunc func() *fasthttp.Request

// Request represents an HTTP request to be sent using the fasthttp client.
// It isn't thread-safe and should be used by a single goroutine.
type Request struct {
	getClient  ClientGeneratorFunc
	getRequest RequestGeneratorFunc
}

type keyValueGenerator struct {
	key   func() string
	value func() string
}

// Send sends the HTTP request using the fasthttp client with a specified timeout.
// It returns the HTTP response or an error if the request fails or times out.
func (r *Request) Send(ctx context.Context, timeout time.Duration) (*fasthttp.Response, error) {
	client := r.getClient()
	request := r.getRequest()
	defer fasthttp.ReleaseRequest(request)

	response := fasthttp.AcquireResponse()
	ch := make(chan error)
	go func() {
		err := client.DoTimeout(request, response, timeout)
		ch <- err
	}()
	select {
	case err := <-ch:
		if err != nil {
			fasthttp.ReleaseResponse(response)
			return nil, err
		}
		return response, nil
	case <-time.After(timeout):
		fasthttp.ReleaseResponse(response)
		return nil, types.ErrTimeout
	case <-ctx.Done():
		return nil, types.ErrInterrupt
	}
}

// newRequest creates a new Request instance based on the provided configuration and clients.
// It initializes a random number generator using the current time and a unique identifier (uid).
// Depending on the number of clients provided, it sets up a function to select the appropriate client.
// It also sets up a function to generate the request based on the provided configuration.
func newRequest(
	requestConfig config.RequestConfig,
	clients []*fasthttp.HostClient,
	uid int64,
) *Request {
	localRand := rand.New(rand.NewSource(time.Now().UnixNano() + uid))

	clientsCount := len(clients)
	if clientsCount < 1 {
		panic("no clients")
	}

	getClient := ClientGeneratorFunc(nil)
	if clientsCount == 1 {
		getClient = getSharedClientFuncSingle(clients[0])
	} else {
		getClient = getSharedClientFuncMultiple(clients, localRand)
	}

	getRequest := getRequestGeneratorFunc(
		requestConfig.URL,
		requestConfig.Params,
		requestConfig.Headers,
		requestConfig.Cookies,
		requestConfig.Method,
		requestConfig.Body,
		localRand,
	)

	requests := &Request{
		getClient:  getClient,
		getRequest: getRequest,
	}

	return requests
}

// getRequestGeneratorFunc returns a RequestGeneratorFunc which generates HTTP requests with the specified parameters.
// The function uses a local random number generator to select bodies, headers, cookies, and parameters if multiple options are provided.
func getRequestGeneratorFunc(
	URL url.URL,
	params types.Params,
	headers types.Headers,
	cookies types.Cookies,
	method string,
	bodies []string,
	localRand *rand.Rand,
) RequestGeneratorFunc {
	getParams := getKeyValueGeneratorFunc(params, localRand)
	getHeaders := getKeyValueGeneratorFunc(headers, localRand)
	getCookies := getKeyValueGeneratorFunc(cookies, localRand)
	getBody := getValueFunc(bodies, newFuncMap(localRand), localRand)

	return func() *fasthttp.Request {
		return newFasthttpRequest(
			URL,
			getParams(),
			getHeaders(),
			getCookies(),
			method,
			getBody(),
		)
	}
}

// newFasthttpRequest creates a new fasthttp.Request object with the provided parameters.
// It sets the request URI, host header, headers, cookies, params, method, and body.
func newFasthttpRequest(
	URL url.URL,
	params []types.KeyValue[string, string],
	headers []types.KeyValue[string, string],
	cookies []types.KeyValue[string, string],
	method string,
	body string,
) *fasthttp.Request {
	request := fasthttp.AcquireRequest()
	request.SetRequestURI(URL.Path)

	// Set the host of the request to the host header
	// If the host header is not set, the request will fail
	// If there is host header in the headers, it will be overwritten
	request.Header.SetHost(URL.Host)
	setRequestParams(request, params)
	setRequestHeaders(request, headers)
	setRequestCookies(request, cookies)
	setRequestMethod(request, method)
	setRequestBody(request, body)
	if URL.Scheme == "https" {
		request.URI().SetScheme("https")
	}

	return request
}

// setRequestParams adds the query parameters of the given request based on the provided key-value pairs.
func setRequestParams(req *fasthttp.Request, params []types.KeyValue[string, string]) {
	for _, param := range params {
		req.URI().QueryArgs().Add(param.Key, param.Value)
	}
}

// setRequestHeaders adds the headers of the given request with the provided key-value pairs.
func setRequestHeaders(req *fasthttp.Request, headers []types.KeyValue[string, string]) {
	for _, header := range headers {
		req.Header.Add(header.Key, header.Value)
	}
}

// setRequestCookies adds the cookies of the given request with the provided key-value pairs.
func setRequestCookies(req *fasthttp.Request, cookies []types.KeyValue[string, string]) {
	for _, cookie := range cookies {
		req.Header.Add("Cookie", cookie.Key+"="+cookie.Value)
	}
}

// setRequestMethod sets the HTTP request method for the given request.
func setRequestMethod(req *fasthttp.Request, method string) {
	req.Header.SetMethod(method)
}

// setRequestBody sets the request body of the given fasthttp.Request object.
// The body parameter is a string that will be converted to a byte slice and set as the request body.
func setRequestBody(req *fasthttp.Request, body string) {
	req.SetBody([]byte(body))
}

// getKeyValueGeneratorFunc creates a function that generates key-value pairs for HTTP requests.
// It takes a slice of key-value pairs where each key maps to a slice of possible values,
// and a random number generator.
//
// If any key has multiple possible values, the function will randomly select one value for each
// call (using the provided random number generator). If all keys have at most one value, the
// function will always return the same set of key-value pairs for efficiency.
func getKeyValueGeneratorFunc[
	T []types.KeyValue[string, string],
](
	keyValueSlice []types.KeyValue[string, []string],
	localRand *rand.Rand,
) func() T {
	keyValueGenerators := make([]keyValueGenerator, len(keyValueSlice))

	funcMap := newFuncMap(localRand)

	for i, kv := range keyValueSlice {
		keyValueGenerators[i] = keyValueGenerator{
			key:   getKeyFunc(kv.Key, funcMap),
			value: getValueFunc(kv.Value, funcMap, localRand),
		}
	}

	return func() T {
		keyValues := make(T, len(keyValueGenerators))
		for i, keyValue := range keyValueGenerators {
			keyValues[i] = types.KeyValue[string, string]{
				Key:   keyValue.key(),
				Value: keyValue.value(),
			}
		}
		return keyValues
	}
}

func getKeyFunc(key string, funcMap template.FuncMap) func() string {
	t, err := template.New("default").Funcs(funcMap).Parse(key)
	if err != nil {
		panic(err)
	}

	return func() string {
		var buf bytes.Buffer
		_ = t.Execute(&buf, nil)
		return buf.String()
	}
}

func getValueFunc(
	values []string,
	funcMap template.FuncMap,
	localRand *rand.Rand,
) func() string {
	templates := make([]*template.Template, len(values))

	for i, value := range values {
		t, err := template.New("default").Funcs(funcMap).Parse(value)
		if err != nil {
			panic(err)
		}
		templates[i] = t
	}

	randomTemplateFunc := utils.RandomValueCycle(templates, localRand)

	return func() string {
		if tmpl := randomTemplateFunc(); tmpl == nil {
			return ""
		} else {
			var buf bytes.Buffer
			_ = tmpl.Execute(&buf, nil)
			return buf.String()
		}
	}
}

func newFuncMap(localRand *rand.Rand) template.FuncMap {
	localFaker := gofakeit.NewFaker(localRand, false)

	return template.FuncMap{
		"upper":       strings.ToUpper,
		"fakeit_Name": localFaker.Name,
	}
}