diff --git a/cmd/cli/main.go b/cmd/cli/main.go index c492843..4b05082 100644 --- a/cmd/cli/main.go +++ b/cmd/cli/main.go @@ -75,7 +75,7 @@ func main() { }), ) - srn.Start(ctx, stopCtrl) + srn.Start(stopCtrl) switch *combinedConfig.Output { case config.ConfigOutputTypeNone: diff --git a/internal/sarin/runner.go b/internal/sarin/runner.go index 921d306..cfaafd6 100644 --- a/internal/sarin/runner.go +++ b/internal/sarin/runner.go @@ -199,6 +199,10 @@ type sarin struct { logError bool logFile string + // jobsCtx ends the run, either when the duration expires or on Ctrl+C. + jobsCtx context.Context //nolint:containedctx + jobsCancel context.CancelFunc + hostClients []*fasthttp.HostClient scriptHTTPClients []*scriptHTTPClient responses *SarinResponseData @@ -272,15 +276,20 @@ func NewSarin( scriptChain := script.NewChain(luaSources, jsSources) + jobsCtx, jobsCancel := context.WithCancel(ctx) + var scriptHTTPClients []*scriptHTTPClient if !scriptChain.IsEmpty() { - scriptHTTPClients, err = newScriptHTTPClients(ctx, proxyURLs, workers) + scriptHTTPClients, err = newScriptHTTPClients(jobsCtx, proxyURLs, workers) if err != nil { + jobsCancel() return nil, err } } srn := &sarin{ + jobsCtx: jobsCtx, + jobsCancel: jobsCancel, workers: workers, requestURL: requestURL, methods: methods, @@ -316,9 +325,7 @@ func (s sarin) GetResponses() *SarinResponseData { return s.responses } -func (s sarin) Start(ctx context.Context, stopCtrl *StopController) { - jobsCtx, jobsCancel := context.WithCancel(ctx) - +func (s sarin) Start(stopCtrl *StopController) { var workersWG sync.WaitGroup jobsCh := make(chan struct{}, max(s.workers, 1)) @@ -382,15 +389,14 @@ func (s sarin) Start(ctx context.Context, stopCtrl *StopController) { 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 go s.streamProgress(streamCtx, stopCtrl, streamCh, totalRequests, &counter, tuiLogChannel, showProgressBar) } // Setup duration-based cancellation - s.setupDurationTimeout(ctx, jobsCancel) + s.setupDurationTimeout() // Distribute jobs to workers. // This blocks until all jobs are sent or the context is canceled. - s.sendJobs(jobsCtx, jobsCh) + s.sendJobs(jobsCh) // Close the jobs channel so workers stop after completing their current job close(jobsCh) @@ -455,31 +461,31 @@ func (s sarin) startWorkers(wg *sync.WaitGroup, jobs <-chan struct{}, counter *a } } -func (s sarin) setupDurationTimeout(ctx context.Context, cancel context.CancelFunc) { +func (s sarin) setupDurationTimeout() { if s.totalDuration != nil { go func() { timer := time.NewTimer(*s.totalDuration) defer timer.Stop() select { case <-timer.C: - cancel() - case <-ctx.Done(): + s.jobsCancel() + case <-s.jobsCtx.Done(): // Context cancelled, cleanup } }() } } -func (s sarin) sendJobs(ctx context.Context, jobs chan<- struct{}) { +func (s sarin) sendJobs(jobs chan<- struct{}) { if s.totalRequests != nil && *s.totalRequests > 0 { for range *s.totalRequests { - if ctx.Err() != nil { + if s.jobsCtx.Err() != nil { break } jobs <- struct{}{} } } else { - for ctx.Err() == nil { + for s.jobsCtx.Err() == nil { jobs <- struct{}{} } } diff --git a/internal/sarin/scripthttp.go b/internal/sarin/scripthttp.go index 2143d4d..de1c567 100644 --- a/internal/sarin/scripthttp.go +++ b/internal/sarin/scripthttp.go @@ -17,6 +17,7 @@ const scriptHTTPDefaultTimeout = 30 * time.Second // scriptHTTPClient sends scripts' http.* requests. It is safe for concurrent use. type scriptHTTPClient struct { + ctx context.Context //nolint:containedctx client *fasthttp.Client insecureClient *fasthttp.Client defaultTimeout time.Duration @@ -29,7 +30,7 @@ var _ script.HTTPDoer = (*scriptHTTPClient)(nil) // - types.ProxyDialError func newScriptHTTPClients(ctx context.Context, proxies []url.URL, maxConns uint) ([]*scriptHTTPClient, error) { if len(proxies) == 0 { - return []*scriptHTTPClient{newScriptHTTPClient(fasthttp.DialDualStackTimeout, scriptHTTPDefaultTimeout, maxConns)}, nil + return []*scriptHTTPClient{newScriptHTTPClient(ctx, fasthttp.DialDualStackTimeout, scriptHTTPDefaultTimeout, maxConns)}, nil } clients := make([]*scriptHTTPClient, 0, len(proxies)) @@ -38,12 +39,12 @@ func newScriptHTTPClients(ctx context.Context, proxies []url.URL, maxConns uint) if err != nil { return nil, types.NewProxyDialError(proxy.String(), err) } - clients = append(clients, newScriptHTTPClient(dial, scriptHTTPDefaultTimeout, maxConns)) + clients = append(clients, newScriptHTTPClient(ctx, dial, scriptHTTPDefaultTimeout, maxConns)) } return clients, nil } -func newScriptHTTPClient(dial fasthttp.DialFuncWithTimeout, defaultTimeout time.Duration, maxConns uint) *scriptHTTPClient { +func newScriptHTTPClient(ctx context.Context, dial fasthttp.DialFuncWithTimeout, defaultTimeout time.Duration, maxConns uint) *scriptHTTPClient { newClient := func(insecure bool) *fasthttp.Client { // Read and write timeouts stay unset because they would cap the per-call timeout, // so dialWithDeadline is what bounds the TLS handshake. @@ -60,6 +61,7 @@ func newScriptHTTPClient(dial fasthttp.DialFuncWithTimeout, defaultTimeout time. } return &scriptHTTPClient{ + ctx: ctx, client: newClient(false), insecureClient: newClient(true), defaultTimeout: defaultTimeout, @@ -91,6 +93,10 @@ func dialWithDeadline(dial fasthttp.DialFuncWithTimeout) fasthttp.DialFuncWithTi // It can return the following errors: // - types.ScriptHTTPRequestError func (c *scriptHTTPClient) Do(r *script.HTTPRequest) (*script.HTTPResponse, error) { + if err := c.ctx.Err(); err != nil { + return nil, types.NewScriptHTTPRequestError(r.Method, r.URL, err) + } + parsedURL, err := url.Parse(r.URL) if err != nil { return nil, types.NewScriptHTTPRequestError(r.Method, r.URL, err)