From 2eebac68c9fb8099eedc0d1f2f3e8ccbcd273177 Mon Sep 17 00:00:00 2001 From: Aykhan Shahsuvarov Date: Fri, 17 Apr 2026 16:03:03 +0400 Subject: [PATCH] feat: add double Ctrl+C hard-kill and split TUI into tui.go First signal cancels the main context; second releases the terminal and exits 130, unblocking shutdown when captcha polls or user scripts ignore context. Also moves the bubbletea models and streamProgress out of sarin.go into a new tui.go. --- cmd/cli/main.go | 14 +- internal/sarin/sarin.go | 292 +------------------------------------- internal/sarin/stop.go | 55 ++++++++ internal/sarin/tui.go | 300 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 365 insertions(+), 296 deletions(-) create mode 100644 internal/sarin/stop.go create mode 100644 internal/sarin/tui.go diff --git a/cmd/cli/main.go b/cmd/cli/main.go index 2ca978c..0530246 100644 --- a/cmd/cli/main.go +++ b/cmd/cli/main.go @@ -15,7 +15,8 @@ import ( func main() { ctx, cancel := context.WithCancel(context.Background()) - go listenForTermination(func() { cancel() }) + stopCtrl := sarin.NewStopController(cancel) + go listenForTermination(stopCtrl.Stop) combinedConfig := config.ReadAllConfigs() @@ -73,7 +74,7 @@ func main() { }), ) - srn.Start(ctx) + srn.Start(ctx, stopCtrl) switch *combinedConfig.Output { case config.ConfigOutputTypeNone: @@ -87,9 +88,10 @@ func main() { } } -func listenForTermination(do func()) { - sigChan := make(chan os.Signal, 1) +func listenForTermination(stop func()) { + sigChan := make(chan os.Signal, 4) signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) - <-sigChan - do() + for range sigChan { + stop() + } } diff --git a/internal/sarin/sarin.go b/internal/sarin/sarin.go index 176b10e..e068915 100644 --- a/internal/sarin/sarin.go +++ b/internal/sarin/sarin.go @@ -5,15 +5,10 @@ import ( "net/url" "os" "strconv" - "strings" "sync" "sync/atomic" "time" - "github.com/charmbracelet/bubbles/progress" - "github.com/charmbracelet/bubbles/spinner" - tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/lipgloss" "github.com/charmbracelet/x/term" "github.com/valyala/fasthttp" "go.aykhans.me/sarin/internal/script" @@ -138,7 +133,7 @@ func (q sarin) GetResponses() *SarinResponseData { return q.responses } -func (q sarin) Start(ctx context.Context) { +func (q sarin) Start(ctx context.Context, stopCtrl *StopController) { jobsCtx, jobsCancel := context.WithCancel(ctx) var workersWG sync.WaitGroup @@ -183,7 +178,7 @@ func (q sarin) Start(ctx context.Context) { if !q.quiet { // Start streaming to terminal //nolint:contextcheck // streamCtx must remain active until all workers complete to ensure all collected data is streamed - go q.streamProgress(streamCtx, jobsCancel, streamCh, totalRequests, &counter, messageChannel) + go q.streamProgress(streamCtx, stopCtrl, streamCh, totalRequests, &counter, messageChannel) } // Setup duration-based cancellation @@ -531,286 +526,3 @@ func (q sarin) sendJobs(ctx context.Context, jobs chan<- struct{}) { } } } - -type tickMsg time.Time - -var ( - helpStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#d1d1d1")) - errorStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#FC5B5B")).Bold(true) - warningStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#FFD93D")).Bold(true) - messageChannelStyle = lipgloss.NewStyle(). - Border(lipgloss.ThickBorder(), false, false, false, true). - BorderForeground(lipgloss.Color("#757575")). - PaddingLeft(1). - Margin(1, 0, 0, 0). - Foreground(lipgloss.Color("#888888")) -) - -type progressModel struct { - progress progress.Model - startTime time.Time - messages []string - counter *atomic.Uint64 - current uint64 - maxValue uint64 - ctx context.Context //nolint:containedctx - cancel context.CancelFunc - cancelling bool -} - -func (m progressModel) Init() tea.Cmd { - return tea.Batch(progressTickCmd()) -} - -func (m progressModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { - switch msg := msg.(type) { - case tea.KeyMsg: - if msg.Type == tea.KeyCtrlC { - m.cancelling = true - m.cancel() - } - return m, nil - - case tea.WindowSizeMsg: - m.progress.Width = max(10, msg.Width-1) - if m.ctx.Err() != nil { - return m, tea.Quit - } - return m, nil - - case runtimeMessage: - var msgBuilder strings.Builder - msgBuilder.WriteString("[") - msgBuilder.WriteString(msg.timestamp.Format("15:04:05")) - msgBuilder.WriteString("] ") - switch msg.level { - case runtimeMessageLevelError: - msgBuilder.WriteString(errorStyle.Render("ERROR: ")) - case runtimeMessageLevelWarning: - msgBuilder.WriteString(warningStyle.Render("WARNING: ")) - } - msgBuilder.WriteString(msg.text) - m.messages = append(m.messages[1:], msgBuilder.String()) - if m.ctx.Err() != nil { - return m, tea.Quit - } - return m, nil - - case tickMsg: - if m.ctx.Err() != nil { - return m, tea.Quit - } - return m, progressTickCmd() - - default: - if m.ctx.Err() != nil { - return m, tea.Quit - } - return m, nil - } -} - -func (m progressModel) View() string { - var messagesBuilder strings.Builder - for i, msg := range m.messages { - if len(msg) > 0 { - messagesBuilder.WriteString(msg) - if i < len(m.messages)-1 { - messagesBuilder.WriteString("\n") - } - } - } - - var finalBuilder strings.Builder - if messagesBuilder.Len() > 0 { - finalBuilder.WriteString(messageChannelStyle.Render(messagesBuilder.String())) - finalBuilder.WriteString("\n") - } - - m.current = m.counter.Load() - finalBuilder.WriteString("\n ") - finalBuilder.WriteString(strconv.FormatUint(m.current, 10)) - finalBuilder.WriteString("/") - finalBuilder.WriteString(strconv.FormatUint(m.maxValue, 10)) - finalBuilder.WriteString(" - ") - finalBuilder.WriteString(time.Since(m.startTime).Round(time.Second / 10).String()) - finalBuilder.WriteString("\n ") - finalBuilder.WriteString(m.progress.ViewAs(float64(m.current) / float64(m.maxValue))) - finalBuilder.WriteString("\n\n ") - if m.cancelling { - finalBuilder.WriteString(helpStyle.Render("Stopping...")) - } else { - finalBuilder.WriteString(helpStyle.Render("Press Ctrl+C to quit")) - } - return finalBuilder.String() -} - -func progressTickCmd() tea.Cmd { - return tea.Tick(time.Millisecond*250, func(t time.Time) tea.Msg { - return tickMsg(t) - }) -} - -var infiniteProgressStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#00D4FF")) - -type infiniteProgressModel struct { - spinner spinner.Model - startTime time.Time - counter *atomic.Uint64 - messages []string - ctx context.Context //nolint:containedctx - quit bool - cancel context.CancelFunc - cancelling bool -} - -func (m infiniteProgressModel) Init() tea.Cmd { - return m.spinner.Tick -} - -func (m infiniteProgressModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { - switch msg := msg.(type) { - case tea.KeyMsg: - if msg.Type == tea.KeyCtrlC { - m.cancelling = true - m.cancel() - } - return m, nil - - case runtimeMessage: - var msgBuilder strings.Builder - msgBuilder.WriteString("[") - msgBuilder.WriteString(msg.timestamp.Format("15:04:05")) - msgBuilder.WriteString("] ") - switch msg.level { - case runtimeMessageLevelError: - msgBuilder.WriteString(errorStyle.Render("ERROR: ")) - case runtimeMessageLevelWarning: - msgBuilder.WriteString(warningStyle.Render("WARNING: ")) - } - msgBuilder.WriteString(msg.text) - m.messages = append(m.messages[1:], msgBuilder.String()) - if m.ctx.Err() != nil { - m.quit = true - return m, tea.Quit - } - return m, nil - - default: - if m.ctx.Err() != nil { - m.quit = true - return m, tea.Quit - } - var cmd tea.Cmd - m.spinner, cmd = m.spinner.Update(msg) - return m, cmd - } -} - -func (m infiniteProgressModel) View() string { - var messagesBuilder strings.Builder - for i, msg := range m.messages { - if len(msg) > 0 { - messagesBuilder.WriteString(msg) - if i < len(m.messages)-1 { - messagesBuilder.WriteString("\n") - } - } - } - - var finalBuilder strings.Builder - if messagesBuilder.Len() > 0 { - finalBuilder.WriteString(messageChannelStyle.Render(messagesBuilder.String())) - finalBuilder.WriteString("\n") - } - - if m.quit { - finalBuilder.WriteString("\n ") - finalBuilder.WriteString(strconv.FormatUint(m.counter.Load(), 10)) - finalBuilder.WriteString(" ") - finalBuilder.WriteString(infiniteProgressStyle.Render("∙∙∙∙∙")) - finalBuilder.WriteString(" ") - finalBuilder.WriteString(time.Since(m.startTime).Round(time.Second / 10).String()) - finalBuilder.WriteString("\n\n") - } else { - finalBuilder.WriteString("\n ") - finalBuilder.WriteString(strconv.FormatUint(m.counter.Load(), 10)) - finalBuilder.WriteString(" ") - finalBuilder.WriteString(m.spinner.View()) - finalBuilder.WriteString(" ") - finalBuilder.WriteString(time.Since(m.startTime).Round(time.Second / 10).String()) - finalBuilder.WriteString("\n\n ") - if m.cancelling { - finalBuilder.WriteString(helpStyle.Render("Stopping...")) - } else { - finalBuilder.WriteString(helpStyle.Render("Press Ctrl+C to quit")) - } - } - return finalBuilder.String() -} - -func (q sarin) streamProgress( - ctx context.Context, - cancel context.CancelFunc, - done chan<- struct{}, - total uint64, - counter *atomic.Uint64, - messageChannel <-chan runtimeMessage, -) { - var program *tea.Program - if total > 0 { - model := progressModel{ - progress: progress.New(progress.WithGradient("#151594", "#00D4FF")), - startTime: time.Now(), - messages: make([]string, 8), - counter: counter, - current: 0, - maxValue: total, - ctx: ctx, - cancel: cancel, - } - - program = tea.NewProgram(model) - } else { - model := infiniteProgressModel{ - spinner: spinner.New( - spinner.WithSpinner( - spinner.Spinner{ - Frames: []string{ - "●∙∙∙∙", - "∙●∙∙∙", - "∙∙●∙∙", - "∙∙∙●∙", - "∙∙∙∙●", - "∙∙∙●∙", - "∙∙●∙∙", - "∙●∙∙∙", - }, - FPS: time.Second / 8, //nolint:mnd - }, - ), - spinner.WithStyle(infiniteProgressStyle), - ), - startTime: time.Now(), - counter: counter, - messages: make([]string, 8), - ctx: ctx, - cancel: cancel, - quit: false, - } - - program = tea.NewProgram(model) - } - - go func() { - for msg := range messageChannel { - program.Send(msg) - } - }() - - if _, err := program.Run(); err != nil { - panic(err) - } - - done <- struct{}{} -} diff --git a/internal/sarin/stop.go b/internal/sarin/stop.go new file mode 100644 index 0000000..f43ea26 --- /dev/null +++ b/internal/sarin/stop.go @@ -0,0 +1,55 @@ +package sarin + +import ( + "fmt" + "os" + "sync" + "sync/atomic" + + tea "github.com/charmbracelet/bubbletea" +) + +const forceExitCode = 130 + +// StopController coordinates a two-stage shutdown. +// +// The first Stop call cancels the supplied context so workers and the job +// loop can drain. The second Stop call restores the terminal (if a bubbletea +// program has been attached) and calls os.Exit(forceExitCode), bypassing any +// in-flight captcha polls, Lua/JS scripts, or HTTP requests that would +// otherwise keep the process alive. +type StopController struct { + count atomic.Int32 + cancel func() + mu sync.Mutex + program *tea.Program +} + +func NewStopController(cancel func()) *StopController { + return &StopController{cancel: cancel} +} + +// AttachProgram registers the active bubbletea program so the terminal state +// can be restored before os.Exit on the forced shutdown path. Pass nil to +// detach once the program has finished. +func (s *StopController) AttachProgram(program *tea.Program) { + s.mu.Lock() + s.program = program + s.mu.Unlock() +} + +func (s *StopController) Stop() { + switch s.count.Add(1) { + case 1: + s.cancel() + case 2: + s.mu.Lock() + p := s.program + s.mu.Unlock() + if p != nil { + _ = p.ReleaseTerminal() + } + fmt.Fprintln(os.Stderr, "killing...") + os.Exit(forceExitCode) + } +} diff --git a/internal/sarin/tui.go b/internal/sarin/tui.go new file mode 100644 index 0000000..534d023 --- /dev/null +++ b/internal/sarin/tui.go @@ -0,0 +1,300 @@ +package sarin + +import ( + "context" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/charmbracelet/bubbles/progress" + "github.com/charmbracelet/bubbles/spinner" + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" +) + +type tickMsg time.Time + +var ( + helpStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#d1d1d1")) + errorStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#FC5B5B")).Bold(true) + warningStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#FFD93D")).Bold(true) + messageChannelStyle = lipgloss.NewStyle(). + Border(lipgloss.ThickBorder(), false, false, false, true). + BorderForeground(lipgloss.Color("#757575")). + PaddingLeft(1). + Margin(1, 0, 0, 0). + Foreground(lipgloss.Color("#888888")) +) + +type progressModel struct { + progress progress.Model + startTime time.Time + messages []string + counter *atomic.Uint64 + current uint64 + maxValue uint64 + ctx context.Context //nolint:containedctx + stop func() + cancelling bool +} + +func (m progressModel) Init() tea.Cmd { + return tea.Batch(progressTickCmd()) +} + +func (m progressModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.KeyMsg: + if msg.Type == tea.KeyCtrlC { + m.cancelling = true + m.stop() + } + return m, nil + + case tea.WindowSizeMsg: + m.progress.Width = max(10, msg.Width-1) + if m.ctx.Err() != nil { + return m, tea.Quit + } + return m, nil + + case runtimeMessage: + var msgBuilder strings.Builder + msgBuilder.WriteString("[") + msgBuilder.WriteString(msg.timestamp.Format("15:04:05")) + msgBuilder.WriteString("] ") + switch msg.level { + case runtimeMessageLevelError: + msgBuilder.WriteString(errorStyle.Render("ERROR: ")) + case runtimeMessageLevelWarning: + msgBuilder.WriteString(warningStyle.Render("WARNING: ")) + } + msgBuilder.WriteString(msg.text) + m.messages = append(m.messages[1:], msgBuilder.String()) + if m.ctx.Err() != nil { + return m, tea.Quit + } + return m, nil + + case tickMsg: + if m.ctx.Err() != nil { + return m, tea.Quit + } + return m, progressTickCmd() + + default: + if m.ctx.Err() != nil { + return m, tea.Quit + } + return m, nil + } +} + +func (m progressModel) View() string { + var messagesBuilder strings.Builder + for i, msg := range m.messages { + if len(msg) > 0 { + messagesBuilder.WriteString(msg) + if i < len(m.messages)-1 { + messagesBuilder.WriteString("\n") + } + } + } + + var finalBuilder strings.Builder + if messagesBuilder.Len() > 0 { + finalBuilder.WriteString(messageChannelStyle.Render(messagesBuilder.String())) + finalBuilder.WriteString("\n") + } + + m.current = m.counter.Load() + finalBuilder.WriteString("\n ") + finalBuilder.WriteString(strconv.FormatUint(m.current, 10)) + finalBuilder.WriteString("/") + finalBuilder.WriteString(strconv.FormatUint(m.maxValue, 10)) + finalBuilder.WriteString(" - ") + finalBuilder.WriteString(time.Since(m.startTime).Round(time.Second / 10).String()) + finalBuilder.WriteString("\n ") + finalBuilder.WriteString(m.progress.ViewAs(float64(m.current) / float64(m.maxValue))) + finalBuilder.WriteString("\n\n ") + if m.cancelling { + finalBuilder.WriteString(helpStyle.Render("Stopping... (Ctrl+C again to force)")) + } else { + finalBuilder.WriteString(helpStyle.Render("Press Ctrl+C to quit")) + } + return finalBuilder.String() +} + +func progressTickCmd() tea.Cmd { + return tea.Tick(time.Millisecond*250, func(t time.Time) tea.Msg { + return tickMsg(t) + }) +} + +var infiniteProgressStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#00D4FF")) + +type infiniteProgressModel struct { + spinner spinner.Model + startTime time.Time + counter *atomic.Uint64 + messages []string + ctx context.Context //nolint:containedctx + quit bool + stop func() + cancelling bool +} + +func (m infiniteProgressModel) Init() tea.Cmd { + return m.spinner.Tick +} + +func (m infiniteProgressModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.KeyMsg: + if msg.Type == tea.KeyCtrlC { + m.cancelling = true + m.stop() + } + return m, nil + + case runtimeMessage: + var msgBuilder strings.Builder + msgBuilder.WriteString("[") + msgBuilder.WriteString(msg.timestamp.Format("15:04:05")) + msgBuilder.WriteString("] ") + switch msg.level { + case runtimeMessageLevelError: + msgBuilder.WriteString(errorStyle.Render("ERROR: ")) + case runtimeMessageLevelWarning: + msgBuilder.WriteString(warningStyle.Render("WARNING: ")) + } + msgBuilder.WriteString(msg.text) + m.messages = append(m.messages[1:], msgBuilder.String()) + if m.ctx.Err() != nil { + m.quit = true + return m, tea.Quit + } + return m, nil + + default: + if m.ctx.Err() != nil { + m.quit = true + return m, tea.Quit + } + var cmd tea.Cmd + m.spinner, cmd = m.spinner.Update(msg) + return m, cmd + } +} + +func (m infiniteProgressModel) View() string { + var messagesBuilder strings.Builder + for i, msg := range m.messages { + if len(msg) > 0 { + messagesBuilder.WriteString(msg) + if i < len(m.messages)-1 { + messagesBuilder.WriteString("\n") + } + } + } + + var finalBuilder strings.Builder + if messagesBuilder.Len() > 0 { + finalBuilder.WriteString(messageChannelStyle.Render(messagesBuilder.String())) + finalBuilder.WriteString("\n") + } + + if m.quit { + finalBuilder.WriteString("\n ") + finalBuilder.WriteString(strconv.FormatUint(m.counter.Load(), 10)) + finalBuilder.WriteString(" ") + finalBuilder.WriteString(infiniteProgressStyle.Render("∙∙∙∙∙")) + finalBuilder.WriteString(" ") + finalBuilder.WriteString(time.Since(m.startTime).Round(time.Second / 10).String()) + finalBuilder.WriteString("\n\n") + } else { + finalBuilder.WriteString("\n ") + finalBuilder.WriteString(strconv.FormatUint(m.counter.Load(), 10)) + finalBuilder.WriteString(" ") + finalBuilder.WriteString(m.spinner.View()) + finalBuilder.WriteString(" ") + finalBuilder.WriteString(time.Since(m.startTime).Round(time.Second / 10).String()) + finalBuilder.WriteString("\n\n ") + if m.cancelling { + finalBuilder.WriteString(helpStyle.Render("Stopping... (Ctrl+C again to force)")) + } else { + finalBuilder.WriteString(helpStyle.Render("Press Ctrl+C to quit")) + } + } + return finalBuilder.String() +} + +func (q sarin) streamProgress( + ctx context.Context, + stopCtrl *StopController, + done chan<- struct{}, + total uint64, + counter *atomic.Uint64, + messageChannel <-chan runtimeMessage, +) { + var program *tea.Program + if total > 0 { + model := progressModel{ + progress: progress.New(progress.WithGradient("#151594", "#00D4FF")), + startTime: time.Now(), + messages: make([]string, 8), + counter: counter, + current: 0, + maxValue: total, + ctx: ctx, + stop: stopCtrl.Stop, + } + + program = tea.NewProgram(model) + } else { + model := infiniteProgressModel{ + spinner: spinner.New( + spinner.WithSpinner( + spinner.Spinner{ + Frames: []string{ + "●∙∙∙∙", + "∙●∙∙∙", + "∙∙●∙∙", + "∙∙∙●∙", + "∙∙∙∙●", + "∙∙∙●∙", + "∙∙●∙∙", + "∙●∙∙∙", + }, + FPS: time.Second / 8, //nolint:mnd + }, + ), + spinner.WithStyle(infiniteProgressStyle), + ), + startTime: time.Now(), + counter: counter, + messages: make([]string, 8), + ctx: ctx, + stop: stopCtrl.Stop, + quit: false, + } + + program = tea.NewProgram(model) + } + + stopCtrl.AttachProgram(program) + defer stopCtrl.AttachProgram(nil) + + go func() { + for msg := range messageChannel { + program.Send(msg) + } + }() + + if _, err := program.Run(); err != nil { + panic(err) + } + + done <- struct{}{} +}