mirror of
https://github.com/aykhans/sarin.git
synced 2026-10-13 21:28:52 +00:00
feat(script): add http() and json to Lua and JS scripts
This commit is contained in:
17 files changed
+990
-140
No files matched your search
@@ -715,7 +715,7 @@ func validateScriptSource(script string) error {
|
||||
return types.ErrScriptEmpty
|
||||
}
|
||||
|
||||
// Not a file/URL reference - it's an inline script
|
||||
// Not a file/URL reference, so it's an inline script
|
||||
if !strings.HasPrefix(script, "@") {
|
||||
return nil
|
||||
}
|
||||
@@ -725,7 +725,7 @@ func validateScriptSource(script string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// It's a file or URL reference - validate the source
|
||||
// It's a file or URL reference, so validate the source
|
||||
source := script[1:] // Remove the @ prefix
|
||||
|
||||
if source == "" {
|
||||
@@ -744,7 +744,7 @@ func validateScriptSource(script string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// It's a file path - basic validation (not empty, checked above)
|
||||
// It's a file path, already checked to be non-empty above
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -822,7 +822,7 @@ func (m printConfigModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
}
|
||||
|
||||
func (m printConfigModel) View() tea.View {
|
||||
// AltScreen and MouseMode were program options in bubbletea v1; in v2 they
|
||||
// AltScreen and MouseMode were program options in bubbletea v1, in v2 they
|
||||
// are view properties, so every return path has to declare them.
|
||||
newView := func(s string) tea.View {
|
||||
v := tea.NewView(s)
|
||||
|
||||
@@ -193,7 +193,7 @@ func validateTemplateURLPath(urlPath string, funcMap template.FuncMap) []types.F
|
||||
|
||||
func ValidateTemplates(config *Config) []types.FieldValidationError {
|
||||
// Create template function map using the same functions as sarin package
|
||||
// Use nil for fileCache during validation - templates are only parsed, not executed
|
||||
// fileCache is nil because templates are only parsed here, not executed
|
||||
randSource := sarin.NewDefaultRandSource()
|
||||
funcMap := sarin.NewDefaultTemplateFuncMap(randSource, nil)
|
||||
|
||||
|
||||
@@ -18,7 +18,8 @@ import (
|
||||
"golang.org/x/net/proxy"
|
||||
)
|
||||
|
||||
type HostClientGenerator func() *fasthttp.HostClient
|
||||
// clientIndexGenerator returns the index of the host client (and proxy) to use next.
|
||||
type clientIndexGenerator func() int
|
||||
|
||||
func safeUintToInt(u uint) int {
|
||||
if u > math.MaxInt {
|
||||
@@ -306,18 +307,10 @@ func (c *bufferedConn) Read(b []byte) (int, error) {
|
||||
return c.reader.Read(b)
|
||||
}
|
||||
|
||||
func NewHostClientGenerator(clients ...*fasthttp.HostClient) HostClientGenerator {
|
||||
switch len(clients) {
|
||||
case 0:
|
||||
hostClient := &fasthttp.HostClient{}
|
||||
return func() *fasthttp.HostClient {
|
||||
return hostClient
|
||||
}
|
||||
case 1:
|
||||
return func() *fasthttp.HostClient {
|
||||
return clients[0]
|
||||
}
|
||||
default:
|
||||
return utilsSlice.RandomCycle(nil, clients...)
|
||||
func newClientIndexGenerator(count int) clientIndexGenerator {
|
||||
indexes := make([]int, count)
|
||||
for i := range indexes {
|
||||
indexes[i] = i
|
||||
}
|
||||
return utilsSlice.RandomCycle(nil, indexes...)
|
||||
}
|
||||
@@ -58,8 +58,7 @@ func (fc *FileCache) GetOrLoad(source string) (*CachedFile, error) {
|
||||
|
||||
file := &CachedFile{Content: content, Filename: filename}
|
||||
|
||||
// LoadOrStore handles race condition - if another goroutine
|
||||
// cached it first, we get theirs (no duplicate storage)
|
||||
// If another goroutine cached it first, LoadOrStore returns theirs
|
||||
actual, _ := fc.cache.LoadOrStore(source, file)
|
||||
return actual.(*CachedFile), nil
|
||||
}
|
||||
|
||||
+45
-56
@@ -199,10 +199,11 @@ type sarin struct {
|
||||
logError bool
|
||||
logFile string
|
||||
|
||||
hostClients []*fasthttp.HostClient
|
||||
responses *SarinResponseData
|
||||
fileCache *FileCache
|
||||
scriptChain *script.Chain
|
||||
hostClients []*fasthttp.HostClient
|
||||
scriptHTTPClients []*scriptHTTPClient
|
||||
responses *SarinResponseData
|
||||
fileCache *FileCache
|
||||
scriptChain *script.Chain
|
||||
}
|
||||
|
||||
// NewSarin creates a new sarin instance for load testing.
|
||||
@@ -248,7 +249,12 @@ func NewSarin(
|
||||
}
|
||||
}
|
||||
|
||||
hostClients, err := newHostClients(ctx, timeout, proxies, workers, requestURL, skipCertVerify)
|
||||
proxyURLs := make([]url.URL, len(proxies))
|
||||
for i, proxy := range proxies {
|
||||
proxyURLs[i] = url.URL(proxy)
|
||||
}
|
||||
|
||||
hostClients, err := NewHostClients(ctx, timeout, proxyURLs, workers, requestURL, skipCertVerify)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -266,28 +272,37 @@ func NewSarin(
|
||||
|
||||
scriptChain := script.NewChain(luaSources, jsSources)
|
||||
|
||||
var scriptHTTPClients []*scriptHTTPClient
|
||||
if !scriptChain.IsEmpty() {
|
||||
scriptHTTPClients, err = newScriptHTTPClients(ctx, proxyURLs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
srn := &sarin{
|
||||
workers: workers,
|
||||
requestURL: requestURL,
|
||||
methods: methods,
|
||||
params: params,
|
||||
headers: headers,
|
||||
cookies: cookies,
|
||||
bodies: bodies,
|
||||
totalRequests: totalRequests,
|
||||
totalDuration: totalDuration,
|
||||
timeout: timeout,
|
||||
showProgress: showProgress,
|
||||
skipCertVerify: skipCertVerify,
|
||||
values: values,
|
||||
collectStats: collectStats,
|
||||
dryRun: dryRun,
|
||||
logInfo: logInfo,
|
||||
logError: logError,
|
||||
logFile: logFile,
|
||||
hostClients: hostClients,
|
||||
fileCache: NewFileCache(time.Second * 10),
|
||||
scriptChain: scriptChain,
|
||||
workers: workers,
|
||||
requestURL: requestURL,
|
||||
methods: methods,
|
||||
params: params,
|
||||
headers: headers,
|
||||
cookies: cookies,
|
||||
bodies: bodies,
|
||||
totalRequests: totalRequests,
|
||||
totalDuration: totalDuration,
|
||||
timeout: timeout,
|
||||
showProgress: showProgress,
|
||||
skipCertVerify: skipCertVerify,
|
||||
values: values,
|
||||
collectStats: collectStats,
|
||||
dryRun: dryRun,
|
||||
logInfo: logInfo,
|
||||
logError: logError,
|
||||
logFile: logFile,
|
||||
hostClients: hostClients,
|
||||
scriptHTTPClients: scriptHTTPClients,
|
||||
fileCache: NewFileCache(time.Second * 10),
|
||||
scriptChain: scriptChain,
|
||||
}
|
||||
|
||||
if collectStats {
|
||||
@@ -364,7 +379,7 @@ func (s sarin) Start(ctx context.Context, stopCtrl *StopController) {
|
||||
}
|
||||
|
||||
// Start workers
|
||||
s.startWorkers(&workersWG, jobsCh, s.hostClients, &counter, sendLog, sendRespLog)
|
||||
s.startWorkers(&workersWG, jobsCh, &counter, sendLog, sendRespLog)
|
||||
|
||||
if runTUI {
|
||||
//nolint:contextcheck // streamCtx must remain active until all workers complete to ensure all collected data is streamed
|
||||
@@ -394,7 +409,7 @@ func (s sarin) Start(ctx context.Context, stopCtrl *StopController) {
|
||||
}
|
||||
|
||||
// newWriterLog builds the loggers that write formatted lines to w (a log file or
|
||||
// stderr). sendLog stays general (it filters by each log's level); sendRespLog
|
||||
// stderr). sendLog stays general (it filters by each log's level), while sendRespLog
|
||||
// only ever emits info, so its decision is baked once into a no-op when off.
|
||||
func (s sarin) newWriterLog(w io.Writer) (runtimeLogger, respLogger) {
|
||||
// log.Logger serializes writes with its own mutex, so concurrent workers
|
||||
@@ -432,36 +447,10 @@ func (s sarin) newChannelLog(ch chan<- runtimeLog) (runtimeLogger, respLogger) {
|
||||
return sendLog, sendRespLog
|
||||
}
|
||||
|
||||
// newHostClients initializes HTTP clients for the given configuration.
|
||||
// It can return the following errors:
|
||||
// - types.ProxyDialError
|
||||
func newHostClients(
|
||||
ctx context.Context,
|
||||
timeout time.Duration,
|
||||
proxies types.Proxies,
|
||||
workers uint,
|
||||
requestURL *url.URL,
|
||||
skipCertVerify bool,
|
||||
) ([]*fasthttp.HostClient, error) {
|
||||
proxiesRaw := make([]url.URL, len(proxies))
|
||||
for i, proxy := range proxies {
|
||||
proxiesRaw[i] = url.URL(proxy)
|
||||
}
|
||||
|
||||
return NewHostClients(
|
||||
ctx,
|
||||
timeout,
|
||||
proxiesRaw,
|
||||
workers,
|
||||
requestURL,
|
||||
skipCertVerify,
|
||||
)
|
||||
}
|
||||
|
||||
func (s sarin) startWorkers(wg *sync.WaitGroup, jobs <-chan struct{}, hostClients []*fasthttp.HostClient, counter *atomic.Uint64, sendLog runtimeLogger, sendRespLog respLogger) {
|
||||
func (s sarin) startWorkers(wg *sync.WaitGroup, jobs <-chan struct{}, counter *atomic.Uint64, sendLog runtimeLogger, sendRespLog respLogger) {
|
||||
for range max(s.workers, 1) {
|
||||
wg.Go(func() {
|
||||
s.Worker(jobs, NewHostClientGenerator(hostClients...), counter, sendLog, sendRespLog)
|
||||
s.Worker(jobs, counter, sendLog, sendRespLog)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
package sarin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"net/url"
|
||||
"time"
|
||||
|
||||
"github.com/valyala/fasthttp"
|
||||
"go.aykhans.me/sarin/internal/script"
|
||||
"go.aykhans.me/sarin/internal/types"
|
||||
)
|
||||
|
||||
// scriptHTTPMaxRedirects is the redirect limit when followRedirects is set.
|
||||
const scriptHTTPMaxRedirects = 10
|
||||
|
||||
// scriptHTTPDefaultTimeout is the default timeout for script requests, independent of -T.
|
||||
const scriptHTTPDefaultTimeout = 30 * time.Second
|
||||
|
||||
// scriptHTTPClient sends scripts' http.* requests. It is safe for concurrent use.
|
||||
type scriptHTTPClient struct {
|
||||
client *fasthttp.Client
|
||||
insecureClient *fasthttp.Client
|
||||
defaultTimeout time.Duration
|
||||
}
|
||||
|
||||
var _ script.HTTPDoer = (*scriptHTTPClient)(nil)
|
||||
|
||||
// newScriptHTTPClients creates one client per proxy, in the same order as NewHostClients.
|
||||
// It can return the following errors:
|
||||
// - types.ProxyDialError
|
||||
func newScriptHTTPClients(ctx context.Context, proxies []url.URL) ([]*scriptHTTPClient, error) {
|
||||
if len(proxies) == 0 {
|
||||
return []*scriptHTTPClient{newScriptHTTPClient(nil, scriptHTTPDefaultTimeout)}, nil
|
||||
}
|
||||
|
||||
clients := make([]*scriptHTTPClient, 0, len(proxies))
|
||||
for _, proxy := range proxies {
|
||||
dialFunc, err := NewProxyDialFunc(ctx, &proxy, scriptHTTPDefaultTimeout)
|
||||
if err != nil {
|
||||
return nil, types.NewProxyDialError(proxy.String(), err)
|
||||
}
|
||||
clients = append(clients, newScriptHTTPClient(dialFunc, scriptHTTPDefaultTimeout))
|
||||
}
|
||||
return clients, nil
|
||||
}
|
||||
|
||||
func newScriptHTTPClient(dialFunc fasthttp.DialFunc, defaultTimeout time.Duration) *scriptHTTPClient {
|
||||
newClient := func(insecure bool) *fasthttp.Client {
|
||||
return &fasthttp.Client{
|
||||
Dial: dialFunc,
|
||||
TLSConfig: &tls.Config{
|
||||
InsecureSkipVerify: insecure, //nolint:gosec
|
||||
},
|
||||
DisableHeaderNamesNormalizing: true,
|
||||
DisablePathNormalizing: true,
|
||||
NoDefaultUserAgentHeader: true,
|
||||
}
|
||||
}
|
||||
|
||||
return &scriptHTTPClient{
|
||||
client: newClient(false),
|
||||
insecureClient: newClient(true),
|
||||
defaultTimeout: defaultTimeout,
|
||||
}
|
||||
}
|
||||
|
||||
// Do sends the request. With followRedirects the timeout applies to each hop.
|
||||
// It can return the following errors:
|
||||
// - types.ScriptHTTPRequestError
|
||||
func (c *scriptHTTPClient) Do(r *script.HTTPRequest) (*script.HTTPResponse, error) {
|
||||
parsedURL, err := url.Parse(r.URL)
|
||||
if err != nil {
|
||||
return nil, types.NewScriptHTTPRequestError(r.Method, r.URL, err)
|
||||
}
|
||||
if (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") || parsedURL.Host == "" {
|
||||
return nil, types.NewScriptHTTPRequestError(r.Method, r.URL, types.ErrScriptHTTPURLInvalid)
|
||||
}
|
||||
|
||||
req := fasthttp.AcquireRequest()
|
||||
resp := fasthttp.AcquireResponse()
|
||||
defer fasthttp.ReleaseRequest(req)
|
||||
defer fasthttp.ReleaseResponse(resp)
|
||||
|
||||
req.SetRequestURI(r.URL)
|
||||
req.Header.SetMethod(r.Method)
|
||||
for key, values := range r.Headers {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
if len(r.Params) > 0 {
|
||||
args := req.URI().QueryArgs()
|
||||
for key, values := range r.Params {
|
||||
for _, value := range values {
|
||||
args.Add(key, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(r.Cookies) > 0 {
|
||||
req.Header.Add("Cookie", cookieHeaderValue(r.Cookies))
|
||||
}
|
||||
|
||||
req.SetBodyString(r.Body)
|
||||
|
||||
timeout := r.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = c.defaultTimeout
|
||||
}
|
||||
|
||||
client := c.client
|
||||
if r.Insecure {
|
||||
client = c.insecureClient
|
||||
}
|
||||
|
||||
req.SetTimeout(timeout)
|
||||
if r.FollowRedirects {
|
||||
err = client.DoRedirects(req, resp, scriptHTTPMaxRedirects)
|
||||
} else {
|
||||
err = client.Do(req, resp)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, types.NewScriptHTTPRequestError(r.Method, r.URL, err)
|
||||
}
|
||||
|
||||
return &script.HTTPResponse{
|
||||
Status: resp.StatusCode(),
|
||||
Headers: collectRespHeaders(resp),
|
||||
Body: string(resp.Body()),
|
||||
}, nil
|
||||
}
|
||||
@@ -97,7 +97,7 @@ func NewDefaultTemplateFuncMap(randSource rand.Source, fileCache *FileCache) tem
|
||||
return string(data), nil
|
||||
},
|
||||
// json_Object builds a JSON object from interleaved key-value pairs and returns it
|
||||
// as a JSON string. Keys must be strings; values may be any JSON-encodable type.
|
||||
// as a JSON string. Keys must be strings, values may be any JSON-encodable type.
|
||||
// Usage: {{ json_Object "name" "Alice" "age" 30 }}
|
||||
"json_Object": func(pairs ...any) (string, error) {
|
||||
if len(pairs)%2 != 0 {
|
||||
@@ -751,7 +751,7 @@ func NewDefaultBodyTemplateFuncMap(
|
||||
|
||||
switch {
|
||||
case strings.HasPrefix(val, "@@"):
|
||||
// Escaped @ - send as literal string without first @
|
||||
// Escaped @: send the string without the first @
|
||||
if err := writer.WriteField(key, val[1:]); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
+46
-18
@@ -31,7 +31,6 @@ func statusCodeToString(code int) string {
|
||||
|
||||
func (s sarin) Worker(
|
||||
jobs <-chan struct{},
|
||||
hostClientGenerator HostClientGenerator,
|
||||
counter *atomic.Uint64,
|
||||
sendLog runtimeLogger,
|
||||
sendRespLog respLogger,
|
||||
@@ -53,6 +52,9 @@ func (s sarin) Worker(
|
||||
defer scriptTransformer.Close()
|
||||
}
|
||||
|
||||
nextClientIndex := newClientIndexGenerator(len(s.hostClients))
|
||||
bindScriptProxy := s.newScriptProxyBinder(scriptTransformer)
|
||||
|
||||
requestGenerator, isDynamic := NewRequestGenerator(
|
||||
s.methods, s.requestURL, s.params, s.headers, s.cookies, s.bodies, s.values, s.fileCache, scriptTransformer,
|
||||
)
|
||||
@@ -60,34 +62,45 @@ func (s sarin) Worker(
|
||||
if s.dryRun {
|
||||
switch {
|
||||
case s.collectStats && isDynamic:
|
||||
s.workerDryRunStatsWithDynamic(jobs, req, requestGenerator, counter, sendLog)
|
||||
s.workerDryRunStatsWithDynamic(jobs, req, requestGenerator, nextClientIndex, bindScriptProxy, counter, sendLog)
|
||||
case s.collectStats && !isDynamic:
|
||||
s.workerDryRunStatsWithStatic(jobs, req, requestGenerator, counter, sendLog)
|
||||
case !s.collectStats && isDynamic:
|
||||
s.workerDryRunNoStatsWithDynamic(jobs, req, requestGenerator, counter, sendLog)
|
||||
s.workerDryRunNoStatsWithDynamic(jobs, req, requestGenerator, nextClientIndex, bindScriptProxy, counter, sendLog)
|
||||
default:
|
||||
s.workerDryRunNoStatsWithStatic(jobs, req, requestGenerator, counter, sendLog)
|
||||
}
|
||||
} else {
|
||||
switch {
|
||||
case s.collectStats && isDynamic:
|
||||
s.workerStatsWithDynamic(jobs, req, resp, requestGenerator, hostClientGenerator, counter, sendLog, sendRespLog)
|
||||
s.workerStatsWithDynamic(jobs, req, resp, requestGenerator, nextClientIndex, bindScriptProxy, counter, sendLog, sendRespLog)
|
||||
case s.collectStats && !isDynamic:
|
||||
s.workerStatsWithStatic(jobs, req, resp, requestGenerator, hostClientGenerator, counter, sendLog, sendRespLog)
|
||||
s.workerStatsWithStatic(jobs, req, resp, requestGenerator, nextClientIndex, counter, sendLog, sendRespLog)
|
||||
case !s.collectStats && isDynamic:
|
||||
s.workerNoStatsWithDynamic(jobs, req, resp, requestGenerator, hostClientGenerator, counter, sendLog, sendRespLog)
|
||||
s.workerNoStatsWithDynamic(jobs, req, resp, requestGenerator, nextClientIndex, bindScriptProxy, counter, sendLog, sendRespLog)
|
||||
default:
|
||||
s.workerNoStatsWithStatic(jobs, req, resp, requestGenerator, hostClientGenerator, counter, sendLog, sendRespLog)
|
||||
s.workerNoStatsWithStatic(jobs, req, resp, requestGenerator, nextClientIndex, counter, sendLog, sendRespLog)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// newScriptProxyBinder returns a func that sends the scripts' http.* calls through proxy i.
|
||||
func (s sarin) newScriptProxyBinder(scriptTransformer *script.Transformer) func(i int) {
|
||||
if scriptTransformer == nil || scriptTransformer.IsEmpty() {
|
||||
return func(int) {}
|
||||
}
|
||||
return func(i int) {
|
||||
scriptTransformer.SetHTTPDoer(s.scriptHTTPClients[i])
|
||||
}
|
||||
}
|
||||
|
||||
func (s sarin) workerStatsWithDynamic(
|
||||
jobs <-chan struct{},
|
||||
req *fasthttp.Request,
|
||||
resp *fasthttp.Response,
|
||||
requestGenerator RequestGenerator,
|
||||
hostClientGenerator HostClientGenerator,
|
||||
nextClientIndex clientIndexGenerator,
|
||||
bindScriptProxy func(int),
|
||||
counter *atomic.Uint64,
|
||||
sendLog runtimeLogger,
|
||||
sendRespLog respLogger,
|
||||
@@ -95,6 +108,9 @@ func (s sarin) workerStatsWithDynamic(
|
||||
for range jobs {
|
||||
req.Reset()
|
||||
|
||||
// Bound before generation so scripts' http.* calls use the main request's proxy.
|
||||
clientIndex := nextClientIndex()
|
||||
bindScriptProxy(clientIndex)
|
||||
if err := requestGenerator(req); err != nil {
|
||||
s.responses.Add(err.Error(), 0)
|
||||
sendLog(runtimeLogLevelError, err.Error())
|
||||
@@ -103,7 +119,7 @@ func (s sarin) workerStatsWithDynamic(
|
||||
}
|
||||
|
||||
startTime := time.Now()
|
||||
err := hostClientGenerator().DoTimeout(req, resp, s.timeout)
|
||||
err := s.hostClients[clientIndex].DoTimeout(req, resp, s.timeout)
|
||||
respDuration := time.Since(startTime)
|
||||
|
||||
if err != nil {
|
||||
@@ -121,13 +137,13 @@ func (s sarin) workerStatsWithStatic(
|
||||
req *fasthttp.Request,
|
||||
resp *fasthttp.Response,
|
||||
requestGenerator RequestGenerator,
|
||||
hostClientGenerator HostClientGenerator,
|
||||
nextClientIndex clientIndexGenerator,
|
||||
counter *atomic.Uint64,
|
||||
sendLog runtimeLogger,
|
||||
sendRespLog respLogger,
|
||||
) {
|
||||
if err := requestGenerator(req); err != nil {
|
||||
// Static request generation failed - record all jobs as errors
|
||||
// Static request generation failed, so record all jobs as errors
|
||||
for range jobs {
|
||||
s.responses.Add(err.Error(), 0)
|
||||
sendLog(runtimeLogLevelError, err.Error())
|
||||
@@ -138,7 +154,7 @@ func (s sarin) workerStatsWithStatic(
|
||||
|
||||
for range jobs {
|
||||
startTime := time.Now()
|
||||
err := hostClientGenerator().DoTimeout(req, resp, s.timeout)
|
||||
err := s.hostClients[nextClientIndex()].DoTimeout(req, resp, s.timeout)
|
||||
respDuration := time.Since(startTime)
|
||||
if err != nil {
|
||||
s.responses.Add(err.Error(), respDuration)
|
||||
@@ -155,20 +171,24 @@ func (s sarin) workerNoStatsWithDynamic(
|
||||
req *fasthttp.Request,
|
||||
resp *fasthttp.Response,
|
||||
requestGenerator RequestGenerator,
|
||||
hostClientGenerator HostClientGenerator,
|
||||
nextClientIndex clientIndexGenerator,
|
||||
bindScriptProxy func(int),
|
||||
counter *atomic.Uint64,
|
||||
sendLog runtimeLogger,
|
||||
sendRespLog respLogger,
|
||||
) {
|
||||
for range jobs {
|
||||
req.Reset()
|
||||
// Bound before generation so scripts' http.* calls use the main request's proxy.
|
||||
clientIndex := nextClientIndex()
|
||||
bindScriptProxy(clientIndex)
|
||||
if err := requestGenerator(req); err != nil {
|
||||
sendLog(runtimeLogLevelError, err.Error())
|
||||
counter.Add(1)
|
||||
continue
|
||||
}
|
||||
startTime := time.Now()
|
||||
err := hostClientGenerator().DoTimeout(req, resp, s.timeout)
|
||||
err := s.hostClients[clientIndex].DoTimeout(req, resp, s.timeout)
|
||||
if err == nil {
|
||||
sendRespLog(time.Since(startTime), resp)
|
||||
}
|
||||
@@ -181,7 +201,7 @@ func (s sarin) workerNoStatsWithStatic(
|
||||
req *fasthttp.Request,
|
||||
resp *fasthttp.Response,
|
||||
requestGenerator RequestGenerator,
|
||||
hostClientGenerator HostClientGenerator,
|
||||
nextClientIndex clientIndexGenerator,
|
||||
counter *atomic.Uint64,
|
||||
sendLog runtimeLogger,
|
||||
sendRespLog respLogger,
|
||||
@@ -189,7 +209,7 @@ func (s sarin) workerNoStatsWithStatic(
|
||||
if err := requestGenerator(req); err != nil {
|
||||
sendLog(runtimeLogLevelError, err.Error())
|
||||
|
||||
// Static request generation failed - just count the jobs without sending
|
||||
// Static request generation failed, so just count the jobs without sending
|
||||
for range jobs {
|
||||
counter.Add(1)
|
||||
}
|
||||
@@ -198,7 +218,7 @@ func (s sarin) workerNoStatsWithStatic(
|
||||
|
||||
for range jobs {
|
||||
startTime := time.Now()
|
||||
err := hostClientGenerator().DoTimeout(req, resp, s.timeout)
|
||||
err := s.hostClients[nextClientIndex()].DoTimeout(req, resp, s.timeout)
|
||||
if err == nil {
|
||||
sendRespLog(time.Since(startTime), resp)
|
||||
}
|
||||
@@ -210,11 +230,15 @@ func (s sarin) workerDryRunStatsWithDynamic(
|
||||
jobs <-chan struct{},
|
||||
req *fasthttp.Request,
|
||||
requestGenerator RequestGenerator,
|
||||
nextClientIndex clientIndexGenerator,
|
||||
bindScriptProxy func(int),
|
||||
counter *atomic.Uint64,
|
||||
sendLog runtimeLogger,
|
||||
) {
|
||||
for range jobs {
|
||||
req.Reset()
|
||||
// Nothing is sent, this only binds a proxy for the scripts' http.* calls.
|
||||
bindScriptProxy(nextClientIndex())
|
||||
startTime := time.Now()
|
||||
if err := requestGenerator(req); err != nil {
|
||||
s.responses.Add(err.Error(), time.Since(startTime))
|
||||
@@ -235,7 +259,7 @@ func (s sarin) workerDryRunStatsWithStatic(
|
||||
sendLog runtimeLogger,
|
||||
) {
|
||||
if err := requestGenerator(req); err != nil {
|
||||
// Static request generation failed - record all jobs as errors
|
||||
// Static request generation failed, so record all jobs as errors
|
||||
for range jobs {
|
||||
s.responses.Add(err.Error(), 0)
|
||||
sendLog(runtimeLogLevelError, err.Error())
|
||||
@@ -254,11 +278,15 @@ func (s sarin) workerDryRunNoStatsWithDynamic(
|
||||
jobs <-chan struct{},
|
||||
req *fasthttp.Request,
|
||||
requestGenerator RequestGenerator,
|
||||
nextClientIndex clientIndexGenerator,
|
||||
bindScriptProxy func(int),
|
||||
counter *atomic.Uint64,
|
||||
sendLog runtimeLogger,
|
||||
) {
|
||||
for range jobs {
|
||||
req.Reset()
|
||||
// Nothing is sent, this only binds a proxy for the scripts' http.* calls.
|
||||
bindScriptProxy(nextClientIndex())
|
||||
if err := requestGenerator(req); err != nil {
|
||||
sendLog(runtimeLogLevelError, err.Error())
|
||||
}
|
||||
|
||||
@@ -91,6 +91,16 @@ func (t *Transformer) Transform(req *RequestData) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetHTTPDoer sets the doer that the scripts' http.* calls are sent through.
|
||||
func (t *Transformer) SetHTTPDoer(doer HTTPDoer) {
|
||||
for _, engine := range t.luaEngines {
|
||||
engine.SetHTTPDoer(doer)
|
||||
}
|
||||
for _, engine := range t.jsEngines {
|
||||
engine.SetHTTPDoer(doer)
|
||||
}
|
||||
}
|
||||
|
||||
// Close releases all engine resources.
|
||||
func (t *Transformer) Close() {
|
||||
for _, engine := range t.luaEngines {
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
package script
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.aykhans.me/sarin/internal/types"
|
||||
)
|
||||
|
||||
// HTTPRequest is a request made by a script through the http.* functions.
|
||||
type HTTPRequest struct {
|
||||
Method string
|
||||
URL string
|
||||
Headers map[string][]string
|
||||
Params map[string][]string
|
||||
Cookies map[string][]string
|
||||
Body string
|
||||
Timeout time.Duration // zero means the doer's default timeout
|
||||
Insecure bool
|
||||
FollowRedirects bool
|
||||
}
|
||||
|
||||
// HTTPResponse is the response returned to a script by the http.* functions.
|
||||
type HTTPResponse struct {
|
||||
Status int
|
||||
Headers map[string][]string
|
||||
Body string
|
||||
}
|
||||
|
||||
// Header returns the first value of the named header, matching the name case-insensitively.
|
||||
func (r *HTTPResponse) Header(name string) (string, bool) {
|
||||
for key, values := range r.Headers {
|
||||
if strings.EqualFold(key, name) && len(values) > 0 {
|
||||
return values[0], true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// Cookie returns the value of the named cookie from the Set-Cookie headers.
|
||||
func (r *HTTPResponse) Cookie(name string) (string, bool) {
|
||||
for key, values := range r.Headers {
|
||||
if !strings.EqualFold(key, "Set-Cookie") {
|
||||
continue
|
||||
}
|
||||
for _, line := range values {
|
||||
if cookie, err := http.ParseSetCookie(line); err == nil && cookie.Name == name {
|
||||
return cookie.Value, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// HTTPDoer sends the HTTP requests that scripts make.
|
||||
type HTTPDoer interface {
|
||||
Do(req *HTTPRequest) (*HTTPResponse, error)
|
||||
}
|
||||
|
||||
// httpDefaultMethod is used when the options set no method.
|
||||
const httpDefaultMethod = http.MethodGet
|
||||
|
||||
// Option keys accepted by the http function.
|
||||
const (
|
||||
httpOptionMethod = "method"
|
||||
httpOptionHeaders = "headers"
|
||||
httpOptionParams = "params"
|
||||
httpOptionCookies = "cookies"
|
||||
httpOptionBody = "body"
|
||||
httpOptionTimeout = "timeout"
|
||||
httpOptionInsecure = "insecure"
|
||||
httpOptionFollowRedirects = "followRedirects"
|
||||
)
|
||||
|
||||
// httpBridge connects http to the doer. Calls only work while transform runs.
|
||||
type httpBridge struct {
|
||||
doer HTTPDoer
|
||||
active bool
|
||||
}
|
||||
|
||||
// It can return the following errors:
|
||||
// - types.ErrScriptHTTPOutsideTransform
|
||||
// - types.ErrScriptHTTPUnavailable
|
||||
// - types.ErrScriptHTTPMethodEmpty
|
||||
// - any error returned by the doer
|
||||
func (b *httpBridge) do(req *HTTPRequest) (*HTTPResponse, error) {
|
||||
if !b.active {
|
||||
return nil, types.ErrScriptHTTPOutsideTransform
|
||||
}
|
||||
if b.doer == nil {
|
||||
return nil, types.ErrScriptHTTPUnavailable
|
||||
}
|
||||
if req.Method == "" {
|
||||
return nil, types.ErrScriptHTTPMethodEmpty
|
||||
}
|
||||
return b.doer.Do(req)
|
||||
}
|
||||
|
||||
// parseHTTPTimeout parses a Go duration string such as "500ms" or "2s".
|
||||
// It can return the following errors:
|
||||
// - types.ScriptHTTPOptionError
|
||||
func parseHTTPTimeout(value string) (time.Duration, error) {
|
||||
timeout, err := time.ParseDuration(value)
|
||||
if err != nil {
|
||||
return 0, types.NewScriptHTTPOptionError(httpOptionTimeout, err)
|
||||
}
|
||||
if timeout <= 0 {
|
||||
return 0, types.NewScriptHTTPOptionError(httpOptionTimeout, types.NewScriptTypeError("a duration greater than 0", value))
|
||||
}
|
||||
return timeout, nil
|
||||
}
|
||||
+35
-7
@@ -11,6 +11,10 @@ import (
|
||||
type JsEngine struct {
|
||||
runtime *goja.Runtime
|
||||
transform goja.Callable
|
||||
jsonParse goja.Callable
|
||||
// errorConstructor is the built-in Error, captured before the script can replace it.
|
||||
errorConstructor goja.Value
|
||||
http httpBridge
|
||||
}
|
||||
|
||||
// NewJsEngine creates a new JavaScript script engine with the given script content.
|
||||
@@ -29,11 +33,17 @@ type JsEngine struct {
|
||||
// - types.ScriptExecutionError
|
||||
func NewJsEngine(scriptContent string) (*JsEngine, error) {
|
||||
vm := goja.New()
|
||||
engine := &JsEngine{runtime: vm}
|
||||
|
||||
// Register the globals before running the script so its functions can use them
|
||||
if err := engine.registerHTTP(); err != nil {
|
||||
return nil, types.NewScriptExecutionError("JavaScript", err)
|
||||
}
|
||||
|
||||
// Execute the script to define the transform function
|
||||
_, err := vm.RunString(scriptContent)
|
||||
if err != nil {
|
||||
return nil, types.NewScriptExecutionError("JavaScript", err)
|
||||
return nil, types.NewScriptExecutionError("JavaScript", jsExceptionError(err))
|
||||
}
|
||||
|
||||
// Get the transform function
|
||||
@@ -47,23 +57,30 @@ func NewJsEngine(scriptContent string) (*JsEngine, error) {
|
||||
return nil, types.NewScriptExecutionError("JavaScript", errors.New("'transform' must be a function"))
|
||||
}
|
||||
|
||||
return &JsEngine{
|
||||
runtime: vm,
|
||||
transform: transform,
|
||||
}, nil
|
||||
engine.transform = transform
|
||||
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
// SetHTTPDoer sets the doer that the script's http.* calls are sent through.
|
||||
func (e *JsEngine) SetHTTPDoer(doer HTTPDoer) {
|
||||
e.http.doer = doer
|
||||
}
|
||||
|
||||
// Transform executes the JavaScript transform function with the given request data.
|
||||
// It can return the following errors:
|
||||
// - types.ScriptExecutionError
|
||||
func (e *JsEngine) Transform(req *RequestData) error {
|
||||
e.http.active = true
|
||||
defer func() { e.http.active = false }()
|
||||
|
||||
// Convert RequestData to JavaScript object
|
||||
reqObj := e.requestDataToObject(req)
|
||||
|
||||
// Call transform(req)
|
||||
result, err := e.transform(goja.Undefined(), reqObj)
|
||||
if err != nil {
|
||||
return types.NewScriptExecutionError("JavaScript", err)
|
||||
return types.NewScriptExecutionError("JavaScript", jsExceptionError(err))
|
||||
}
|
||||
|
||||
// Update RequestData from the returned object
|
||||
@@ -79,6 +96,8 @@ func (e *JsEngine) Close() {
|
||||
// goja doesn't have an explicit close method, but we can help GC
|
||||
e.runtime = nil
|
||||
e.transform = nil
|
||||
e.jsonParse = nil
|
||||
e.errorConstructor = nil
|
||||
}
|
||||
|
||||
// requestDataToObject converts RequestData to a goja Value (JavaScript object).
|
||||
@@ -190,9 +209,18 @@ func (e *JsEngine) objectToStringSliceMap(obj *goja.Object) map[string][]string
|
||||
}
|
||||
result[key] = values
|
||||
} else {
|
||||
// Single value - wrap in slice
|
||||
// Single value, wrap it in a slice
|
||||
result[key] = []string{v.String()}
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// jsExceptionError keeps the thrown value's message without goja's stack frame.
|
||||
func jsExceptionError(err error) error {
|
||||
var exception *goja.Exception
|
||||
if !errors.As(err, &exception) {
|
||||
return err
|
||||
}
|
||||
return errors.New(exception.Value().String())
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
package script
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"github.com/dop251/goja"
|
||||
"go.aykhans.me/sarin/internal/types"
|
||||
)
|
||||
|
||||
// registerHTTP defines the global http function.
|
||||
func (e *JsEngine) registerHTTP() error {
|
||||
rt := e.runtime
|
||||
|
||||
jsonParse, ok := goja.AssertFunction(rt.Get("JSON").ToObject(rt).Get("parse"))
|
||||
if !ok {
|
||||
return errors.New("JSON.parse is not available")
|
||||
}
|
||||
e.jsonParse = jsonParse
|
||||
e.errorConstructor = rt.Get("Error")
|
||||
|
||||
return rt.Set("http", func(call goja.FunctionCall) goja.Value {
|
||||
return e.httpCall(call.Argument(0), call.Argument(1))
|
||||
})
|
||||
}
|
||||
|
||||
func (e *JsEngine) httpCall(urlValue, optsValue goja.Value) goja.Value {
|
||||
rt := e.runtime
|
||||
|
||||
url, ok := urlValue.Export().(string)
|
||||
if !ok {
|
||||
panic(rt.NewTypeError("http: url must be a string, got %s", jsTypeName(urlValue)))
|
||||
}
|
||||
req := &HTTPRequest{Method: httpDefaultMethod, URL: url}
|
||||
|
||||
if !goja.IsUndefined(optsValue) && !goja.IsNull(optsValue) {
|
||||
opts, ok := optsValue.(*goja.Object)
|
||||
if !ok {
|
||||
panic(rt.NewTypeError("http: options must be an object, got %s", jsTypeName(optsValue)))
|
||||
}
|
||||
if err := e.parseHTTPOptions(opts, req); err != nil {
|
||||
e.throw(err)
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := e.http.do(req)
|
||||
if err != nil {
|
||||
e.throw(err)
|
||||
}
|
||||
|
||||
return e.httpResponseToObject(resp)
|
||||
}
|
||||
|
||||
// parseHTTPOptions fills req from a JavaScript options object.
|
||||
// It can return the following errors:
|
||||
// - types.ScriptHTTPOptionError
|
||||
func (e *JsEngine) parseHTTPOptions(opts *goja.Object, req *HTTPRequest) error {
|
||||
for _, name := range opts.Keys() {
|
||||
value := opts.Get(name)
|
||||
if goja.IsUndefined(value) {
|
||||
continue
|
||||
}
|
||||
|
||||
switch name {
|
||||
case httpOptionMethod:
|
||||
method, ok := value.Export().(string)
|
||||
if !ok {
|
||||
return types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("string", jsTypeName(value)))
|
||||
}
|
||||
|
||||
req.Method = method
|
||||
case httpOptionHeaders, httpOptionParams, httpOptionCookies:
|
||||
object, ok := value.(*goja.Object)
|
||||
if !ok {
|
||||
return types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("object", jsTypeName(value)))
|
||||
}
|
||||
|
||||
values := e.objectToStringSliceMap(object)
|
||||
switch name {
|
||||
case httpOptionHeaders:
|
||||
req.Headers = values
|
||||
case httpOptionParams:
|
||||
req.Params = values
|
||||
default:
|
||||
req.Cookies = values
|
||||
}
|
||||
case httpOptionBody:
|
||||
body, ok := value.Export().(string)
|
||||
if !ok {
|
||||
return types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("string", jsTypeName(value)))
|
||||
}
|
||||
|
||||
req.Body = body
|
||||
case httpOptionTimeout:
|
||||
timeout, ok := value.Export().(string)
|
||||
if !ok {
|
||||
return types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("duration string", jsTypeName(value)))
|
||||
}
|
||||
|
||||
var err error
|
||||
if req.Timeout, err = parseHTTPTimeout(timeout); err != nil {
|
||||
return err
|
||||
}
|
||||
case httpOptionInsecure, httpOptionFollowRedirects:
|
||||
flag, ok := value.Export().(bool)
|
||||
if !ok {
|
||||
return types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("boolean", jsTypeName(value)))
|
||||
}
|
||||
|
||||
if name == httpOptionInsecure {
|
||||
req.Insecure = flag
|
||||
} else {
|
||||
req.FollowRedirects = flag
|
||||
}
|
||||
default:
|
||||
return types.NewScriptHTTPOptionError(name, types.ErrScriptHTTPUnknownOption)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// httpResponseToObject converts a response to a JavaScript object.
|
||||
func (e *JsEngine) httpResponseToObject(resp *HTTPResponse) *goja.Object {
|
||||
rt := e.runtime
|
||||
obj := rt.NewObject()
|
||||
|
||||
headers := rt.NewObject()
|
||||
for key, values := range resp.Headers {
|
||||
_ = headers.Set(key, e.stringSliceToArray(values))
|
||||
}
|
||||
|
||||
lookup := func(find func(string) (string, bool)) func(goja.FunctionCall) goja.Value {
|
||||
return func(call goja.FunctionCall) goja.Value {
|
||||
name, ok := call.Argument(0).Export().(string)
|
||||
if !ok {
|
||||
panic(rt.NewTypeError("name must be a string, got %s", jsTypeName(call.Argument(0))))
|
||||
}
|
||||
if value, ok := find(name); ok {
|
||||
return rt.ToValue(value)
|
||||
}
|
||||
return goja.Null()
|
||||
}
|
||||
}
|
||||
|
||||
_ = obj.Set("status", resp.Status)
|
||||
_ = obj.Set("headers", headers)
|
||||
_ = obj.Set("body", resp.Body)
|
||||
_ = obj.Set("header", lookup(resp.Header))
|
||||
_ = obj.Set("cookie", lookup(resp.Cookie))
|
||||
_ = obj.Set("json", func(goja.FunctionCall) goja.Value {
|
||||
value, err := e.jsonParse(goja.Undefined(), rt.ToValue(resp.Body))
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return value
|
||||
})
|
||||
|
||||
return obj
|
||||
}
|
||||
|
||||
// jsTypeName returns a short type name for error messages.
|
||||
func jsTypeName(value goja.Value) string {
|
||||
switch {
|
||||
case value == nil || goja.IsUndefined(value):
|
||||
return "undefined"
|
||||
case goja.IsNull(value):
|
||||
return "null"
|
||||
}
|
||||
switch value.Export().(type) {
|
||||
case string:
|
||||
return "string"
|
||||
case bool:
|
||||
return "boolean"
|
||||
case int64, float64:
|
||||
return "number"
|
||||
case []any:
|
||||
return "array"
|
||||
default:
|
||||
if _, ok := goja.AssertFunction(value); ok {
|
||||
return "function"
|
||||
}
|
||||
return "object"
|
||||
}
|
||||
}
|
||||
|
||||
// throw raises err as a plain JavaScript Error instead of a GoError.
|
||||
func (e *JsEngine) throw(err error) {
|
||||
errorValue, newErr := e.runtime.New(e.errorConstructor, e.runtime.ToValue(err.Error()))
|
||||
if newErr != nil {
|
||||
panic(e.runtime.NewGoError(err))
|
||||
}
|
||||
panic(errorValue)
|
||||
}
|
||||
+29
-33
@@ -11,6 +11,7 @@ import (
|
||||
type LuaEngine struct {
|
||||
state *lua.LState
|
||||
transform *lua.LFunction
|
||||
http httpBridge
|
||||
}
|
||||
|
||||
// NewLuaEngine creates a new Lua script engine with the given script content.
|
||||
@@ -29,6 +30,11 @@ type LuaEngine struct {
|
||||
// - types.ScriptExecutionError
|
||||
func NewLuaEngine(scriptContent string) (*LuaEngine, error) {
|
||||
L := lua.NewState()
|
||||
engine := &LuaEngine{state: L}
|
||||
|
||||
// Register the globals before running the script so its functions can use them
|
||||
engine.registerHTTP()
|
||||
registerLuaJSON(L)
|
||||
|
||||
// Execute the script to define the transform function
|
||||
if err := L.DoString(scriptContent); err != nil {
|
||||
@@ -42,17 +48,23 @@ func NewLuaEngine(scriptContent string) (*LuaEngine, error) {
|
||||
L.Close()
|
||||
return nil, types.ErrScriptTransformMissing
|
||||
}
|
||||
engine.transform = transform.(*lua.LFunction)
|
||||
|
||||
return &LuaEngine{
|
||||
state: L,
|
||||
transform: transform.(*lua.LFunction),
|
||||
}, nil
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
// SetHTTPDoer sets the doer that the script's http.* calls are sent through.
|
||||
func (e *LuaEngine) SetHTTPDoer(doer HTTPDoer) {
|
||||
e.http.doer = doer
|
||||
}
|
||||
|
||||
// Transform executes the Lua transform function with the given request data.
|
||||
// It can return the following errors:
|
||||
// - types.ScriptExecutionError
|
||||
func (e *LuaEngine) Transform(req *RequestData) error {
|
||||
e.http.active = true
|
||||
defer func() { e.http.active = false }()
|
||||
|
||||
// Convert RequestData to Lua table
|
||||
reqTable := e.requestDataToTable(req)
|
||||
|
||||
@@ -93,39 +105,23 @@ func (e *LuaEngine) requestDataToTable(req *RequestData) *lua.LTable {
|
||||
t.RawSetString("path", lua.LString(req.Path))
|
||||
t.RawSetString("body", lua.LString(req.Body))
|
||||
|
||||
// Headers (map[string][]string -> table of arrays)
|
||||
headers := L.NewTable()
|
||||
for k, values := range req.Headers {
|
||||
arr := L.NewTable()
|
||||
t.RawSetString("headers", stringSliceMapToTable(L, req.Headers))
|
||||
t.RawSetString("params", stringSliceMapToTable(L, req.Params))
|
||||
t.RawSetString("cookies", stringSliceMapToTable(L, req.Cookies))
|
||||
|
||||
return t
|
||||
}
|
||||
|
||||
// stringSliceMapToTable converts a Go map[string][]string to a Lua table of arrays.
|
||||
func stringSliceMapToTable(state *lua.LState, m map[string][]string) *lua.LTable {
|
||||
t := state.NewTable()
|
||||
for k, values := range m {
|
||||
arr := state.NewTable()
|
||||
for _, v := range values {
|
||||
arr.Append(lua.LString(v))
|
||||
}
|
||||
headers.RawSetString(k, arr)
|
||||
t.RawSetString(k, arr)
|
||||
}
|
||||
t.RawSetString("headers", headers)
|
||||
|
||||
// Params (map[string][]string -> table of arrays)
|
||||
params := L.NewTable()
|
||||
for k, values := range req.Params {
|
||||
arr := L.NewTable()
|
||||
for _, v := range values {
|
||||
arr.Append(lua.LString(v))
|
||||
}
|
||||
params.RawSetString(k, arr)
|
||||
}
|
||||
t.RawSetString("params", params)
|
||||
|
||||
// Cookies (map[string][]string -> table of arrays)
|
||||
cookies := L.NewTable()
|
||||
for k, values := range req.Cookies {
|
||||
arr := L.NewTable()
|
||||
for _, v := range values {
|
||||
arr.Append(lua.LString(v))
|
||||
}
|
||||
cookies.RawSetString(k, arr)
|
||||
}
|
||||
t.RawSetString("cookies", cookies)
|
||||
|
||||
return t
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
package script
|
||||
|
||||
import (
|
||||
lua "github.com/yuin/gopher-lua"
|
||||
"go.aykhans.me/sarin/internal/types"
|
||||
)
|
||||
|
||||
// registerHTTP defines the global http function.
|
||||
func (e *LuaEngine) registerHTTP() {
|
||||
L := e.state
|
||||
L.SetGlobal("http", L.NewFunction(func(L *lua.LState) int {
|
||||
return e.httpCall(L, 1)
|
||||
}))
|
||||
}
|
||||
|
||||
// httpCall sends the request whose URL is at stack position urlIndex and pushes the response.
|
||||
func (e *LuaEngine) httpCall(state *lua.LState, urlIndex int) int {
|
||||
req := &HTTPRequest{
|
||||
Method: httpDefaultMethod,
|
||||
URL: state.CheckString(urlIndex),
|
||||
}
|
||||
|
||||
switch opts := state.Get(urlIndex + 1).(type) {
|
||||
case *lua.LNilType:
|
||||
case *lua.LTable:
|
||||
if err := e.parseHTTPOptions(opts, req); err != nil {
|
||||
state.RaiseError("%s", err.Error())
|
||||
}
|
||||
default:
|
||||
state.ArgError(urlIndex+1, "options must be a table")
|
||||
}
|
||||
|
||||
resp, err := e.http.do(req)
|
||||
if err != nil {
|
||||
state.RaiseError("%s", err.Error())
|
||||
}
|
||||
|
||||
state.Push(httpResponseToTable(state, resp))
|
||||
return 1
|
||||
}
|
||||
|
||||
// parseHTTPOptions fills req from a Lua options table.
|
||||
// It can return the following errors:
|
||||
// - types.ScriptHTTPOptionError
|
||||
func (e *LuaEngine) parseHTTPOptions(opts *lua.LTable, req *HTTPRequest) error {
|
||||
var err error
|
||||
opts.ForEach(func(k, v lua.LValue) {
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
key, ok := k.(lua.LString)
|
||||
if !ok {
|
||||
err = types.NewScriptHTTPOptionError(k.String(), types.ErrScriptHTTPUnknownOption)
|
||||
return
|
||||
}
|
||||
|
||||
switch name := string(key); name {
|
||||
case httpOptionMethod:
|
||||
method, ok := v.(lua.LString)
|
||||
if !ok {
|
||||
err = types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("string", v.Type().String()))
|
||||
return
|
||||
}
|
||||
req.Method = string(method)
|
||||
case httpOptionHeaders, httpOptionParams, httpOptionCookies:
|
||||
table, ok := v.(*lua.LTable)
|
||||
if !ok {
|
||||
err = types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("table", v.Type().String()))
|
||||
return
|
||||
}
|
||||
values := e.tableToStringSliceMap(table)
|
||||
switch name {
|
||||
case httpOptionHeaders:
|
||||
req.Headers = values
|
||||
case httpOptionParams:
|
||||
req.Params = values
|
||||
default:
|
||||
req.Cookies = values
|
||||
}
|
||||
case httpOptionBody:
|
||||
body, ok := v.(lua.LString)
|
||||
if !ok {
|
||||
err = types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("string", v.Type().String()))
|
||||
return
|
||||
}
|
||||
req.Body = string(body)
|
||||
case httpOptionTimeout:
|
||||
timeout, ok := v.(lua.LString)
|
||||
if !ok {
|
||||
err = types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("duration string", v.Type().String()))
|
||||
return
|
||||
}
|
||||
req.Timeout, err = parseHTTPTimeout(string(timeout))
|
||||
case httpOptionInsecure, httpOptionFollowRedirects:
|
||||
flag, ok := v.(lua.LBool)
|
||||
if !ok {
|
||||
err = types.NewScriptHTTPOptionError(name, types.NewScriptTypeError("boolean", v.Type().String()))
|
||||
return
|
||||
}
|
||||
if name == httpOptionInsecure {
|
||||
req.Insecure = bool(flag)
|
||||
} else {
|
||||
req.FollowRedirects = bool(flag)
|
||||
}
|
||||
default:
|
||||
err = types.NewScriptHTTPOptionError(name, types.ErrScriptHTTPUnknownOption)
|
||||
}
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// httpResponseToTable converts a response to a Lua table.
|
||||
func httpResponseToTable(state *lua.LState, resp *HTTPResponse) *lua.LTable {
|
||||
t := state.NewTable()
|
||||
t.RawSetString("status", lua.LNumber(resp.Status))
|
||||
t.RawSetString("headers", stringSliceMapToTable(state, resp.Headers))
|
||||
t.RawSetString("body", lua.LString(resp.Body))
|
||||
|
||||
lookup := func(find func(string) (string, bool)) *lua.LFunction {
|
||||
return state.NewFunction(func(L *lua.LState) int {
|
||||
// Last argument, so both res:header() and res.header() work.
|
||||
top := L.GetTop()
|
||||
if top == 0 {
|
||||
L.ArgError(1, "name expected")
|
||||
}
|
||||
if value, ok := find(L.CheckString(top)); ok {
|
||||
L.Push(lua.LString(value))
|
||||
} else {
|
||||
L.Push(lua.LNil)
|
||||
}
|
||||
return 1
|
||||
})
|
||||
}
|
||||
t.RawSetString("header", lookup(resp.Header))
|
||||
t.RawSetString("cookie", lookup(resp.Cookie))
|
||||
|
||||
return t
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
package script
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"math"
|
||||
"strings"
|
||||
|
||||
lua "github.com/yuin/gopher-lua"
|
||||
"go.aykhans.me/sarin/internal/types"
|
||||
)
|
||||
|
||||
// registerLuaJSON defines the global json table with encode and decode.
|
||||
func registerLuaJSON(state *lua.LState) {
|
||||
module := state.NewTable()
|
||||
module.RawSetString("encode", state.NewFunction(luaJSONEncode))
|
||||
module.RawSetString("decode", state.NewFunction(luaJSONDecode))
|
||||
state.SetGlobal("json", module)
|
||||
}
|
||||
|
||||
func luaJSONEncode(state *lua.LState) int {
|
||||
value, err := luaToGo(state.CheckAny(1), make(map[*lua.LTable]bool))
|
||||
if err != nil {
|
||||
state.RaiseError("%s", types.NewScriptJSONError("encode", err).Error())
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
encoder := json.NewEncoder(&buf)
|
||||
encoder.SetEscapeHTML(false)
|
||||
if err := encoder.Encode(value); err != nil {
|
||||
state.RaiseError("%s", types.NewScriptJSONError("encode", err).Error())
|
||||
}
|
||||
|
||||
state.Push(lua.LString(strings.TrimSuffix(buf.String(), "\n")))
|
||||
return 1
|
||||
}
|
||||
|
||||
func luaJSONDecode(state *lua.LState) int {
|
||||
var value any
|
||||
if err := json.Unmarshal([]byte(state.CheckString(1)), &value); err != nil {
|
||||
state.RaiseError("%s", types.NewScriptJSONError("decode", err).Error())
|
||||
}
|
||||
|
||||
state.Push(goToLua(state, value))
|
||||
return 1
|
||||
}
|
||||
|
||||
// luaToGo converts a Lua value to a JSON-encodable Go value.
|
||||
// It can return the following errors:
|
||||
// - types.ErrScriptJSONCycle
|
||||
// - types.ScriptTypeError
|
||||
func luaToGo(value lua.LValue, visiting map[*lua.LTable]bool) (any, error) {
|
||||
switch v := value.(type) {
|
||||
case *lua.LNilType:
|
||||
return nil, nil
|
||||
case lua.LBool:
|
||||
return bool(v), nil
|
||||
case lua.LNumber:
|
||||
return float64(v), nil
|
||||
case lua.LString:
|
||||
return string(v), nil
|
||||
case *lua.LTable:
|
||||
if visiting[v] {
|
||||
return nil, types.ErrScriptJSONCycle
|
||||
}
|
||||
visiting[v] = true
|
||||
defer delete(visiting, v)
|
||||
|
||||
if length, ok := luaArrayLength(v); ok {
|
||||
array := make([]any, length)
|
||||
for i := range length {
|
||||
item, err := luaToGo(v.RawGetInt(i+1), visiting)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
array[i] = item
|
||||
}
|
||||
return array, nil
|
||||
}
|
||||
|
||||
object := make(map[string]any)
|
||||
var err error
|
||||
v.ForEach(func(key, item lua.LValue) {
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
switch key.(type) {
|
||||
case lua.LString, lua.LNumber:
|
||||
default:
|
||||
err = types.NewScriptTypeError("string or number object key", key.Type().String())
|
||||
return
|
||||
}
|
||||
object[key.String()], err = luaToGo(item, visiting)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return object, nil
|
||||
default:
|
||||
return nil, types.NewScriptTypeError("nil, boolean, number, string or table", value.Type().String())
|
||||
}
|
||||
}
|
||||
|
||||
// luaArrayLength reports whether t is a non-empty sequence with keys exactly 1..n.
|
||||
func luaArrayLength(t *lua.LTable) (int, bool) {
|
||||
count := 0
|
||||
isArray := true
|
||||
maxKey := 0.0
|
||||
t.ForEach(func(key, _ lua.LValue) {
|
||||
count++
|
||||
number, ok := key.(lua.LNumber)
|
||||
if !ok || float64(number) < 1 || float64(number) != math.Trunc(float64(number)) {
|
||||
isArray = false
|
||||
return
|
||||
}
|
||||
maxKey = max(maxKey, float64(number))
|
||||
})
|
||||
if !isArray || count == 0 || maxKey != float64(count) {
|
||||
return 0, false
|
||||
}
|
||||
return count, true
|
||||
}
|
||||
|
||||
// goToLua converts a value produced by encoding/json to a Lua value.
|
||||
func goToLua(state *lua.LState, value any) lua.LValue {
|
||||
switch v := value.(type) {
|
||||
case bool:
|
||||
return lua.LBool(v)
|
||||
case float64:
|
||||
return lua.LNumber(v)
|
||||
case string:
|
||||
return lua.LString(v)
|
||||
case []any:
|
||||
t := state.CreateTable(len(v), 0)
|
||||
for i, item := range v {
|
||||
t.RawSetInt(i+1, goToLua(state, item))
|
||||
}
|
||||
return t
|
||||
case map[string]any:
|
||||
t := state.CreateTable(0, len(v))
|
||||
for key, item := range v {
|
||||
t.RawSetString(key, goToLua(state, item))
|
||||
}
|
||||
return t
|
||||
default:
|
||||
return lua.LNil
|
||||
}
|
||||
}
|
||||
@@ -31,6 +31,9 @@ type Engine interface {
|
||||
// The script should modify the RequestData and return it.
|
||||
Transform(req *RequestData) error
|
||||
|
||||
// SetHTTPDoer sets the doer that the script's http.* calls are sent through.
|
||||
SetHTTPDoer(doer HTTPDoer)
|
||||
|
||||
// Close releases any resources held by the engine.
|
||||
Close()
|
||||
}
|
||||
@@ -69,7 +72,7 @@ func LoadSource(ctx context.Context, source string, engineType EngineType) (*Sou
|
||||
|
||||
switch {
|
||||
case strings.HasPrefix(source, "@@"):
|
||||
// Escaped @ - it's an inline script starting with literal @
|
||||
// Escaped @: an inline script starting with a literal @
|
||||
content = source[1:] // Remove first @, keep the rest
|
||||
case strings.HasPrefix(source, "@"):
|
||||
// File or URL reference
|
||||
@@ -125,7 +128,7 @@ func ValidateScript(ctx context.Context, source string, engineType EngineType) e
|
||||
return err
|
||||
}
|
||||
|
||||
// Try to create an engine - this validates syntax and transform function
|
||||
// Creating an engine validates the syntax and the transform function
|
||||
var engine Engine
|
||||
switch engineType {
|
||||
case EngineTypeLua:
|
||||
@@ -140,7 +143,7 @@ func ValidateScript(ctx context.Context, source string, engineType EngineType) e
|
||||
return err
|
||||
}
|
||||
|
||||
// Clean up the engine - we only needed it for validation
|
||||
// The engine was only needed for validation
|
||||
engine.Close()
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -401,8 +401,88 @@ var (
|
||||
ErrScriptTransformMissing = errors.New("script must define a global 'transform' function")
|
||||
ErrScriptTransformReturnObject = errors.New("transform function must return an object")
|
||||
ErrScriptURLNoHost = errors.New("script URL must have a host")
|
||||
ErrScriptHTTPOutsideTransform = errors.New("http functions can only be called while transform is running")
|
||||
ErrScriptHTTPUnavailable = errors.New("http client is not available")
|
||||
ErrScriptHTTPMethodEmpty = errors.New("http method cannot be empty")
|
||||
ErrScriptHTTPURLInvalid = errors.New("URL must be an absolute http or https URL")
|
||||
ErrScriptHTTPUnknownOption = errors.New("unknown option")
|
||||
ErrScriptJSONCycle = errors.New("cannot encode a table that contains itself")
|
||||
)
|
||||
|
||||
type ScriptTypeError struct {
|
||||
Expected string
|
||||
Got string
|
||||
}
|
||||
|
||||
func NewScriptTypeError(expected, got string) ScriptTypeError {
|
||||
return ScriptTypeError{expected, got}
|
||||
}
|
||||
|
||||
func (e ScriptTypeError) Error() string {
|
||||
return fmt.Sprintf("expected %s, got %s", e.Expected, e.Got)
|
||||
}
|
||||
|
||||
type ScriptHTTPOptionError struct {
|
||||
Option string
|
||||
Err error
|
||||
}
|
||||
|
||||
func NewScriptHTTPOptionError(option string, err error) ScriptHTTPOptionError {
|
||||
if err == nil {
|
||||
err = errNoError
|
||||
}
|
||||
return ScriptHTTPOptionError{option, err}
|
||||
}
|
||||
|
||||
func (e ScriptHTTPOptionError) Error() string {
|
||||
return fmt.Sprintf("http option %q: %v", e.Option, e.Err)
|
||||
}
|
||||
|
||||
func (e ScriptHTTPOptionError) Unwrap() error {
|
||||
return e.Err
|
||||
}
|
||||
|
||||
type ScriptHTTPRequestError struct {
|
||||
Method string
|
||||
URL string
|
||||
Err error
|
||||
}
|
||||
|
||||
func NewScriptHTTPRequestError(method, url string, err error) ScriptHTTPRequestError {
|
||||
if err == nil {
|
||||
err = errNoError
|
||||
}
|
||||
return ScriptHTTPRequestError{method, url, err}
|
||||
}
|
||||
|
||||
func (e ScriptHTTPRequestError) Error() string {
|
||||
return fmt.Sprintf("http %s %s: %v", e.Method, e.URL, e.Err)
|
||||
}
|
||||
|
||||
func (e ScriptHTTPRequestError) Unwrap() error {
|
||||
return e.Err
|
||||
}
|
||||
|
||||
type ScriptJSONError struct {
|
||||
Op string
|
||||
Err error
|
||||
}
|
||||
|
||||
func NewScriptJSONError(op string, err error) ScriptJSONError {
|
||||
if err == nil {
|
||||
err = errNoError
|
||||
}
|
||||
return ScriptJSONError{op, err}
|
||||
}
|
||||
|
||||
func (e ScriptJSONError) Error() string {
|
||||
return fmt.Sprintf("json.%s: %v", e.Op, e.Err)
|
||||
}
|
||||
|
||||
func (e ScriptJSONError) Unwrap() error {
|
||||
return e.Err
|
||||
}
|
||||
|
||||
type ScriptLoadError struct {
|
||||
Source string
|
||||
Err error
|
||||
|
||||
Reference in new issue
Block a user