dodo/requests/request.go

package requests

import (
	"context"
	"math/rand"
	"net/url"
	"time"

	"github.com/aykhans/dodo/config"
	"github.com/aykhans/dodo/types"
	"github.com/aykhans/dodo/utils"
	"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
}

// 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 {
	bodiesLen := len(bodies)
	getBody := func() string { return "" }
	if bodiesLen == 1 {
		getBody = func() string { return bodies[0] }
	} else if bodiesLen > 1 {
		getBody = utils.RandomValueCycle(bodies, localRand)
	}

	getParams := getKeyValueGeneratorFunc(params, localRand)
	getHeaders := getKeyValueGeneratorFunc(headers, localRand)
	getCookies := getKeyValueGeneratorFunc(cookies, 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)
	}
	if req.Header.UserAgent() == nil {
		req.Header.SetUserAgent(config.DefaultUserAgent)
	}
}

// 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 {
	getKeyValueSlice := []map[string]func() string{}
	isRandom := false

	for _, kv := range keyValueSlice {
		valuesLen := len(kv.Value)

		getValueFunc := func() string { return "" }
		if valuesLen == 1 {
			getValueFunc = func() string { return kv.Value[0] }
		} else if valuesLen > 1 {
			getValueFunc = utils.RandomValueCycle(kv.Value, localRand)
			isRandom = true
		}

		getKeyValueSlice = append(
			getKeyValueSlice,
			map[string]func() string{kv.Key: getValueFunc},
		)
	}

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