diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 91e80136..37808278 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -48,6 +48,8 @@ type Runner struct { pdcpWriter *pdcp.UploadWriter } +const maxInputScanTokenSize = 4 * 1024 * 1024 + // New creates a new runner from provided configuration options func New(options *clients.Options) (*Runner, error) { // Disable coloring of log output if asked by user @@ -422,7 +424,9 @@ func (r *Runner) processInputElementWorker(inputs chan taskInput, wg *sync.WaitG func (r *Runner) normalizeAndQueueInputs(inputs chan taskInput) error { // Process Normal Inputs for _, text := range r.options.Inputs { - r.processInputItem(text, inputs) + for _, entry := range splitInputEntries(text) { + r.processInputItem(entry, inputs) + } } if r.options.InputList != "" { @@ -437,25 +441,42 @@ func (r *Runner) normalizeAndQueueInputs(inputs chan taskInput) error { }() scanner := bufio.NewScanner(file) + scanner.Buffer(make([]byte, 0, 64*1024), maxInputScanTokenSize) for scanner.Scan() { - text := scanner.Text() - if text != "" { - r.processInputItem(text, inputs) + for _, entry := range splitInputEntries(scanner.Text()) { + r.processInputItem(entry, inputs) } } + if err := scanner.Err(); err != nil { + return errkit.Wrap(err, "could not read input file") + } } if r.hasStdin { scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 0, 64*1024), maxInputScanTokenSize) for scanner.Scan() { - text := scanner.Text() - if text != "" { - r.processInputItem(text, inputs) + for _, entry := range splitInputEntries(scanner.Text()) { + r.processInputItem(entry, inputs) } } + if err := scanner.Err(); err != nil { + return errkit.Wrap(err, "could not read stdin") + } } return nil } +func splitInputEntries(text string) []string { + var entries []string + for _, entry := range strings.Split(text, ",") { + entry = strings.TrimSpace(entry) + if entry != "" { + entries = append(entries, entry) + } + } + return entries +} + // resolveFQDN resolves a FQDN and returns the IP addresses func (r *Runner) resolveFQDN(target string) ([]string, error) { // If the host is a Domain, then perform resolution and discover all IP diff --git a/internal/runner/runner_test.go b/internal/runner/runner_test.go index b4ed856e..ad215aa1 100644 --- a/internal/runner/runner_test.go +++ b/internal/runner/runner_test.go @@ -38,6 +38,107 @@ func Test_InputDomain_processInputItem(t *testing.T) { require.ElementsMatch(t, expected, got, "could not get correct taskInputs") } +func TestSplitInputEntries(t *testing.T) { + testCases := []struct { + name string + input string + expected []string + }{ + { + name: "single input remains unchanged", + input: "example.com", + expected: []string{"example.com"}, + }, + { + name: "comma separated inputs are split", + input: "example.com,one.one.one.one,192.168.1.1", + expected: []string{"example.com", "one.one.one.one", "192.168.1.1"}, + }, + { + name: "whitespace and empty entries are ignored", + input: " example.com, , one.one.one.one ,, ", + expected: []string{"example.com", "one.one.one.one"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + got := splitInputEntries(tc.input) + require.Equal(t, tc.expected, got) + }) + } +} + +func TestNormalizeAndQueueInputsFromFile(t *testing.T) { + file, err := os.CreateTemp(t.TempDir(), "tlsx-inputs-*.txt") + require.NoError(t, err) + defer file.Close() + + _, err = file.WriteString("example.com,one.one.one.one\n192.168.1.1\n") + require.NoError(t, err) + + options := &clients.Options{ + InputList: file.Name(), + Ports: []string{"443"}, + } + runner := &Runner{options: options} + + inputs := make(chan taskInput, 10) + err = runner.normalizeAndQueueInputs(inputs) + require.NoError(t, err) + close(inputs) + + var got []taskInput + for task := range inputs { + got = append(got, task) + } + + expected := []taskInput{ + {host: "example.com", port: "443"}, + {host: "one.one.one.one", port: "443"}, + {host: "192.168.1.1", port: "443"}, + } + require.ElementsMatch(t, expected, got) +} + +func TestNormalizeAndQueueInputsFromStdin(t *testing.T) { + file, err := os.CreateTemp(t.TempDir(), "tlsx-stdin-*.txt") + require.NoError(t, err) + + _, err = file.WriteString("example.com,one.one.one.one\n") + require.NoError(t, err) + _, err = file.Seek(0, 0) + require.NoError(t, err) + defer file.Close() + + originalStdin := os.Stdin + os.Stdin = file + defer func() { + os.Stdin = originalStdin + }() + + options := &clients.Options{ + Ports: []string{"443"}, + } + runner := &Runner{options: options, hasStdin: true, hasStdinSet: true} + + inputs := make(chan taskInput, 10) + err = runner.normalizeAndQueueInputs(inputs) + require.NoError(t, err) + close(inputs) + + var got []taskInput + for task := range inputs { + got = append(got, task) + } + + expected := []taskInput{ + {host: "example.com", port: "443"}, + {host: "one.one.one.one", port: "443"}, + } + require.ElementsMatch(t, expected, got) +} + func Test_InputForMultipleIps_processInputItem(t *testing.T) { options := &clients.Options{ Ports: []string{"443"},