Skip to content

Commit f1668f9

Browse files
DnFreddieDnFreddie
DnFreddie
and
DnFreddie
authored
feat: handle stdin piped input (#71)
Co-authored-by: DnFreddie <[email protected]>
1 parent 8157518 commit f1668f9

File tree

4 files changed

+216
-25
lines changed

4 files changed

+216
-25
lines changed

‎internal/probe.go‎

Lines changed: 43 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ type Probe struct {
2020
Logger *slog.Logger
2121
// Channel to send back partial results.
2222
ReportCh chan *Report
23+
// URLs (HTTP), host/port strings (TCP) or domains (DNS).
24+
Input []string
2325
}
2426

2527
// Ensures the probe setup is correct.
@@ -76,34 +78,51 @@ func (p Probe) Do(ctx context.Context) error {
7678
p.Logger.Debug(
7779
"New protocol", "count", count, "protocol", proto,
7880
)
79-
var errMessage string
80-
rhost, extra, err := proto.Probe("")
81-
if err != nil {
82-
errMessage = err.Error()
81+
if len(p.Input) == 0 {
82+
// Probe default list of urls
83+
rhost, extra, err := proto.Probe("")
84+
report := Report{
85+
ProtocolID: proto.String(),
86+
Time: time.Since(start),
87+
Error: err,
88+
RHost: rhost,
89+
Extra: extra,
90+
}
91+
p.Logger.Debug(
92+
"Sending report back",
93+
"count", count, "report", report,
94+
)
95+
p.ReportCh <- &report
96+
time.Sleep(p.Delay)
97+
continue
8398
}
84-
report := Report{
85-
ProtocolID: proto.String(),
86-
RHost: rhost,
87-
Time: time.Since(start),
88-
Error: errMessage,
89-
Extra: extra,
99+
100+
// iterate over User provided inputs
101+
for _, addr := range p.Input {
102+
rhost, extra, err := proto.Probe(addr)
103+
report := Report{
104+
ProtocolID: proto.String(),
105+
Time: time.Since(start),
106+
Error: err,
107+
RHost: rhost,
108+
Extra: extra,
109+
}
110+
p.Logger.Debug(
111+
"Sending report back for address",
112+
"count", count, "address", addr, "report", report,
113+
)
114+
p.ReportCh <- &report
115+
time.Sleep(p.Delay)
90116
}
91-
p.Logger.Debug(
92-
"Sending report back",
93-
"count", count, "report", report,
94-
)
95-
p.ReportCh <- &report
96-
time.Sleep(p.Delay)
97117
}
118+
time.Sleep(p.Delay)
98119
}
99-
p.Logger.Debug(
100-
"Iteration finished", "count", count, "p.Count", p.Count,
101-
)
102-
count++
103-
if count == p.Count {
104-
p.Logger.Debug("Count limit reached", "count", count)
105-
return nil
106-
}
120+
}
121+
p.Logger.Debug("Iteration finished", "count", count, "p.Count", p.Count)
122+
count++
123+
if count == p.Count {
124+
p.Logger.Debug("Count limit reached", "count", count)
125+
return nil
107126
}
108127
}
109128
}

‎internal/stdin.go‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package internal
2+
3+
import (
4+
"bytes"
5+
"fmt"
6+
"net/url"
7+
"os"
8+
"regexp"
9+
"strings"
10+
)
11+
12+
const (
13+
domainPattern = `^[a-zA-Z0-9][a-zA-Z0-9-]{1,61}[a-zA-Z0-9]\.[a-zA-Z]{2,}$`
14+
ipPattern = `^(\d{1,3}\.){3}\d{1,3}$`
15+
)
16+
17+
func ReadStdin() (string, error) {
18+
info, err := os.Stdin.Stat()
19+
if err != nil {
20+
return "", fmt.Errorf("failed to retrieve stdin : %w", err)
21+
}
22+
if (info.Mode() & os.ModeCharDevice) != 0 {
23+
return "", nil
24+
}
25+
buf := &bytes.Buffer{}
26+
_, err = buf.ReadFrom(os.Stdin)
27+
if err != nil {
28+
return "", fmt.Errorf("reading from stdin: %w", err)
29+
}
30+
return buf.String(), nil
31+
}
32+
33+
var (
34+
ipRegex = regexp.MustCompile(ipPattern)
35+
domainRegex = regexp.MustCompile(domainPattern)
36+
)
37+
38+
// Ensures the inputs are correct
39+
func validateInput(addr string) bool {
40+
_, err := url.ParseRequestURI(addr)
41+
if err == nil {
42+
return true
43+
}
44+
if ipRegex.MatchString(addr) || domainRegex.MatchString(addr) {
45+
return true
46+
}
47+
return false
48+
}
49+
func ProcessInputs(s string) ([]string, error) {
50+
var inputs []string
51+
for _, input := range strings.Fields(s) {
52+
if validateInput(input) {
53+
inputs = append(inputs, input)
54+
} else {
55+
return inputs, fmt.Errorf("invalid address format: %s", input)
56+
}
57+
}
58+
return inputs, nil
59+
}

‎internal/stdin_test.go‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
package internal
2+
3+
import (
4+
"os"
5+
"testing"
6+
)
7+
8+
// Tests [validateInput] func
9+
func TestValidateInput(t *testing.T) {
10+
entries := []struct {
11+
input string
12+
expected bool
13+
}{
14+
{"192.180.33.25", true},
15+
{"example.com", true},
16+
{"http://192.180.33.25", true},
17+
{"https://example.org/path?query=123", true},
18+
{"test-domain.org", true},
19+
{"256.256.256.256", true},
20+
{"invalid@domain", false},
21+
{"not-a-domain", false},
22+
}
23+
for _, entry := range entries {
24+
t.Run(entry.input, func(t *testing.T) {
25+
result := validateInput(entry.input)
26+
if result != entry.expected {
27+
t.Errorf(
28+
"validateInput(%q) = %v; want %v",
29+
entry.input,
30+
result,
31+
entry.expected,
32+
)
33+
}
34+
})
35+
}
36+
}
37+
38+
// Tests [ReadStdin] and [ProcessInputs] func
39+
func TestReadAndProcessInputs(t *testing.T) {
40+
tests := []struct {
41+
name string
42+
input string
43+
wantInputs []string
44+
wantErr bool
45+
}{
46+
{
47+
name: "single valid url",
48+
input: "http://192.168.1.1\n",
49+
wantInputs: []string{"http://192.168.1.1"},
50+
wantErr: false,
51+
},
52+
{
53+
name: "multiple urls with whitespace",
54+
input: " http://192.168.1.1 https://google.com \n",
55+
wantInputs: []string{"http://192.168.1.1", "https://google.com"},
56+
wantErr: false,
57+
},
58+
{
59+
name: "mixed valid and invalid",
60+
input: "http://192.168.1.1 not_valid example.com\n",
61+
wantInputs: []string{"http://192.168.1.1"},
62+
wantErr: true,
63+
},
64+
{
65+
name: "empty input",
66+
input: "",
67+
wantInputs: []string{},
68+
wantErr: false,
69+
},
70+
}
71+
for _, tt := range tests {
72+
t.Run(tt.name, func(t *testing.T) {
73+
r, w, _ := os.Pipe()
74+
defer r.Close()
75+
os.Stdin = r
76+
w.Write([]byte(tt.input))
77+
w.Close()
78+
got, readErr := ReadStdin()
79+
if readErr != nil {
80+
t.Errorf("ReadStdin() error = %v", readErr)
81+
return
82+
}
83+
inputs, processErr := ProcessInputs(got)
84+
if (processErr != nil) != tt.wantErr {
85+
t.Errorf(
86+
"ProcessInputs() error = %v, wantErr %v",
87+
processErr,
88+
tt.wantErr,
89+
)
90+
return
91+
}
92+
if len(inputs) != len(tt.wantInputs) {
93+
t.Errorf("got %v, want %v", inputs, tt.wantInputs)
94+
return
95+
}
96+
for i := range inputs {
97+
if inputs[i] != tt.wantInputs[i] {
98+
t.Errorf("got %v, want %v", inputs, tt.wantInputs)
99+
return
100+
}
101+
}
102+
})
103+
}
104+
}

‎main.go‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,14 @@ func main() {
4242
Level: lvl,
4343
}))
4444
var opts internal.Options
45+
stdin, err := internal.ReadStdin()
46+
if err != nil {
47+
fatal(fmt.Errorf("reading stdin: %w", err))
48+
}
49+
inputs, err := internal.ProcessInputs(stdin)
50+
if err != nil {
51+
fatal(fmt.Errorf("failed to process the inputs: %w", err))
52+
}
4553
opts.Parse()
4654
if opts.Debug {
4755
lvl.Set(slog.LevelDebug)
@@ -98,6 +106,7 @@ func main() {
98106
Delay: opts.Delay,
99107
Logger: logger,
100108
ReportCh: reportCh,
109+
Input: inputs,
101110
}
102111
var format internal.Format
103112
switch {
@@ -126,7 +135,7 @@ func main() {
126135
}
127136
}()
128137
logger.Debug("Running ...", "setup", probe)
129-
err := probe.Do(ctx)
138+
err = probe.Do(ctx)
130139
if err != nil {
131140
fatal(fmt.Errorf("running probe: %w", err))
132141
}

0 commit comments

Comments
 (0)