diff --git a/cmd/icmperf/main.go b/cmd/icmperf/main.go index 034e851..c1d7085 100644 --- a/cmd/icmperf/main.go +++ b/cmd/icmperf/main.go @@ -1,42 +1,58 @@ package main import ( - "context" - "net" - "time" + "fmt" "github.com/alecthomas/kong" - tea "github.com/charmbracelet/bubbletea" - "golang.org/x/net/icmp" + "github.com/prometheus-community/pro-bing" "github.com/b4nst/icmperf/pkg/cli" - "github.com/b4nst/icmperf/pkg/model" - "github.com/b4nst/icmperf/pkg/pinger" "github.com/b4nst/icmperf/pkg/recorder" + "github.com/b4nst/icmperf/pkg/session" ) func main() { cli := cli.CLI{} ktx := kong.Parse(&cli) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - pinger := pinger.NewPinger(cli.BindAddr, cli.MTU, cli.Timeout) - record := recorder.NewRecord() - pinger.OnRecv(func(m *icmp.Message, t time.Time) error { - body := m.Body.(*icmp.Echo) - id := uint64(body.ID)<<32 | uint64(body.Seq) - record.PacketReceived(id, len(body.Data), t) - return nil - }) - - peer := net.UDPAddr(cli.Target) - m := model.NewModel(pinger, record, &peer, cli.MTU, cli.Duration) - if err := pinger.Start(ctx); err != nil { - ktx.FatalIfErrorf(err) + + pingers := make([]*probing.Pinger, 0, len(cli.DataSizes)+1) + + // Latenncy pinger + p, err := probing.NewPinger(cli.Target) + ktx.FatalIfErrorf(err) + if cli.Duration > 0 { + p.Timeout = cli.Duration + } else { + p.Count = cli.Count } + p.SetPrivileged(cli.Privileged) + p.Size = 24 // Minimum size + pingers = append(pingers, p) - if _, err := tea.NewProgram(m).Run(); err != nil { + // Data pingers + for _, ds := range cli.DataSizes { + p, err := probing.NewPinger(cli.Target) ktx.FatalIfErrorf(err) + if cli.Duration > 0 { + p.Timeout = cli.Duration + } else { + p.Count = cli.Count + } + p.Size = ds + p.SetPrivileged(cli.Privileged) + pingers = append(pingers, p) + } + + s := session.NewSession(pingers) + ktx.FatalIfErrorf(s.Run()) + + stats := s.Statistics() + stat, err := recorder.ProcessStats(stats) + ktx.FatalIfErrorf(err) + + for _, s := range stats { + fmt.Printf("%s\n", s) } + fmt.Println("- - - - - - - - - - - - - - - - - - - - - - - - -") + fmt.Printf("%s\n", stat) } diff --git a/go.mod b/go.mod index 01054ee..51454bc 100644 --- a/go.mod +++ b/go.mod @@ -7,8 +7,11 @@ require ( github.com/charmbracelet/bubbles v0.19.0 github.com/charmbracelet/bubbletea v1.1.0 github.com/dustin/go-humanize v1.0.1 - github.com/emirpasic/gods v1.18.1 + github.com/montanaflynn/stats v0.7.1 + github.com/prometheus-community/pro-bing v0.4.1 + github.com/stretchr/testify v1.9.0 golang.org/x/net v0.28.0 + golang.org/x/sync v0.8.0 ) require ( @@ -17,7 +20,9 @@ require ( github.com/charmbracelet/lipgloss v0.13.0 // indirect github.com/charmbracelet/x/ansi v0.2.3 // indirect github.com/charmbracelet/x/term v0.2.0 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect + github.com/google/uuid v1.6.0 // indirect github.com/lucasb-eyer/go-colorful v1.2.0 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-localereader v0.0.1 // indirect @@ -25,8 +30,9 @@ require ( github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect github.com/muesli/cancelreader v0.2.2 // indirect github.com/muesli/termenv v0.15.2 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/rivo/uniseg v0.4.7 // indirect - golang.org/x/sync v0.8.0 // indirect golang.org/x/sys v0.24.0 // indirect golang.org/x/text v0.17.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 62364a4..e8aa722 100644 --- a/go.sum +++ b/go.sum @@ -18,12 +18,14 @@ github.com/charmbracelet/x/ansi v0.2.3 h1:VfFN0NUpcjBRd4DnKfRaIRo53KRgey/nhOoEqo github.com/charmbracelet/x/ansi v0.2.3/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw= github.com/charmbracelet/x/term v0.2.0 h1:cNB9Ot9q8I711MyZ7myUR5HFWL/lc3OpU8jZ4hwm0x0= github.com/charmbracelet/x/term v0.2.0/go.mod h1:GVxgxAbjUrmpvIINHIQnJJKpMlHiZ4cktEQCN6GWyF0= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= -github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc= -github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ= github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= @@ -34,15 +36,23 @@ github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2J github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/montanaflynn/stats v0.7.1 h1:etflOAAHORrCC44V+aR6Ftzort912ZU+YLiSTuV8eaE= +github.com/montanaflynn/stats v0.7.1/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= github.com/muesli/termenv v0.15.2 h1:GohcuySI0QmI3wN8Ok9PtKGkgkFIk7y6Vpb5PvrY+Wo= github.com/muesli/termenv v0.15.2/go.mod h1:Epx+iuz8sNs7mNKhxzH4fWXGNpZwUaJKRS1noLXviQ8= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus-community/pro-bing v0.4.1 h1:aMaJwyifHZO0y+h8+icUz0xbToHbia0wdmzdVZ+Kl3w= +github.com/prometheus-community/pro-bing v0.4.1/go.mod h1:aLsw+zqCaDoa2RLVVSX3+UiCkBBXTMtZC3c7EkfWnAE= github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE= golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg= golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ= @@ -53,3 +63,7 @@ golang.org/x/sys v0.24.0 h1:Twjiwq9dn6R1fQcyiK+wQyHWfaz/BJB+YIpzU/Cv3Xg= golang.org/x/sys v0.24.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc= golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go index 50ed2c6..bad5583 100644 --- a/pkg/cli/cli.go +++ b/pkg/cli/cli.go @@ -2,25 +2,30 @@ package cli import ( "fmt" - "net" "time" ) -type Target net.UDPAddr +type DataSizes []int -func (t *Target) UnmarshalText(text []byte) error { - ip, err := net.LookupIP(string(text)) - if err != nil { - return fmt.Errorf("failed to lookup target: %w", err) +func (ds *DataSizes) Validate() error { + if len(*ds) == 0 { + return fmt.Errorf("data sizes must have at least one element") + } + + for _, size := range *ds { + if size < 24 { + return fmt.Errorf("data size must be greater than 0") + } } - t.IP = ip[0] return nil } type CLI struct { - Target Target `arg:"" help:"The target host to ping."` - MTU int `help:"The maximum transmission unit of your interface." short:"m" default:"1500"` - Timeout time.Duration `help:"The timeout for each ping." short:"t" default:"5s"` - Duration time.Duration `help:"The duration of the test." short:"d" default:"30s"` - BindAddr string `help:"The address to bind the ICMP listener to." default:"0.0.0.0" short:"l"` + Timeout time.Duration `help:"The timeout for each ping." short:"t" default:"5s"` + Duration time.Duration `help:"The duration of the test." short:"d" xor:"duration"` + Count int `help:"The number of pings to send." short:"c" default:"10" xor:"duration"` + Privileged bool `help:"Use privileged ICMP sockets." default:"false" short:"p"` + DataSizes []int `help:"The size of the data to send." short:"s" default:"1024"` + + Target string `arg:"" help:"The target host to ping."` } diff --git a/pkg/model/model.go b/pkg/model/model.go index 760afb9..740fd69 100644 --- a/pkg/model/model.go +++ b/pkg/model/model.go @@ -1,8 +1,8 @@ package model import ( + "fmt" "net" - "strconv" "strings" "time" @@ -20,11 +20,8 @@ var ( ) type Model struct { - pinger *pinger.Pinger record *recorder.Record - peer *net.UDPAddr payload []byte - stats *recorder.Stats duration time.Duration timer timer.Model @@ -33,11 +30,8 @@ type Model struct { func NewModel(pinger *pinger.Pinger, record *recorder.Record, peer *net.UDPAddr, mtu int, duration time.Duration) *Model { return &Model{ - pinger: pinger, record: record, - peer: peer, payload: make([]byte, mtu-28), - stats: nil, duration: duration, timer: timer.NewWithInterval(duration, 200*time.Millisecond), @@ -45,9 +39,7 @@ func NewModel(pinger *pinger.Pinger, record *recorder.Record, peer *net.UDPAddr, } } -type latencyTick time.Time type pingTick time.Time -type statsMsg *recorder.Stats func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { cmds := []tea.Cmd{} @@ -62,23 +54,11 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.timer, tc = m.timer.Update(msg) pc = m.progress.SetPercent(1.0 - m.timer.Timeout.Seconds()/m.duration.Seconds()) cmds = append(cmds, tc, pc) - case timer.TimeoutMsg: - cmds = append(cmds, m.progress.SetPercent(1.0)) - cmds = append(cmds, m.statsCmd()) case tea.WindowSizeMsg: m.progress.Width = min(msg.Width-4, maxWidth) - case latencyTick: - if !m.timer.Timedout() { - cmds = append(cmds, m.pingLatency()) - } - case pingTick: - if !m.timer.Timedout() { - cmds = append(cmds, m.pingBandwidth()) - } - case statsMsg: - m.stats = msg + case error: + fmt.Println(msg) cmds = append(cmds, tea.Quit) - case progress.FrameMsg: progressModel, cmd := m.progress.Update(msg) m.progress = progressModel.(progress.Model) @@ -91,10 +71,6 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { func (m *Model) View() string { view := strings.Builder{} - view.WriteString("Probing peer at ") - view.WriteString(m.peer.IP.String()) - view.WriteString("\n\n") - view.WriteString("Probe size: ") view.WriteString(humanize.Bytes(uint64(len(m.payload)))) view.WriteString(" ") @@ -105,52 +81,9 @@ func (m *Model) View() string { view.WriteString(m.progress.View()) view.WriteString("\n\n") - if m.stats != nil { - view.WriteString("Bandwidth: ") - view.WriteString(humanize.Bytes(uint64(m.stats.Bandwidth()))) - view.WriteString("/s ") - - view.WriteString("Latency: ") - view.WriteString(m.stats.Latency().Round(time.Millisecond).String()) - view.WriteString(" ") - - view.WriteString("Loss rate: ") - view.WriteString(strconv.FormatFloat(m.stats.PacketLoss()*100, 'f', 2, 64)) - view.WriteString("%\n") - } - return view.String() } func (m *Model) Init() tea.Cmd { - return tea.Batch(m.pingLatency(), m.pingBandwidth(), m.timer.Init()) -} - -func (m *Model) pingLatency() tea.Cmd { - return tea.Tick(1*time.Second, func(t time.Time) tea.Msg { - sendAndRecord(m.pinger, m.peer, []byte{}, m.record) - return latencyTick(t) - }) -} - -func (m *Model) pingBandwidth() tea.Cmd { - return tea.Tick(500*time.Millisecond, func(t time.Time) tea.Msg { - sendAndRecord(m.pinger, m.peer, m.payload, m.record) - return pingTick(t) - }) -} - -func (m *Model) statsCmd() tea.Cmd { - return func() tea.Msg { - return statsMsg(m.record.Stats()) - } -} - -func sendAndRecord(p *pinger.Pinger, peer *net.UDPAddr, payload []byte, r *recorder.Record) error { - id, t, err := p.Send(peer, payload) - if err != nil { - return err - } - r.PacketSent(id, len(payload), t) - return nil + return tea.Batch(m.timer.Init()) } diff --git a/pkg/pinger/pinger.go b/pkg/pinger/pinger.go index c99857d..15978b1 100644 --- a/pkg/pinger/pinger.go +++ b/pkg/pinger/pinger.go @@ -2,6 +2,7 @@ package pinger import ( "context" + "fmt" "net" "os" "sync/atomic" @@ -19,7 +20,7 @@ type Pinger struct { seq atomic.Uint64 id int - recvCb func(*icmp.Message, time.Time) error + recvbuf []byte } func NewPinger(addr string, mtu int, timeout time.Duration) *Pinger { @@ -29,6 +30,8 @@ func NewPinger(addr string, mtu int, timeout time.Duration) *Pinger { MTU: mtu, seq: atomic.Uint64{}, id: os.Getpid() & 0xffff, + + recvbuf: make([]byte, mtu), } } @@ -38,44 +41,30 @@ func (p *Pinger) Start(ctx context.Context) error { return err } p.conn = conn - go p.listen(ctx) return nil } -func (p *Pinger) OnRecv(cb func(*icmp.Message, time.Time) error) { - p.recvCb = cb -} - -func (p *Pinger) listen(ctx context.Context) { - buf := make([]byte, p.MTU) - for { - if ctx.Err() != nil { - return - } - err := p.conn.SetReadDeadline(time.Now().Add(p.timeout)) - if err != nil { - // TODO: log error - continue - } - n, _, err := p.conn.ReadFrom(buf) - if err != nil { - // TODO: log error - continue - } - received := time.Now() - parsed, err := icmp.ParseMessage(1, buf[:n]) - if err != nil { - // TODO: log error - continue - } - if err := p.recvCb(parsed, received); err != nil { - // TODO: log error - continue - } +// Reply waits for an ICMP reply and fills the message with the response. +func (p *Pinger) Reply(message *icmp.Message) (time.Time, error) { + err := p.conn.SetReadDeadline(time.Now().Add(p.timeout)) + if err != nil { + return time.Time{}, fmt.Errorf("failed to set read deadline: %w", err) } + n, _, err := p.conn.ReadFrom(p.recvbuf) + if err != nil { + return time.Time{}, fmt.Errorf("failed to read from connection: %w", err) + } + received := time.Now() + parsed, err := icmp.ParseMessage(1, p.recvbuf[:n]) + if err != nil { + return time.Time{}, fmt.Errorf("failed to parse message: %w", err) + } + *message = *parsed + + return received, nil } -func (p *Pinger) Send(peer *net.UDPAddr, payload []byte) (uint64, time.Time, error) { +func (p *Pinger) Echo(peer *net.UDPAddr, payload []byte) (uint64, time.Time, error) { body := &icmp.Echo{ ID: p.id, Seq: int(p.seq.Add(1)),