├── README.md ├── cgroup_linux.go ├── cgroup_os.go ├── cmd └── main.go ├── example └── demo.go ├── go.mod ├── go.sum ├── netflow.go ├── netflow_test.go ├── netstat.go ├── netstat_test.go ├── process.go ├── process_test.go ├── ringqueue.go ├── types.go ├── user.go └── utils.go /README.md: -------------------------------------------------------------------------------- 1 | ## go-netflow 2 | 3 | go-netflow, capture process in/out traffic, similar to c Nethogs. 4 | 5 | [使用 golang 实现进程级流量监控](https://github.com/rfyiamcool/notes/blob/main/netflow.md) 6 | 7 | ### refer 8 | 9 | refer logic design link 10 | 11 | - [https://zhuanlan.zhihu.com/p/49981590](https://zhuanlan.zhihu.com/p/49981590) 12 | 13 | refer nethogs source link 14 | 15 | - [https://github.com/raboof/nethogs](https://github.com/raboof/nethogs) 16 | 17 | ### dep 18 | 19 | ``` 20 | yum install libpcap 21 | yum install libpcap-devel 22 | ``` 23 | 24 | ### cli usage 25 | 26 | netflow cli run: 27 | 28 | ``` 29 | go run cmd/main.go 30 | ``` 31 | 32 | stdout: 33 | 34 | ```text 35 | +---------+-------+------------------------------------------------+--------+--------+---------+---------+----------+ 36 | | PID | NAME | EXE | INODES | SUM IN | SUM OUT | IN RATE | OUT RATE | 37 | +---------+-------+------------------------------------------------+--------+--------+---------+---------+----------+ 38 | | 2256431 | Wget | /usr/bin/wget | 1 | 1.0 MB | 0 B | 339 kB | 0 B | 39 | +---------+-------+------------------------------------------------+--------+--------+---------+---------+----------+ 40 | | 2257200 | Wrk | /usr/bin/wrk | 5 | 2.0 MB | 16 kB | 653 kB | 5.2 kB | 41 | +---------+-------+------------------------------------------------+--------+--------+---------+---------+----------+ 42 | | 3707954 | Java | /usr/lib/jvm/java-7-openjdk-amd64/jre/bin/java | 10 | 457 B | 648 B | 152 B | 216 B | 43 | +---------+-------+------------------------------------------------+--------+--------+---------+---------+----------+ 44 | | 2245136 | Wget | /usr/bin/wget | 1 | 444 kB | 0 B | 148 kB | 0 B | 45 | +---------+-------+------------------------------------------------+--------+--------+---------+---------+----------+ 46 | | 2034103 | Nginx | /usr/sbin/nginx | 41 | 0 B | 0 B | 0 B | 0 B | 47 | +---------+-------+------------------------------------------------+--------+--------+---------+---------+----------+ 48 | ``` 49 | 50 | ### sdk simple usage: 51 | 52 | ```go 53 | package main 54 | 55 | import ( 56 | "encoding/json" 57 | "fmt" 58 | "time" 59 | 60 | "github.com/rfyiamcool/go-netflow" 61 | ) 62 | 63 | func main() { 64 | nf, err := netflow.New( 65 | netflow.WithCaptureTimeout(5 * time.Second), 66 | ) 67 | if err != nil { 68 | panic(err) 69 | } 70 | 71 | err = nf.Start() 72 | if err != nil { 73 | panic(err) 74 | } 75 | defer nf.Stop() 76 | 77 | <-nf.Done() 78 | 79 | var ( 80 | limit = 5 81 | recentSec = 5 82 | ) 83 | 84 | rank, err := nf.GetProcessRank(limit, recentSec) 85 | if err != nil { 86 | panic(err) 87 | } 88 | 89 | bs, err := json.MarshalIndent(rank, "", " ") 90 | if err != nil { 91 | panic(err) 92 | } 93 | 94 | fmt.Println(string(bs)) 95 | } 96 | ``` 97 | 98 | ### how to use sdk of go-netflow: 99 | 100 | #### set pcap filename 101 | 102 | Don't save pcap file by default. 103 | 104 | `WithStorePcap` option is used to save pcap file, use `tcpdump -nnr {filename}` command to read pcap file. 105 | 106 | ``` 107 | WithStorePcap(fpath string) 108 | ``` 109 | 110 | #### set custom pcap bpf filter. 111 | 112 | ``` 113 | WithPcapFilter(filter string) 114 | ``` 115 | 116 | #### set custom pcap bpf filter. 117 | 118 | example: 119 | 120 | - host xiaorui.cc and port 80 121 | - src host 123.56.223.52 and (dst port 3389 or 22) 122 | 123 | ``` 124 | WithPcapFilter(filter string) 125 | ``` 126 | 127 | #### limit netflow cpu/mem resource. 128 | 129 | ``` 130 | WithLimitCgroup(cpu float64, mem int) 131 | ``` 132 | 133 | #### set time to capturing packet. 134 | 135 | ``` 136 | WithCaptureTimeout(dur time.Duration) 137 | ``` 138 | 139 | #### set time to rescan process and inode data. 140 | 141 | ``` 142 | WithSyncInterval(dur time.Duration) 143 | ``` 144 | 145 | #### set the number of worker to consume pcap queue. 146 | 147 | ``` 148 | WithWorkerNum(num int) 149 | ``` 150 | 151 | #### set custom context. 152 | 153 | ``` 154 | WithCtx(ctx context.Context) 155 | ``` 156 | 157 | #### set custom devices to capture. 158 | 159 | ``` 160 | WithBindDevices(devs []string) 161 | ``` 162 | 163 | #### set pcap queue size. if the queue is full, new packet is thrown away. 164 | 165 | ``` 166 | WithQueueSize(size int) 167 | ``` 168 | 169 | ### types 170 | 171 | netflow.Interface 172 | 173 | ```go 174 | type Interface interface { 175 | Start() error 176 | Stop() 177 | Done() <-chan struct{} 178 | LoadCounter() int64 179 | GetProcessRank(int, int) ([]*Process, error) 180 | } 181 | ``` 182 | 183 | netflow.Process 184 | 185 | ```go 186 | type Process struct { 187 | Name string 188 | Pid string 189 | Exe string 190 | State string 191 | Inodes []string 192 | TrafficStats *trafficStatsEntry 193 | Ring []*trafficEntry 194 | } 195 | ``` 196 | 197 | netflow.trafficStatsEntry 198 | 199 | ```go 200 | type trafficStatsEntry struct { 201 | In int64 `json:"in"` 202 | Out int64 `json:"out"` 203 | InRate int64 `json:"in_rate"` 204 | OutRate int64 `json:"out_rate"` 205 | InputEWMA int64 `json:"input_ewma" valid:"-"` 206 | OutputEWMA int64 `json:"output_ewma" valid:"-"` 207 | } 208 | ``` 209 | 210 | netflow.trafficEntry 211 | 212 | ```go 213 | type trafficEntry struct { 214 | Timestamp int64 `json:"timestamp"` 215 | In int64 `json:"in"` 216 | Out int64 `json:"out"` 217 | } 218 | ``` 219 | -------------------------------------------------------------------------------- /cgroup_linux.go: -------------------------------------------------------------------------------- 1 | // +build linux 2 | 3 | package netflow 4 | 5 | import ( 6 | "fmt" 7 | "github.com/containerd/cgroups" 8 | specs "github.com/opencontainers/runtime-spec/specs-go" 9 | "github.com/spf13/cast" 10 | ) 11 | 12 | type cgroupsLimiter struct { 13 | controls []cgroups.Cgroup 14 | } 15 | 16 | // configure 17 | func (r *cgroupsLimiter) configure(pid int, core float64, mbn int) error { 18 | const ( 19 | cpuUnit = 10000 20 | memUnit = 1024 * 1024 21 | ) 22 | 23 | if core <= 0 { 24 | core = 1 25 | } 26 | 27 | var ( 28 | quota int64 = int64(core * cpuUnit) // core * 1u 29 | period uint64 = 10000 // 1u 30 | mem int64 = int64(mbn * memUnit) 31 | ) 32 | 33 | cfg := &specs.LinuxResources{ 34 | CPU: &specs.LinuxCPU{ 35 | Period: &period, 36 | Quota: "a, 37 | }, 38 | } 39 | 40 | if mem != 0 { 41 | cfg.Memory = &specs.LinuxMemory{ 42 | Limit: &mem, 43 | } 44 | } 45 | 46 | // file as /sys/fs/cgroup/cpu/netflow/... 47 | cgroupPath := "/netflow" 48 | control, err := cgroups.New(cgroups.V1, cgroups.StaticPath(cgroupPath), cfg) 49 | if err != nil { 50 | return err 51 | } 52 | 53 | r.controls = append(r.controls, control) 54 | err = control.Add(cgroups.Process{Pid: cast.ToInt(pid)}) 55 | return err 56 | } 57 | 58 | // free 59 | func (r *cgroupsLimiter) free() error { 60 | for _, ctrl := range r.controls { 61 | ctrl.Delete() 62 | } 63 | fmt.Println("exit") 64 | return nil 65 | } 66 | -------------------------------------------------------------------------------- /cgroup_os.go: -------------------------------------------------------------------------------- 1 | // +build !linux 2 | 3 | package netflow 4 | 5 | import ( 6 | "errors" 7 | ) 8 | 9 | type cgroupsLimiter struct { 10 | } 11 | 12 | func (r *cgroupsLimiter) free() error { 13 | return nil 14 | } 15 | 16 | func (r *cgroupsLimiter) configure(pid int, core float64, mb int) error { 17 | return errors.New("don't support cgroup") 18 | } 19 | -------------------------------------------------------------------------------- /cmd/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "os/signal" 7 | "syscall" 8 | "time" 9 | 10 | "github.com/dustin/go-humanize" 11 | "github.com/fatih/color" 12 | "github.com/olekukonko/tablewriter" 13 | log "github.com/sirupsen/logrus" 14 | "github.com/spf13/cast" 15 | 16 | "github.com/rfyiamcool/go-netflow" 17 | ) 18 | 19 | var ( 20 | nf netflow.Interface 21 | 22 | yellow = color.New(color.FgYellow).SprintFunc() 23 | red = color.New(color.FgRed).SprintFunc() 24 | info = color.New(color.FgGreen).SprintFunc() 25 | blue = color.New(color.FgBlue).SprintFunc() 26 | magenta = color.New(color.FgHiMagenta).SprintFunc() 27 | ) 28 | 29 | func start() { 30 | var err error 31 | 32 | nf, err = netflow.New() 33 | if err != nil { 34 | log.Fatal(err) 35 | } 36 | 37 | err = nf.Start() 38 | if err != nil { 39 | log.Fatal(err) 40 | } 41 | 42 | var ( 43 | recentRankLimit = 10 44 | 45 | sigch = make(chan os.Signal, 1) 46 | ticker = time.NewTicker(3 * time.Second) 47 | timeout = time.NewTimer(300 * time.Second) 48 | ) 49 | 50 | signal.Notify(sigch, 51 | syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT, 52 | syscall.SIGHUP, syscall.SIGUSR1, syscall.SIGUSR2, 53 | ) 54 | 55 | defer func() { 56 | nf.Stop() 57 | }() 58 | 59 | go func() { 60 | for { 61 | <-ticker.C 62 | rank, err := nf.GetProcessRank(recentRankLimit, 3) 63 | if err != nil { 64 | log.Errorf("GetProcessRank failed, err: %s", err.Error()) 65 | continue 66 | } 67 | 68 | clear() 69 | showTable(rank) 70 | } 71 | }() 72 | 73 | for { 74 | select { 75 | case <-sigch: 76 | return 77 | 78 | case <-timeout.C: 79 | return 80 | } 81 | } 82 | } 83 | 84 | func stop() { 85 | if nf == nil { 86 | return 87 | } 88 | 89 | nf.Stop() 90 | } 91 | 92 | const thold = 1024 * 1024 // 1mb 93 | 94 | func clear() { 95 | fmt.Printf("\x1b[2J") 96 | } 97 | 98 | func showTable(ps []*netflow.Process) { 99 | table := tablewriter.NewWriter(os.Stdout) 100 | table.SetHeader([]string{"pid", "name", "exe", "inodes", "sum_in", "sum_out", "in_rate", "out_rate"}) 101 | table.SetRowLine(true) 102 | 103 | items := [][]string{} 104 | for _, po := range ps { 105 | inRate := humanBytes(po.TrafficStats.InRate) 106 | if po.TrafficStats.InRate > int64(thold) { 107 | inRate = red(inRate) 108 | } 109 | 110 | outRate := humanBytes(po.TrafficStats.OutRate) 111 | if po.TrafficStats.OutRate > int64(thold) { 112 | outRate = red(outRate) 113 | } 114 | 115 | item := []string{ 116 | po.Pid, 117 | po.Name, 118 | po.Exe, 119 | cast.ToString(po.InodeCount), 120 | humanBytes(po.TrafficStats.In), 121 | humanBytes(po.TrafficStats.Out), 122 | inRate + "/s", 123 | outRate + "/s", 124 | } 125 | 126 | items = append(items, item) 127 | } 128 | 129 | table.AppendBulk(items) 130 | table.Render() 131 | } 132 | 133 | func humanBytes(n int64) string { 134 | return humanize.Bytes(uint64(n)) 135 | } 136 | 137 | func main() { 138 | log.Info("start netflow sniffer") 139 | 140 | start() 141 | stop() 142 | 143 | log.Info("netflow sniffer exit") 144 | } 145 | -------------------------------------------------------------------------------- /example/demo.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "time" 7 | 8 | "github.com/rfyiamcool/go-netflow" 9 | ) 10 | 11 | func main() { 12 | nf, err := netflow.New( 13 | netflow.WithCaptureTimeout(5 * time.Second), 14 | ) 15 | if err != nil { 16 | panic(err) 17 | } 18 | 19 | err = nf.Start() 20 | if err != nil { 21 | panic(err) 22 | } 23 | defer nf.Stop() 24 | 25 | <-nf.Done() 26 | 27 | var ( 28 | limit = 5 29 | recentSec = 5 30 | ) 31 | 32 | rank, err := nf.GetProcessRank(limit, recentSec) 33 | if err != nil { 34 | panic(err) 35 | } 36 | 37 | bs, err := json.MarshalIndent(rank, "", " ") 38 | if err != nil { 39 | panic(err) 40 | } 41 | 42 | fmt.Println(string(bs)) 43 | } 44 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/rfyiamcool/go-netflow 2 | 3 | go 1.16 4 | 5 | require ( 6 | github.com/containerd/cgroups v1.0.2 7 | github.com/dustin/go-humanize v1.0.0 8 | github.com/fatih/color v1.13.0 9 | github.com/google/gopacket v1.1.19 10 | github.com/mattn/go-runewidth v0.0.13 // indirect 11 | github.com/olekukonko/tablewriter v0.0.5 12 | github.com/opencontainers/runtime-spec v1.0.2 13 | github.com/sirupsen/logrus v1.8.1 14 | github.com/spf13/cast v1.4.1 15 | github.com/stretchr/testify v1.7.0 16 | golang.org/x/sync v0.0.0-20210220032951-036812b2e83c 17 | ) 18 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= 2 | github.com/cilium/ebpf v0.4.0/go.mod h1:4tRaxcgiL706VnOzHOdBlY8IEAIdxINsQBcU4xJJXRs= 3 | github.com/containerd/cgroups v1.0.2 h1:mZBclaSgNDfPWtfhj2xJY28LZ9nYIgzB0pwSURPl6JM= 4 | github.com/containerd/cgroups v1.0.2/go.mod h1:qpbpJ1jmlqsR9f2IyaLPsdkCdnt0rbDVqIDlhuu5tRY= 5 | github.com/coreos/go-systemd/v22 v22.3.2 h1:D9/bQk5vlXQFZ6Kwuu6zaiXJ9oTPe68++AzAJc1DzSI= 6 | github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= 7 | github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= 8 | github.com/cpuguy83/go-md2man/v2 v2.0.0/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU= 9 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 10 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 11 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 12 | github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw= 13 | github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= 14 | github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= 15 | github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= 16 | github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w= 17 | github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk= 18 | github.com/frankban/quicktest v1.11.3/go.mod h1:wRf/ReqHper53s+kmmSZizM8NamnL3IM0I9ntUbOk+k= 19 | github.com/godbus/dbus/v5 v5.0.4 h1:9349emZab16e7zQvpmsbtjc18ykshndd8y2PG3sgJbA= 20 | github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= 21 | github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= 22 | github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= 23 | github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 24 | github.com/google/gopacket v1.1.19 h1:ves8RnFZPGiFnTS0uPQStjwru6uO6h+nlr9j6fL7kF8= 25 | github.com/google/gopacket v1.1.19/go.mod h1:iJ8V8n6KS+z2U1A8pUwu8bW5SyEMkXJB8Yo/Vo+TKTo= 26 | github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= 27 | github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= 28 | github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= 29 | github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= 30 | github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= 31 | github.com/mattn/go-colorable v0.1.9 h1:sqDoxXbdeALODt0DAeJCVp38ps9ZogZEAXjus69YV3U= 32 | github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc= 33 | github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU= 34 | github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y= 35 | github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= 36 | github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI= 37 | github.com/mattn/go-runewidth v0.0.13 h1:lTGmDsbAYt5DmK6OnoV7EuIF1wEIFAcxld6ypU4OSgU= 38 | github.com/mattn/go-runewidth v0.0.13/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= 39 | github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec= 40 | github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY= 41 | github.com/opencontainers/runtime-spec v1.0.2 h1:UfAcuLBJB9Coz72x1hgl8O5RVzTdNiaglX6v2DM6FI0= 42 | github.com/opencontainers/runtime-spec v1.0.2/go.mod h1:jwyrGlmzljRJv/Fgzds9SsS/C5hL+LL3ko9hs6T5lQ0= 43 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 44 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 45 | github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY= 46 | github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= 47 | github.com/russross/blackfriday/v2 v2.0.1/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= 48 | github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc= 49 | github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE= 50 | github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0= 51 | github.com/spf13/cast v1.4.1 h1:s0hze+J0196ZfEMTs80N7UlFt0BDuQ7Q+JDnHiMWKdA= 52 | github.com/spf13/cast v1.4.1/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE= 53 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 54 | github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= 55 | github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 56 | github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= 57 | github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 58 | github.com/urfave/cli v1.22.2/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0= 59 | github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= 60 | github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= 61 | golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= 62 | golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= 63 | golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= 64 | golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= 65 | golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= 66 | golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= 67 | golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= 68 | golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= 69 | golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= 70 | golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= 71 | golang.org/x/net v0.0.0-20201021035429-f5854403a974 h1:IX6qOQeG5uLjB/hjjwjedwfjND0hgjPMMyO1RoIXQNI= 72 | golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= 73 | golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 74 | golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 75 | golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 76 | golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ= 77 | golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 78 | golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 79 | golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 80 | golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 81 | golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 82 | golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 83 | golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 84 | golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= 85 | golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c h1:F1jZWGFhYfh0Ci55sIpILtKKK8p3i2/krTr0H1rg74I= 86 | golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 87 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 88 | golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= 89 | golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= 90 | golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= 91 | golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28= 92 | golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= 93 | golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= 94 | golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 95 | golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 96 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 97 | golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 98 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 99 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 100 | gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= 101 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= 102 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 103 | -------------------------------------------------------------------------------- /netflow.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "net" 8 | "os" 9 | "strings" 10 | "sync/atomic" 11 | "time" 12 | 13 | "github.com/google/gopacket" 14 | "github.com/google/gopacket/layers" 15 | "github.com/google/gopacket/pcap" 16 | "github.com/google/gopacket/pcapgo" 17 | "golang.org/x/sync/errgroup" 18 | ) 19 | 20 | var ( 21 | errNotFound = errors.New("not found") 22 | ) 23 | 24 | type sideOption int 25 | 26 | const ( 27 | inputSide sideOption = iota 28 | outputSide 29 | ) 30 | 31 | type Netflow struct { 32 | ctx context.Context 33 | cancel context.CancelFunc 34 | 35 | connInodeHash *Mapping 36 | processHash *processController 37 | workerNum int 38 | qsize int 39 | 40 | // for update action 41 | delayQueue chan *delayEntry 42 | packetQueue chan gopacket.Packet 43 | 44 | bindIPs map[string]nullObject // read only 45 | bindDevices map[string]nullObject // read only 46 | counter int64 47 | captureTimeout time.Duration 48 | syncInterval time.Duration 49 | pcapFilter string // for pcap filter 50 | 51 | pcapFileName string 52 | pcapFile *os.File 53 | pcapWriter *pcapgo.Writer 54 | 55 | // for debug 56 | debugMode bool 57 | logger LoggerInterface 58 | 59 | // for cgroup 60 | cpuCore float64 61 | memMB int 62 | 63 | exitFunc []func() 64 | timer *time.Timer 65 | } 66 | 67 | type optionFunc func(*Netflow) error 68 | 69 | // WithPcapFilter set custom pcap filter 70 | // filter: "port 80", "src host xiaorui.cc and port 80" 71 | func WithPcapFilter(filter string) optionFunc { 72 | return func(o *Netflow) error { 73 | if len(filter) == 0 { 74 | return nil 75 | } 76 | 77 | st := strings.TrimSpace(filter) 78 | if strings.HasPrefix(st, "and") { 79 | return errors.New("invalid pcap filter") 80 | } 81 | 82 | o.pcapFilter = filter 83 | return nil 84 | } 85 | } 86 | 87 | func WithOpenDebug() optionFunc { 88 | return func(o *Netflow) error { 89 | o.debugMode = true 90 | return nil 91 | } 92 | } 93 | 94 | // WithLimitCgroup use cgroup to limit cpu and mem, param cpu's unit is cpu core num , mem's unit is MB 95 | func WithLimitCgroup(cpu float64, mem int) optionFunc { 96 | return func(o *Netflow) error { 97 | o.cpuCore = cpu 98 | o.memMB = mem 99 | return nil 100 | } 101 | } 102 | 103 | func WithLogger(logger LoggerInterface) optionFunc { 104 | return func(o *Netflow) error { 105 | o.logger = logger 106 | return nil 107 | } 108 | } 109 | 110 | func WithStorePcap(fpath string) optionFunc { 111 | return func(o *Netflow) error { 112 | o.pcapFileName = fpath 113 | return nil 114 | } 115 | } 116 | 117 | func WithCaptureTimeout(dur time.Duration) optionFunc { 118 | // capture name 119 | if dur > defaultCaptureTimeout { 120 | dur = defaultCaptureTimeout 121 | } 122 | 123 | return func(o *Netflow) error { 124 | o.captureTimeout = dur 125 | return nil 126 | } 127 | } 128 | 129 | func WithSyncInterval(dur time.Duration) optionFunc { 130 | return func(o *Netflow) error { 131 | if dur <= 0 { 132 | return errors.New("invalid sync interval") 133 | } 134 | 135 | o.syncInterval = dur 136 | return nil 137 | } 138 | } 139 | 140 | func WithWorkerNum(num int) optionFunc { 141 | if num <= 0 { 142 | num = defaultWorkerNum // default 143 | } 144 | 145 | return func(o *Netflow) error { 146 | o.workerNum = num 147 | return nil 148 | } 149 | } 150 | 151 | func WithCtx(ctx context.Context) optionFunc { 152 | return func(o *Netflow) error { 153 | cctx, cancel := context.WithCancel(ctx) 154 | o.ctx = cctx 155 | o.cancel = cancel 156 | return nil 157 | } 158 | } 159 | 160 | func WithBindIPs(ips []string) optionFunc { 161 | return func(o *Netflow) error { 162 | if len(ips) == 0 { 163 | return errors.New("invalid ips") 164 | } 165 | 166 | mm := make(map[string]nullObject, 10) 167 | for _, ip := range ips { 168 | mm[ip] = nullObject{} 169 | } 170 | 171 | o.bindIPs = mm 172 | return nil 173 | } 174 | } 175 | 176 | func WithBindDevices(devs []string) optionFunc { 177 | return func(o *Netflow) error { 178 | if len(devs) == 0 { 179 | return errors.New("invalid devs") 180 | } 181 | 182 | mm := make(map[string]nullObject, 10) 183 | for _, dev := range devs { 184 | mm[dev] = nullObject{} 185 | } 186 | 187 | o.bindDevices = mm 188 | return nil 189 | } 190 | } 191 | 192 | func WithQueueSize(size int) optionFunc { 193 | if size < 1000 { 194 | size = defaultQueueSize 195 | } 196 | 197 | return func(o *Netflow) error { 198 | o.qsize = size 199 | return nil 200 | } 201 | } 202 | 203 | const ( 204 | defaultQueueSize = 20000 // 2w 205 | defaultWorkerNum = 1 // usually one worker is enough. 206 | defaultSyncInterval = time.Duration(1 * time.Second) 207 | defaultCaptureTimeout = time.Duration(300 * time.Second) 208 | ) 209 | 210 | type Interface interface { 211 | // start netflow 212 | Start() error 213 | 214 | // stop netflow 215 | Stop() 216 | 217 | // sum packet 218 | LoadCounter() int64 219 | 220 | // when ctx.cancel() or timeout, notify done. 221 | Done() <-chan struct{} 222 | 223 | // GetProcessRank 224 | // param limit, size of data returned. 225 | // param recentSeconds, the average of the last few seconds' value. 226 | GetProcessRank(limit int, recentSeconds int) ([]*Process, error) 227 | } 228 | 229 | func New(opts ...optionFunc) (Interface, error) { 230 | var ( 231 | ctx, cancel = context.WithCancel(context.Background()) 232 | ) 233 | 234 | ips, devs := parseIpaddrsAndDevices() 235 | nf := &Netflow{ 236 | ctx: ctx, 237 | cancel: cancel, 238 | bindIPs: ips, 239 | bindDevices: devs, 240 | qsize: defaultQueueSize, 241 | workerNum: defaultWorkerNum, 242 | captureTimeout: defaultCaptureTimeout, 243 | syncInterval: defaultSyncInterval, 244 | debugMode: false, 245 | logger: &logger{}, 246 | } 247 | 248 | for _, opt := range opts { 249 | err := opt(nf) 250 | if err != nil { 251 | return nil, err 252 | } 253 | } 254 | 255 | nf.packetQueue = make(chan gopacket.Packet, nf.qsize) 256 | nf.delayQueue = make(chan *delayEntry, nf.qsize) 257 | nf.connInodeHash = NewMapping() 258 | nf.processHash = NewProcessController(nf.ctx) 259 | 260 | return nf, nil 261 | } 262 | 263 | func (nf *Netflow) Done() <-chan struct{} { 264 | return nf.ctx.Done() 265 | } 266 | 267 | func (nf *Netflow) GetProcessRank(limit int, recentSeconds int) ([]*Process, error) { 268 | if recentSeconds > maxRingSize { 269 | return nil, errors.New("windows interval must <= 15") 270 | } 271 | 272 | nf.processHash.Sort(recentSeconds) 273 | prank := nf.processHash.GetRank(limit) 274 | return prank, nil 275 | } 276 | 277 | func (nf *Netflow) incrCounter() { 278 | atomic.AddInt64(&nf.counter, 1) 279 | } 280 | 281 | func (nf *Netflow) LoadCounter() int64 { 282 | return atomic.LoadInt64(&nf.counter) 283 | } 284 | 285 | func (nf *Netflow) configureCgroups() error { 286 | if nf.cpuCore == 0 && nf.memMB == 0 { 287 | return nil 288 | } 289 | 290 | cg := cgroupsLimiter{} 291 | pid := os.Getpid() 292 | 293 | err := cg.configure(pid, nf.cpuCore, nf.memMB) 294 | nf.exitFunc = append(nf.exitFunc, func() { 295 | cg.free() 296 | }) 297 | 298 | return err 299 | } 300 | 301 | func (nf *Netflow) configurePersist() error { 302 | if len(nf.pcapFileName) == 0 { 303 | return nil 304 | } 305 | 306 | f, err := os.Create(nf.pcapFileName) 307 | if err != nil { 308 | return err 309 | } 310 | 311 | nf.pcapFile = f 312 | nf.pcapWriter = pcapgo.NewWriter(f) 313 | nf.pcapWriter.WriteFileHeader(1024, layers.LinkTypeEthernet) 314 | return nil 315 | } 316 | 317 | func (nf *Netflow) Start() error { 318 | var err error 319 | err = nf.configurePersist() 320 | if err != nil { 321 | return err 322 | } 323 | 324 | // linux cpu/mem by cgroup 325 | err = nf.configureCgroups() 326 | if err != nil { 327 | return err 328 | } 329 | 330 | // start workers 331 | go nf.startResourceSyncer() 332 | go nf.startNetworkSniffer() 333 | 334 | return nil 335 | } 336 | 337 | func (nf *Netflow) Stop() { 338 | nf.cancel() 339 | nf.finalize() 340 | 341 | if nf.pcapFile != nil { 342 | nf.pcapFile.Close() 343 | } 344 | } 345 | 346 | func (nf *Netflow) finalize() { 347 | if nf.timer != nil { 348 | nf.timer.Stop() 349 | } 350 | 351 | for _, fn := range nf.exitFunc { 352 | fn() 353 | } 354 | } 355 | 356 | func (nf *Netflow) startResourceSyncer() { 357 | var ( 358 | ticker = time.NewTicker(nf.syncInterval) 359 | entry *delayEntry 360 | lastTime time.Time 361 | ) 362 | 363 | // first run at the beginning 364 | nf.rescanResouce() 365 | 366 | for { 367 | select { 368 | case <-nf.ctx.Done(): 369 | return 370 | case <-ticker.C: 371 | nf.rescanResouce() 372 | lastTime = time.Now() 373 | 374 | // after rescan, handle undo entries 375 | for { 376 | if entry == nil { 377 | entry = nf.consumeDelayQueue() 378 | } 379 | // queue is empty 380 | if entry == nil { 381 | break 382 | } 383 | 384 | // only hanlde entry before rescan. 385 | if entry.timestamp.After(lastTime) { 386 | break 387 | } 388 | 389 | nf.handleDelayEntry(entry) 390 | entry = nil 391 | } 392 | } 393 | } 394 | } 395 | 396 | func (nf *Netflow) rescanResouce() error { 397 | var wg errgroup.Group 398 | 399 | wg.Go(func() error { 400 | return nf.rescanConns() 401 | }) 402 | wg.Go(func() error { 403 | return nf.rescanProcessInodes() 404 | }) 405 | 406 | return wg.Wait() 407 | } 408 | 409 | func (nf *Netflow) rescanProcessInodes() error { 410 | return nf.processHash.Rescan() 411 | } 412 | 413 | func (nf *Netflow) rescanConns() error { 414 | conns, err := netstat("tcp") 415 | if err != nil { 416 | return err 417 | } 418 | 419 | for _, conn := range conns { 420 | nf.connInodeHash.Add(conn.Addr, conn.Inode) 421 | nf.connInodeHash.Add(conn.ReverseAddr, conn.Inode) 422 | } 423 | 424 | return nil 425 | } 426 | 427 | func (nf *Netflow) captureDevice(dev string) { 428 | handler, err := buildPcapHandler(dev, nf.captureTimeout, nf.pcapFilter) 429 | if err != nil { 430 | return 431 | } 432 | 433 | defer func() { 434 | handler.Close() 435 | }() 436 | 437 | packetSource := gopacket.NewPacketSource( 438 | handler, 439 | handler.LinkType(), 440 | ) 441 | 442 | for { 443 | select { 444 | case <-nf.ctx.Done(): 445 | return 446 | 447 | case pkt := <-packetSource.Packets(): 448 | nf.enqueue(pkt) 449 | } 450 | } 451 | } 452 | 453 | func (nf *Netflow) enqueue(pkt gopacket.Packet) { 454 | select { 455 | case nf.packetQueue <- pkt: 456 | nf.incrCounter() 457 | return 458 | default: 459 | nf.logError("queue overflow, current size: ", len(nf.packetQueue)) 460 | } 461 | } 462 | 463 | func (nf *Netflow) dequeue() gopacket.Packet { 464 | select { 465 | case pkt := <-nf.packetQueue: 466 | return pkt 467 | 468 | case <-nf.ctx.Done(): 469 | return nil 470 | } 471 | } 472 | 473 | func (nf *Netflow) loopHandlePacket() { 474 | for { 475 | pkt := nf.dequeue() 476 | if pkt == nil { 477 | return // ctx.Done 478 | } 479 | 480 | nf.handlePacket(pkt) 481 | } 482 | } 483 | 484 | func (nf *Netflow) handlePacket(packet gopacket.Packet) { 485 | // var ( 486 | // ethLayer layers.Ethernet 487 | // ipLayer layers.IPv4 488 | // tcpLayer layers.TCP 489 | 490 | // layerTypes = []gopacket.LayerType{} 491 | // ) 492 | 493 | // parser := gopacket.NewDecodingLayerParser( 494 | // layers.LayerTypeEthernet, 495 | // ðLayer, 496 | // &ipLayer, 497 | // &tcpLayer, 498 | // ) 499 | 500 | // err := parser.DecodeLayers(packet.Data(), &layerTypes) 501 | // if err != nil { 502 | // continue 503 | // } 504 | 505 | // get ipLayer 506 | _ipLayer := packet.Layer(layers.LayerTypeIPv4) 507 | if _ipLayer == nil { 508 | return 509 | } 510 | ipLayer, _ := _ipLayer.(*layers.IPv4) 511 | 512 | // get tcpLayer 513 | _tcpLayer := packet.Layer(layers.LayerTypeTCP) 514 | if _tcpLayer == nil { 515 | return 516 | } 517 | tcpLayer, _ := _tcpLayer.(*layers.TCP) 518 | 519 | var ( 520 | side sideOption 521 | 522 | localIP = ipLayer.SrcIP 523 | localPort = tcpLayer.SrcPort 524 | 525 | remoteIP = ipLayer.DstIP 526 | remotePort = tcpLayer.DstPort 527 | ) 528 | 529 | if nf.isBindIPs(ipLayer.SrcIP.String()) { 530 | side = outputSide 531 | } else { 532 | side = inputSide 533 | } 534 | 535 | // length := len(packet.Data()) // ip header + tcp header + tcp payload 536 | length := len(tcpLayer.Payload) 537 | addr := spliceAddr(localIP, localPort, remoteIP, remotePort) 538 | nf.increaseTraffic(addr, int64(length), side) 539 | 540 | if nf.pcapFile != nil { 541 | nf.pcapWriter.WritePacket(packet.Metadata().CaptureInfo, packet.Data()) 542 | } 543 | 544 | // fmt.Println(">>>>", addr, len(packet.Data()), len(tcpLayer.Payload), side) 545 | } 546 | 547 | func (nf *Netflow) logDebug(msg ...interface{}) { 548 | if !nf.debugMode { 549 | return 550 | } 551 | 552 | nf.logger.Debug(msg...) 553 | } 554 | 555 | func (nf *Netflow) logError(msg ...interface{}) { 556 | if !nf.debugMode { 557 | return 558 | } 559 | 560 | nf.logger.Error(msg...) 561 | } 562 | 563 | func (nf *Netflow) isBindIPs(ipa string) bool { 564 | _, ok := nf.bindIPs[ipa] 565 | return ok 566 | } 567 | 568 | func (nf *Netflow) startNetworkSniffer() { 569 | for dev := range nf.bindDevices { 570 | go nf.captureDevice(dev) 571 | } 572 | 573 | for i := 0; i < nf.workerNum; i++ { 574 | go nf.loopHandlePacket() 575 | } 576 | 577 | nf.timer = time.AfterFunc(nf.captureTimeout, 578 | func() { 579 | nf.Stop() 580 | }, 581 | ) 582 | } 583 | 584 | type delayEntry struct { 585 | // meta 586 | timestamp time.Time 587 | times int 588 | 589 | // data 590 | addr string 591 | length int64 592 | side sideOption 593 | } 594 | 595 | func (nf *Netflow) pushDelayQueue(de *delayEntry) { 596 | select { 597 | case nf.delayQueue <- de: 598 | default: 599 | // if q is full, drain actively . 600 | } 601 | } 602 | 603 | func (nf *Netflow) consumeDelayQueue() *delayEntry { 604 | select { 605 | case <-nf.ctx.Done(): 606 | return nil 607 | 608 | case den := <-nf.delayQueue: 609 | return den 610 | 611 | default: 612 | return nil 613 | } 614 | } 615 | 616 | func (nf *Netflow) handleDelayEntry(entry *delayEntry) error { 617 | proc, err := nf.getProcessByAddr(entry.addr) 618 | if err != nil { 619 | return err 620 | } 621 | 622 | nf.increaseProcessTraffic(proc, entry.length, entry.side) 623 | return nil 624 | } 625 | 626 | func (nf *Netflow) getProcessByAddr(addr string) (*Process, error) { 627 | inode := nf.connInodeHash.Get(addr) 628 | if len(inode) == 0 { 629 | // not found, to rescan 630 | nf.logDebug("not found inode ", addr) 631 | return nil, errNotFound 632 | } 633 | 634 | proc := nf.processHash.GetProcessByInode(inode) 635 | if proc == nil { 636 | // not found, to rescan 637 | nf.logDebug("not found proc ", addr) 638 | return nil, errNotFound 639 | } 640 | 641 | return proc, nil 642 | } 643 | 644 | func (nf *Netflow) increaseProcessTraffic(proc *Process, length int64, side sideOption) error { 645 | switch side { 646 | case inputSide: 647 | proc.IncreaseInput(length) 648 | case outputSide: 649 | proc.IncreaseOutput(length) 650 | } 651 | return nil 652 | } 653 | 654 | func (nf *Netflow) increaseTraffic(addr string, length int64, side sideOption) error { 655 | proc, err := nf.getProcessByAddr(addr) 656 | if err != nil { 657 | den := &delayEntry{ 658 | timestamp: time.Now(), 659 | times: 0, 660 | addr: addr, 661 | length: length, 662 | side: side, 663 | } 664 | nf.pushDelayQueue(den) 665 | return err 666 | } 667 | 668 | nf.increaseProcessTraffic(proc, length, side) 669 | return nil 670 | } 671 | 672 | type nullObject = struct{} 673 | 674 | func parseIpaddrsAndDevices() (map[string]nullObject, map[string]nullObject) { 675 | devs, err := pcap.FindAllDevs() 676 | if err != nil { 677 | return nil, nil 678 | } 679 | 680 | var ( 681 | bindIPs = map[string]nullObject{} 682 | devNames = map[string]nullObject{} 683 | ) 684 | 685 | for _, dev := range devs { 686 | for _, addr := range dev.Addresses { 687 | if addr.IP.IsMulticast() { 688 | continue 689 | } 690 | bindIPs[addr.IP.String()] = struct{}{} 691 | } 692 | 693 | if strings.HasPrefix(dev.Name, "eth") { 694 | devNames[dev.Name] = nullObject{} 695 | continue 696 | } 697 | 698 | if strings.HasPrefix(dev.Name, "em") { 699 | devNames[dev.Name] = nullObject{} 700 | continue 701 | } 702 | 703 | if strings.HasPrefix(dev.Name, "lo") { 704 | devNames[dev.Name] = nullObject{} 705 | continue 706 | } 707 | 708 | if strings.HasPrefix(dev.Name, "bond") { 709 | devNames[dev.Name] = nullObject{} 710 | continue 711 | } 712 | } 713 | return bindIPs, devNames 714 | } 715 | 716 | func buildPcapHandler(device string, timeout time.Duration, pfilter string) (*pcap.Handle, error) { 717 | var ( 718 | snapshotLen int32 = 65536 719 | promisc bool = false 720 | ) 721 | 722 | // if packet captured size >= snapshotLength or 1 second's timer is expired, call user layer. 723 | handler, err := pcap.OpenLive(device, snapshotLen, promisc, time.Second) 724 | if err != nil { 725 | return nil, err 726 | } 727 | 728 | var filter = "tcp and (not broadcast and not multicast)" 729 | if len(pfilter) != 0 { 730 | filter = fmt.Sprintf("%s and %s", filter, pfilter) 731 | } 732 | 733 | err = handler.SetBPFFilter(filter) 734 | if err != nil { 735 | return nil, err 736 | } 737 | 738 | return handler, nil 739 | } 740 | 741 | func spliceAddr(sip net.IP, sport layers.TCPPort, dip net.IP, dport layers.TCPPort) string { 742 | return fmt.Sprintf("%s:%d_%s:%d", sip, sport, dip, dport) 743 | } 744 | -------------------------------------------------------------------------------- /netflow_test.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | func TestCgroupRun(t *testing.T) { 11 | nf, err := New(WithLimitCgroup(0.5, 0)) 12 | assert.Equal(t, nil, err) 13 | nf.Start() 14 | 15 | time.Sleep(60 * time.Second) 16 | nf.Stop() 17 | time.Sleep(60 * time.Second) 18 | } 19 | -------------------------------------------------------------------------------- /netstat.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "fmt" 5 | "io/ioutil" 6 | "os" 7 | "strconv" 8 | "strings" 9 | "time" 10 | ) 11 | 12 | const ( 13 | procTCPFile = "/proc/net/tcp" 14 | procUDPFile = "/proc/net/udp" 15 | procTCP6File = "/proc/net/tcp6" 16 | procUDP6File = "/proc/net/udp6" 17 | 18 | EstablishedSymbol = "01" 19 | ListenSymbol = "0A" 20 | ) 21 | 22 | type ConnectionItem struct { 23 | Addr string `json:"addr" valid:"-"` 24 | ReverseAddr string `json:"reverse_addr" valid:"-"` 25 | SrcIP string `json:"ip"` 26 | SrcPort string `json:"port"` 27 | DestIP string `json:"foreignip"` 28 | DestPort string `json:"foreignport"` 29 | State string `json:"state"` 30 | 31 | TxQueue int `json:"tx_queue" valid:"-"` 32 | RxQueue int `json:"rx_queue" valid:"-"` 33 | Timer int8 `json:"timer" valid:"-"` 34 | TimerDuration time.Duration `json:"timer_duration" valid:"-"` 35 | Rto time.Duration // retransmission timeout 36 | Uid int 37 | Uname string 38 | Timeout time.Duration 39 | Inode string `json:"inode"` 40 | Raw string `json:"raw"` 41 | } 42 | 43 | func (ci *ConnectionItem) GetAddr() string { 44 | return ci.Addr 45 | } 46 | 47 | func parseNetworkLines(tp string) ([]string, error) { 48 | var pf string 49 | 50 | switch tp { 51 | case "tcp": 52 | pf = procTCPFile 53 | case "udp": 54 | pf = procUDPFile 55 | case "tcp6": 56 | pf = procTCP6File 57 | case "udp6": 58 | pf = procUDP6File 59 | default: 60 | pf = procTCPFile 61 | } 62 | 63 | data, err := ioutil.ReadFile(pf) 64 | if err != nil { 65 | return nil, err 66 | } 67 | 68 | lines := strings.Split(string(data), "\n") 69 | return lines[1 : len(lines)-1], nil 70 | } 71 | 72 | func hex2dec(hexstr string) string { 73 | i, _ := strconv.ParseInt(hexstr, 16, 0) 74 | return strconv.FormatInt(i, 10) 75 | } 76 | 77 | func hex2ip(hexstr string) (string, string) { 78 | var ip string 79 | if len(hexstr) != 8 { 80 | err := "parse error" 81 | return ip, err 82 | } 83 | 84 | i1, _ := strconv.ParseInt(hexstr[6:8], 16, 0) 85 | i2, _ := strconv.ParseInt(hexstr[4:6], 16, 0) 86 | i3, _ := strconv.ParseInt(hexstr[2:4], 16, 0) 87 | i4, _ := strconv.ParseInt(hexstr[0:2], 16, 0) 88 | ip = fmt.Sprintf("%d.%d.%d.%d", i1, i2, i3, i4) 89 | 90 | return ip, "" 91 | } 92 | 93 | func parseAddr(str string) (string, string) { 94 | l := strings.Split(str, ":") 95 | if len(l) != 2 { 96 | return str, "" 97 | } 98 | 99 | ip, err := hex2ip(l[0]) 100 | if err != "" { 101 | return str, "" 102 | } 103 | 104 | return ip, hex2dec(l[1]) 105 | } 106 | 107 | // convert hexadecimal to decimal. 108 | func hexToDec(h string) int64 { 109 | d, err := strconv.ParseInt(h, 16, 32) 110 | if err != nil { 111 | fmt.Println(err) 112 | os.Exit(1) 113 | } 114 | 115 | return d 116 | } 117 | 118 | // remove empty data from line 119 | func removeEmpty(array []string) []string { 120 | var columns []string 121 | for _, i := range array { 122 | if i == "" { 123 | continue 124 | } 125 | columns = append(columns, i) 126 | } 127 | return columns 128 | } 129 | 130 | type filterFunc func() 131 | 132 | func netstat(t string) ([]*ConnectionItem, error) { 133 | var ( 134 | conns []*ConnectionItem 135 | ) 136 | 137 | data, err := parseNetworkLines(t) 138 | if err != nil { 139 | return nil, err 140 | } 141 | 142 | for _, line := range data { 143 | pp := getConnectionItem(line) 144 | if pp == nil { 145 | continue 146 | } 147 | 148 | conns = append(conns, pp) 149 | } 150 | 151 | return conns, nil 152 | } 153 | 154 | func getConnectionItem(line string) *ConnectionItem { 155 | // local ip and port 156 | source := removeEmpty(strings.Split(strings.TrimSpace(line), " ")) 157 | 158 | // only notice ESTAB and listen state 159 | if source[3] != EstablishedSymbol && source[3] != ListenSymbol { 160 | return nil 161 | } 162 | 163 | // ignore local listenning records 164 | destIP, destPort := parseAddr(source[2]) 165 | if destIP == "0.0.0.0" { 166 | return nil 167 | } 168 | 169 | // source ip and port 170 | ip, port := parseAddr(source[1]) 171 | 172 | // connection info 173 | stateNum, _ := strconv.ParseInt(source[3], 16, 32) 174 | state := states[int(stateNum)] 175 | 176 | // parse tx, rx queue size 177 | tcpQueue := strings.Split(source[4], ":") 178 | txq, err := strconv.ParseInt(tcpQueue[0], 16, 32) // tx queue size 179 | if err != nil { 180 | return nil 181 | } 182 | 183 | rxq, err := strconv.ParseInt(tcpQueue[1], 16, 32) // rx queue size 184 | if err != nil { 185 | return nil 186 | } 187 | 188 | // socket uid 189 | uid, err := strconv.Atoi(source[7]) 190 | if err != nil { 191 | return nil 192 | } 193 | 194 | // get user name by uid 195 | uname := getUserByUID(source[7]) 196 | 197 | // socket inode 198 | inode := source[9] 199 | 200 | // tcp 4 fileds 201 | addr := ip + ":" + port + "_" + destIP + ":" + destPort 202 | raddr := destIP + ":" + destPort + "_" + ip + ":" + port 203 | 204 | cc := &ConnectionItem{ 205 | Addr: addr, 206 | ReverseAddr: raddr, 207 | State: state, 208 | SrcIP: ip, 209 | SrcPort: port, 210 | DestIP: destIP, 211 | DestPort: destPort, 212 | Inode: inode, 213 | TxQueue: int(txq), 214 | RxQueue: int(rxq), 215 | Uid: uid, 216 | Uname: uname, 217 | Raw: line, 218 | } 219 | return cc 220 | } 221 | 222 | // Tcp func Get a slice of Process type with TCP data 223 | func Tcp() []*ConnectionItem { 224 | data, _ := netstat("tcp") 225 | return data 226 | } 227 | 228 | // Udp func Get a slice of Process type with UDP data 229 | func Udp() []*ConnectionItem { 230 | data, _ := netstat("udp") 231 | return data 232 | } 233 | 234 | // Tcp6 func Get a slice of Process type with TCP6 data 235 | func Tcp6() []*ConnectionItem { 236 | data, _ := netstat("tcp6") 237 | return data 238 | } 239 | 240 | // Udp6 func Get a slice of Process type with UDP6 data 241 | func Udp6() []*ConnectionItem { 242 | data, _ := netstat("udp6") 243 | return data 244 | } 245 | -------------------------------------------------------------------------------- /netstat_test.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "os" 7 | "testing" 8 | 9 | "github.com/stretchr/testify/assert" 10 | ) 11 | 12 | // func TestGetInfos(t *testing.T) { 13 | // conns, _ := netstat("tcp") 14 | // bs, err := json.MarshalIndent(conns, "", " ") 15 | // assert.Equal(t, nil, err) 16 | // t.Log(string(bs)) 17 | 18 | // for _, conn := range conns { 19 | // ConnInodeHash.Add(conn.Addr, conn.Inode) 20 | // } 21 | // fmt.Println("conns", ConnInodeHash.String()) 22 | 23 | // ps, err := GetProcesses() 24 | // assert.Equal(t, nil, err) 25 | // for _, p := range ps { 26 | // for _, inode := range p.Inodes { 27 | // InodePidsHash.Add(inode, p.Pid) 28 | // } 29 | // } 30 | 31 | // fmt.Println("processes", InodePidsHash.String()) 32 | // } 33 | 34 | func TestParseAddr(t *testing.T) { 35 | ip, port := parseAddr("2E0010AC:E898") 36 | 37 | assert.Equal(t, ip, "172.16.0.46") 38 | assert.Equal(t, port, "59544") 39 | } 40 | 41 | func TestHandleFile(t *testing.T) { 42 | for _, fdn := range []string{"0", "1", "2", "3", "4"} { 43 | name := "/proc/994998/fd/" + fdn 44 | 45 | kk, err := os.Readlink(name) 46 | fmt.Println(kk, err) 47 | 48 | // 只是针对 unix socket 49 | // fileinfo, err := os.Lstat(name) 50 | // assert.Equal(t, nil, err) 51 | // name = "/tmp/mongodb-27017.sock" 52 | // fmt.Println("socket: ", fileinfo.Mode()&os.ModeSocket == os.ModeSocket) 53 | // fmt.Println("symlink: ", fileinfo.Mode()&os.ModeSymlink == os.ModeSymlink) 54 | 55 | // if fileinfo.Mode()&os.ModeSocket == os.ModeSocket { 56 | // fmt.Println("this is socket") 57 | // } 58 | // if fileinfo.Mode() == os.ModeSocket { 59 | // fmt.Println("this is socket") 60 | // } 61 | // fmt.Println(fileinfo.Mode().String()) 62 | 63 | // stat, _ := fileinfo.Sys().(*syscall.Stat_t) 64 | 65 | // fmt.Println("socket222: ", stat.Mode, fileinfo.Name()) 66 | // fmt.Println("socket333: ", fileinfo.Mode()) 67 | // fmt.Println("socket444: ", KindFromFileInfo(fileinfo) == SOCKET) 68 | // fmt.Println("socket555: ", fileinfo.Mode()&os.ModeSocket) 69 | } 70 | } 71 | 72 | // https://www.cyub.vip/2020/11/22/Go%E8%AF%AD%E8%A8%80%E5%AE%9E%E7%8E%B0%E7%AE%80%E6%98%93%E7%89%88netstat%E5%91%BD%E4%BB%A4/ 73 | 74 | type Kind int 75 | 76 | const ( 77 | DIR Kind = iota 78 | LINK 79 | PIPE 80 | SOCKET 81 | DEV 82 | FILE 83 | ) 84 | 85 | func KindFromFileInfo(fileInfo os.FileInfo) Kind { 86 | if fileInfo.IsDir() { 87 | return DIR 88 | } 89 | 90 | if fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink { 91 | return LINK 92 | } 93 | 94 | if fileInfo.Mode()&os.ModeNamedPipe == os.ModeNamedPipe { 95 | return PIPE 96 | } 97 | 98 | if fileInfo.Mode()&os.ModeSocket == os.ModeSocket { 99 | return SOCKET 100 | } 101 | 102 | if fileInfo.Mode()&os.ModeDevice == os.ModeDevice { 103 | return DEV 104 | } 105 | 106 | return FILE 107 | } 108 | 109 | func (self Kind) MarshalJSON() ([]byte, error) { 110 | return json.Marshal(self.String()) 111 | } 112 | 113 | func (self *Kind) UnmarshalJSON(data []byte) error { 114 | var s string 115 | if err := json.Unmarshal(data, &s); err != nil { 116 | return err 117 | } 118 | switch s { 119 | case "DIR": 120 | *self = DIR 121 | case "LINK": 122 | *self = LINK 123 | case "PIPE": 124 | *self = PIPE 125 | case "SOCKET": 126 | *self = SOCKET 127 | case "DEV": 128 | *self = DEV 129 | case "FILE": 130 | *self = FILE 131 | default: 132 | return fmt.Errorf("invalid Kind: '%s'", s) 133 | } 134 | return nil 135 | } 136 | 137 | func (self Kind) String() string { 138 | names := []string{ 139 | "DIR", "LINK", "PIPE", "SOCKET", "DEV", "FILE", 140 | } 141 | 142 | // FIXME: bound check? 143 | return names[self] 144 | } 145 | -------------------------------------------------------------------------------- /process.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "fmt" 7 | "os" 8 | "path/filepath" 9 | "regexp" 10 | "sort" 11 | "strings" 12 | "sync" 13 | "time" 14 | ) 15 | 16 | var ( 17 | maxRingSize = 15 18 | ) 19 | 20 | type Process struct { 21 | Name string `json:"name"` 22 | Pid string `json:"pid"` 23 | Exe string `json:"exe"` 24 | State string `json:"state"` 25 | InodeCount int `json:"inode_count"` 26 | TrafficStats *trafficStatsEntry `json:"traffic_stats"` 27 | 28 | // todo: use ringbuffer array to reduce gc cost. 29 | Ring []*trafficEntry `json:"ring"` 30 | 31 | inodes []string 32 | revision int 33 | } 34 | 35 | func (p *Process) getLastTrafficEntry() *trafficEntry { 36 | if len(p.Ring) == 0 { 37 | return nil 38 | } 39 | return p.Ring[len(p.Ring)-1] 40 | } 41 | 42 | func (p *Process) analyseStats(sec int) { 43 | var ( 44 | stats = new(trafficStatsEntry) 45 | thold = time.Now().Add(-time.Duration(sec) * time.Second).Unix() 46 | ) 47 | 48 | // avoid x / 0 to raise exception 49 | if sec == 0 { 50 | return 51 | } 52 | 53 | for _, item := range p.Ring { 54 | if item.Timestamp < thold { 55 | continue 56 | } 57 | stats.In += item.In 58 | stats.Out += item.Out 59 | } 60 | 61 | stats.InRate = stats.In / int64(sec) 62 | stats.OutRate = stats.Out / int64(sec) 63 | p.TrafficStats = stats 64 | } 65 | 66 | func (po *Process) shrink() { 67 | if len(po.Ring) >= maxRingSize { 68 | po.Ring = po.Ring[1:] // reduce size 69 | } 70 | } 71 | 72 | func (po *Process) IncreaseInput(n int64) { 73 | now := time.Now().Unix() 74 | if len(po.Ring) == 0 { 75 | item := &trafficEntry{ 76 | Timestamp: now, 77 | In: n, 78 | } 79 | po.Ring = append(po.Ring, item) 80 | return 81 | } 82 | 83 | po.shrink() 84 | 85 | item := po.Ring[len(po.Ring)-1] 86 | if item.Timestamp == now { 87 | item.In += n 88 | return 89 | } 90 | 91 | item = &trafficEntry{ 92 | Timestamp: now, 93 | In: n, 94 | } 95 | po.Ring = append(po.Ring, item) 96 | } 97 | 98 | // IncreaseOutput 99 | func (po *Process) IncreaseOutput(n int64) { 100 | // todo: format code 101 | now := time.Now().Unix() 102 | if len(po.Ring) == 0 { 103 | item := &trafficEntry{ 104 | Timestamp: now, 105 | Out: n, 106 | } 107 | po.Ring = append(po.Ring, item) 108 | return 109 | } 110 | 111 | po.shrink() 112 | 113 | item := po.Ring[len(po.Ring)-1] 114 | if item.Timestamp == now { 115 | item.Out += n 116 | return 117 | } 118 | 119 | item = &trafficEntry{ 120 | Timestamp: now, 121 | Out: n, 122 | } 123 | po.Ring = append(po.Ring, item) 124 | } 125 | 126 | func (p *Process) copy() *Process { 127 | return &Process{ 128 | Name: p.Name, 129 | Pid: p.Pid, 130 | Exe: p.Exe, 131 | State: p.State, 132 | InodeCount: p.InodeCount, 133 | TrafficStats: &trafficStatsEntry{ 134 | In: p.TrafficStats.In, 135 | Out: p.TrafficStats.Out, 136 | InRate: p.TrafficStats.InRate, 137 | OutRate: p.TrafficStats.OutRate, 138 | }, 139 | Ring: p.Ring, 140 | } 141 | } 142 | 143 | type trafficEntry struct { 144 | Timestamp int64 `json:"timestamp"` 145 | In int64 `json:"in"` 146 | Out int64 `json:"out"` 147 | } 148 | 149 | type trafficStatsEntry struct { 150 | In int64 `json:"in"` 151 | Out int64 `json:"out"` 152 | InRate int64 `json:"in_rate"` 153 | OutRate int64 `json:"out_rate"` 154 | InputEWMA int64 `json:"input_ewma" valid:"-"` 155 | OutputEWMA int64 `json:"output_ewma" valid:"-"` 156 | } 157 | 158 | func GetProcesses() (map[string]*Process, error) { 159 | // to improve performance 160 | files, err := filepath.Glob("/proc/[0-9]*/fd/[0-9]*") 161 | if err != nil { 162 | return nil, err 163 | } 164 | 165 | var ( 166 | ppm = make(map[string]*Process, 1000) 167 | label = "socket:[" 168 | ) 169 | 170 | for _, fpath := range files { 171 | rules := []string{"fd/0", "fd/1", "fd/2"} 172 | if matchStringSuffix(fpath, rules) { 173 | continue 174 | } 175 | 176 | name, _ := os.Readlink(fpath) 177 | 178 | if !strings.HasPrefix(name, label) { 179 | continue 180 | } 181 | 182 | var ( 183 | pid = strings.Split(fpath, "/")[2] 184 | inode = name[len(label) : len(name)-1] 185 | ) 186 | 187 | po := ppm[pid] 188 | if po != nil { // has 189 | po.inodes = append(po.inodes, inode) 190 | po.InodeCount = len(po.inodes) 191 | continue 192 | } 193 | 194 | exe := getProcessExe(pid) 195 | pname := getProcessName(exe) 196 | ppm[pid] = &Process{ 197 | Pid: pid, 198 | inodes: []string{inode}, 199 | InodeCount: 1, 200 | Name: pname, 201 | Exe: exe, 202 | TrafficStats: new(trafficStatsEntry), 203 | } 204 | } 205 | 206 | return ppm, nil 207 | } 208 | 209 | type processController struct { 210 | sync.RWMutex 211 | 212 | ctx context.Context 213 | cancel context.CancelFunc 214 | 215 | // key -> pid, val -> process 216 | dict map[string]*Process 217 | revision int 218 | 219 | // key -> inode_num, val -> pid_num 220 | inodePidMap map[string]string 221 | 222 | // cache 223 | sortedProcesses sortedProcesses 224 | } 225 | 226 | func NewProcessController(ctx context.Context) *processController { 227 | var ( 228 | size = 1000 229 | ) 230 | 231 | cctx, cancel := context.WithCancel(ctx) 232 | return &processController{ 233 | ctx: cctx, 234 | cancel: cancel, 235 | dict: make(map[string]*Process, size), 236 | inodePidMap: make(map[string]string, size), 237 | } 238 | } 239 | 240 | func (pm *processController) GetRank(limit int) []*Process { 241 | pm.RLock() 242 | defer pm.RUnlock() 243 | 244 | src := pm.sortedProcesses 245 | if len(src) > limit { 246 | src = pm.sortedProcesses[:limit] 247 | } 248 | 249 | // copy object 250 | res := []*Process{} 251 | for _, item := range src { 252 | res = append(res, item.copy()) 253 | } 254 | return src 255 | } 256 | 257 | func (pm *processController) Sort(sec int) []*Process { 258 | pm.RLock() 259 | defer pm.RUnlock() 260 | 261 | pos := sortedProcesses{} 262 | for _, po := range pm.dict { 263 | po.analyseStats(sec) 264 | pos = append(pos, po) 265 | } 266 | 267 | sort.Sort(pos) 268 | pm.sortedProcesses = pos 269 | 270 | return pos 271 | } 272 | 273 | func (pm *processController) Add(pid string, p *Process) { 274 | pm.Lock() 275 | defer pm.Unlock() 276 | 277 | pm.dict[pid] = p 278 | } 279 | 280 | func (pm *processController) Get(pid string) *Process { 281 | pm.RLock() 282 | defer pm.RUnlock() 283 | 284 | return pm.dict[pid] 285 | } 286 | 287 | func (pm *processController) GetProcessByInode(inode string) *Process { 288 | pm.RLock() 289 | defer pm.RUnlock() 290 | 291 | pid, ok := pm.inodePidMap[inode] 292 | if !ok { 293 | return nil 294 | } 295 | 296 | return pm.dict[pid] 297 | } 298 | 299 | func (pm *processController) delete(pid string) { 300 | pm.Lock() 301 | defer pm.Unlock() 302 | 303 | delete(pm.dict, pid) 304 | } 305 | 306 | func (pm *processController) readIterator(fn func(*Process)) { 307 | pm.RLock() 308 | defer pm.RUnlock() 309 | 310 | for _, po := range pm.dict { 311 | fn(po) 312 | } 313 | } 314 | 315 | func (pm *processController) anyIterator(fn func(*Process)) { 316 | pm.Lock() 317 | defer pm.Unlock() 318 | 319 | for _, po := range pm.dict { 320 | fn(po) 321 | } 322 | } 323 | 324 | func (pm *processController) copy() map[string]*Process { 325 | ndict := make(map[string]*Process, len(pm.dict)) 326 | 327 | pm.RLock() 328 | defer pm.RUnlock() 329 | 330 | for k, v := range pm.dict { 331 | ndict[k] = v 332 | } 333 | return ndict 334 | } 335 | 336 | func (pm *processController) AsyncRun() { 337 | go pm.Run() 338 | } 339 | 340 | func (pm *processController) Run() { 341 | var ( 342 | interval = 5 * time.Second 343 | ticker = time.NewTicker(interval) 344 | ) 345 | 346 | pm.Rescan() 347 | 348 | for { 349 | select { 350 | case <-pm.ctx.Done(): 351 | return 352 | 353 | case <-ticker.C: 354 | pm.Rescan() 355 | } 356 | } 357 | } 358 | 359 | func (pm *processController) Stop() { 360 | pm.cancel() 361 | } 362 | 363 | func (pm *processController) sortNetflow() string { 364 | bs, _ := json.MarshalIndent(pm.dict, "", " ") 365 | return string(bs) 366 | } 367 | 368 | func (pm *processController) analyse() error { 369 | pm.RLock() 370 | defer pm.RUnlock() 371 | 372 | for pid, po := range pm.dict { 373 | fmt.Println(pid, po) 374 | } 375 | 376 | return nil 377 | } 378 | 379 | func (pm *processController) Rescan() error { 380 | ps, err := GetProcesses() 381 | if err != nil { 382 | return err 383 | } 384 | 385 | pm.Lock() 386 | defer pm.Unlock() 387 | 388 | pm.revision++ 389 | 390 | // add new pid 391 | for pid, po := range ps { 392 | pp, ok := pm.dict[pid] 393 | if ok { 394 | pp.inodes = po.inodes 395 | continue // alread exist 396 | } 397 | 398 | pm.dict[pid] = po 399 | } 400 | 401 | // del old pid 402 | for pid, _ := range pm.dict { 403 | _, ok := ps[pid] 404 | if ok { 405 | continue 406 | } 407 | 408 | delete(pm.dict, pid) 409 | } 410 | 411 | // inode -> pid 412 | inodePidMap := make(map[string]string, 1000) 413 | for pid, po := range ps { 414 | for _, inode := range po.inodes { 415 | inodePidMap[inode] = pid 416 | } 417 | } 418 | pm.inodePidMap = inodePidMap // obj reset 419 | 420 | return nil 421 | } 422 | 423 | func (pm *processController) Reset() { 424 | pm.dict = make(map[string]*Process, 1000) 425 | pm.inodePidMap = make(map[string]string, 1000) 426 | } 427 | 428 | // getProcessExe 429 | func getProcessExe(pid string) string { 430 | exe := fmt.Sprintf("/proc/%s/exe", pid) 431 | path, _ := os.Readlink(exe) 432 | return path 433 | } 434 | 435 | // getProcessName 436 | func getProcessName(exe string) string { 437 | n := strings.Split(exe, "/") 438 | name := n[len(n)-1] 439 | return strings.Title(name) 440 | } 441 | 442 | // findPid unuse 443 | func findPid(inode string) string { 444 | pid := "-" 445 | 446 | d, err := filepath.Glob("/proc/[0-9]*/fd/[0-9]*") 447 | if err != nil { 448 | fmt.Println(err) 449 | os.Exit(1) 450 | } 451 | 452 | re := regexp.MustCompile(inode) 453 | for _, item := range d { 454 | path, _ := os.Readlink(item) 455 | out := re.FindString(path) 456 | if len(out) != 0 { 457 | pid = strings.Split(item, "/")[2] 458 | } 459 | } 460 | return pid 461 | } 462 | 463 | type sortedProcesses []*Process 464 | 465 | func (s sortedProcesses) Len() int { 466 | return len(s) 467 | } 468 | 469 | func (s sortedProcesses) Less(i, j int) bool { 470 | val1 := s[i].TrafficStats.In + s[i].TrafficStats.Out 471 | val2 := s[j].TrafficStats.In + s[j].TrafficStats.Out 472 | return val1 > val2 473 | } 474 | 475 | func (s sortedProcesses) Swap(i, j int) { 476 | s[i], s[j] = s[j], s[i] 477 | } 478 | -------------------------------------------------------------------------------- /process_test.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "fmt" 7 | "sort" 8 | "testing" 9 | "time" 10 | 11 | "github.com/stretchr/testify/assert" 12 | ) 13 | 14 | func TestStringSuffix(t *testing.T) { 15 | cur := "/proc/123123/fd/3" 16 | b := matchStringSuffix(cur, []string{"fd/0", "fd/1", "fd2"}) 17 | assert.Equal(t, false, b) 18 | 19 | cur = "/proc/123123/fd/2" 20 | b = matchStringSuffix(cur, []string{"fd/0", "fd/1", "fd/2"}) 21 | assert.Equal(t, true, b) 22 | 23 | cur = "/proc/123123/fd/11" 24 | b = matchStringSuffix(cur, []string{"fd/0", "fd/1", "fd2"}) 25 | assert.Equal(t, false, b) 26 | } 27 | 28 | func TestGetProcesses(t *testing.T) { 29 | pps, err := GetProcesses() 30 | assert.Equal(t, nil, err) 31 | 32 | bs, err := json.MarshalIndent(pps, "", " ") 33 | assert.Equal(t, nil, err) 34 | 35 | fmt.Printf("%+v", string(bs)) 36 | } 37 | 38 | func TestProcessHash(t *testing.T) { 39 | pm := NewProcessController(context.Background()) 40 | go pm.Run() 41 | 42 | time.Sleep(3 * time.Second) 43 | 44 | t.Log(MarshalIndent(pm.dict)) 45 | t.Log(MarshalIndent(pm.inodePidMap)) 46 | pm.Stop() 47 | } 48 | 49 | func TestProcessShrink(t *testing.T) { 50 | po := Process{} 51 | for i := 0; i < 15; i++ { 52 | po.IncreaseInput(10) 53 | time.Sleep(1 * time.Second) 54 | } 55 | assert.Equal(t, 10, len(po.Ring)) 56 | } 57 | 58 | func TestProcessAnalyse(t *testing.T) { 59 | po := Process{} 60 | po.IncreaseInput(10) 61 | time.Sleep(1 * time.Second) 62 | 63 | po.IncreaseInput(50) 64 | time.Sleep(1 * time.Second) 65 | 66 | po.IncreaseInput(100) 67 | time.Sleep(1 * time.Second) 68 | 69 | // in 70 | po.IncreaseInput(50) 71 | po.IncreaseInput(50) 72 | po.IncreaseInput(50) 73 | // out 74 | po.IncreaseOutput(50) 75 | po.IncreaseOutput(50) 76 | time.Sleep(1 * time.Second) 77 | 78 | assert.EqualValues(t, 150, po.getLastTrafficEntry().In) 79 | assert.EqualValues(t, 100, po.getLastTrafficEntry().Out) 80 | 81 | t.Log(MarshalIndent(po)) 82 | 83 | po.analyseStats(2) 84 | t.Log(MarshalIndent(po)) 85 | } 86 | 87 | func TestProcessAnalyse2(t *testing.T) { 88 | po := Process{} 89 | 90 | for i := 0; i < 20; i++ { 91 | po.IncreaseInput(10) 92 | time.Sleep(1 * time.Second) 93 | } 94 | 95 | po.analyseStats(2) 96 | assert.EqualValues(t, 20, po.TrafficStats.In) 97 | } 98 | 99 | func TestSortedProcesses(t *testing.T) { 100 | pps := []*Process{} 101 | 102 | p1 := &Process{ 103 | Pid: "111", 104 | TrafficStats: &trafficStatsEntry{ 105 | In: 100, 106 | Out: 100, 107 | InRate: 100, 108 | OutRate: 100, 109 | }, 110 | } 111 | 112 | p2 := &Process{ 113 | Pid: "222", 114 | TrafficStats: &trafficStatsEntry{ 115 | In: 200, 116 | Out: 200, 117 | InRate: 200, 118 | OutRate: 200, 119 | }, 120 | } 121 | 122 | p3 := &Process{ 123 | Pid: "333", 124 | TrafficStats: &trafficStatsEntry{ 125 | In: 300, 126 | Out: 300, 127 | InRate: 300, 128 | OutRate: 300, 129 | }, 130 | } 131 | 132 | p4 := &Process{ 133 | Pid: "444", 134 | TrafficStats: &trafficStatsEntry{ 135 | In: 400, 136 | Out: 400, 137 | InRate: 400, 138 | OutRate: 400, 139 | }, 140 | } 141 | 142 | // don't insert in order ! 143 | pps = append(pps, p4) 144 | pps = append(pps, p1) 145 | pps = append(pps, p3) 146 | pps = append(pps, p2) 147 | 148 | sort.Sort(sortedProcesses(pps)) 149 | 150 | t.Log(MarshalIndent(pps)) 151 | 152 | link := []*Process{p4, p3, p2, p1} // desc sort 153 | for idx := range link { 154 | assert.Equal(t, link[idx].Pid, pps[idx].Pid) 155 | } 156 | } 157 | -------------------------------------------------------------------------------- /ringqueue.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "errors" 5 | "sync" 6 | ) 7 | 8 | var ( 9 | ErrQueueFull = errors.New("queue is full") 10 | ) 11 | 12 | type ringQueue struct { 13 | sync.RWMutex 14 | buf []interface{} 15 | head, tail, count int 16 | } 17 | 18 | func newRingQueue(size int) *ringQueue { 19 | return &ringQueue{ 20 | buf: make([]interface{}, size), 21 | } 22 | } 23 | 24 | func (q *ringQueue) Length() int { 25 | return q.count 26 | } 27 | 28 | func (q *ringQueue) Add(elem interface{}) error { 29 | q.Lock() 30 | defer q.Unlock() 31 | 32 | if q.count == len(q.buf) { 33 | return ErrQueueFull 34 | } 35 | 36 | q.count++ 37 | q.buf[q.tail] = elem 38 | 39 | if q.tail+1 < len(q.buf) { 40 | q.tail++ 41 | } 42 | if len(q.buf) == q.count { 43 | q.tail = 0 44 | } 45 | return nil 46 | } 47 | 48 | func (q *ringQueue) Peek() interface{} { 49 | q.RLock() 50 | defer q.RUnlock() 51 | 52 | if q.count <= 0 { 53 | return nil 54 | } 55 | return q.buf[q.head] 56 | } 57 | 58 | func (q *ringQueue) Remove() interface{} { 59 | q.Lock() 60 | defer q.Unlock() 61 | 62 | if q.count <= 0 { 63 | return nil 64 | } 65 | 66 | ret := q.buf[q.head] 67 | q.buf[q.head] = nil 68 | 69 | q.head = (q.head + 1) & (len(q.buf) - 1) 70 | q.count-- 71 | 72 | return ret 73 | } 74 | -------------------------------------------------------------------------------- /types.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "encoding/json" 5 | "fmt" 6 | "sync" 7 | ) 8 | 9 | const ( 10 | TCP_ESTABLISHED = iota + 1 11 | TCP_SYN_SENT 12 | TCP_SYN_RECV 13 | TCP_FIN_WAIT1 14 | TCP_FIN_WAIT2 15 | TCP_TIME_WAIT 16 | TCP_CLOSE 17 | TCP_CLOSE_WAIT 18 | TCP_LAST_ACK 19 | TCP_LISTEN 20 | TCP_CLOSING 21 | //TCP_NEW_SYN_RECV 22 | //TCP_MAX_STATES 23 | ) 24 | 25 | var states = map[int]string{ 26 | TCP_ESTABLISHED: "ESTABLISHED", 27 | TCP_SYN_SENT: "SYN_SENT", 28 | TCP_SYN_RECV: "SYN_RECV", 29 | TCP_FIN_WAIT1: "FIN_WAIT1", 30 | TCP_FIN_WAIT2: "FIN_WAIT2", 31 | TCP_TIME_WAIT: "TIME_WAIT", 32 | TCP_CLOSE: "CLOSE", 33 | TCP_CLOSE_WAIT: "CLOSE_WAIT", 34 | TCP_LAST_ACK: "LAST_ACK", 35 | TCP_LISTEN: "LISTEN", 36 | TCP_CLOSING: "CLOSING", 37 | //TCP_NEW_SYN_RECV: "NEW_SYN_RECV", 38 | //TCP_MAX_STATES: "MAX_STATES", 39 | } 40 | 41 | // https://github.com/torvalds/linux/blob/master/include/net/tcp_states.h 42 | var StateMapping = map[string]string{ 43 | "01": "ESTABLISHED", 44 | "02": "SYN_SENT", 45 | "03": "SYN_RECV", 46 | "04": "FIN_WAIT1", 47 | "05": "FIN_WAIT2", 48 | "06": "TIME_WAIT", 49 | "07": "CLOSE", 50 | "08": "CLOSE_WAIT", 51 | "09": "LAST_ACK", 52 | "0A": "LISTEN", 53 | "0B": "CLOSING", 54 | } 55 | 56 | type Mapping struct { 57 | cb func() 58 | dict map[string]string 59 | sync.RWMutex 60 | } 61 | 62 | func NewMapping() *Mapping { 63 | size := 1000 64 | return &Mapping{ 65 | dict: make(map[string]string, size), 66 | } 67 | } 68 | 69 | func (m *Mapping) Handle() { 70 | } 71 | 72 | func (m *Mapping) Add(k, v string) { 73 | m.Lock() 74 | defer m.Unlock() 75 | 76 | m.dict[k] = v 77 | } 78 | 79 | func (m *Mapping) Get(k string) string { 80 | m.RLock() 81 | defer m.RUnlock() 82 | 83 | return m.dict[k] 84 | } 85 | 86 | func (m *Mapping) Delete(k string) { 87 | m.Lock() 88 | defer m.Unlock() 89 | 90 | delete(m.dict, k) 91 | } 92 | 93 | func (m *Mapping) String() string { 94 | m.RLock() 95 | defer m.RUnlock() 96 | 97 | bs, _ := json.Marshal(m.dict) 98 | return string(bs) 99 | } 100 | 101 | type Null struct{} 102 | 103 | type LoggerInterface interface { 104 | Debug(...interface{}) 105 | Error(...interface{}) 106 | } 107 | 108 | var defaultLogger string 109 | 110 | type logger struct{} 111 | 112 | func (l *logger) Debug(msg ...interface{}) { 113 | fmt.Println(msg...) 114 | } 115 | 116 | func (l *logger) Error(msg ...interface{}) { 117 | fmt.Println(msg...) 118 | } 119 | -------------------------------------------------------------------------------- /user.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "bufio" 5 | "io" 6 | "os" 7 | "strings" 8 | ) 9 | 10 | var ( 11 | systemUsers map[string]string 12 | ) 13 | 14 | func init() { 15 | loadSystemUsersInfo() 16 | } 17 | 18 | func loadSystemUsersInfo() { 19 | f, err := os.Open("/etc/passwd") 20 | if err != nil { 21 | return 22 | } 23 | defer f.Close() 24 | 25 | systemUsers = make(map[string]string) 26 | bf := bufio.NewReader(f) 27 | for { 28 | line, err := bf.ReadString('\n') 29 | if err != nil { 30 | if err == io.EOF { 31 | break 32 | } 33 | } 34 | 35 | items := strings.Split(line, ":") 36 | if len(items) != 2 { 37 | return 38 | } 39 | 40 | systemUsers[items[2]] = items[0] 41 | } 42 | } 43 | 44 | func getUserByUID(uid string) string { 45 | return systemUsers[uid] 46 | } 47 | -------------------------------------------------------------------------------- /utils.go: -------------------------------------------------------------------------------- 1 | package netflow 2 | 3 | import ( 4 | "encoding/json" 5 | "strings" 6 | ) 7 | 8 | func matchStringSuffix(s string, mv []string) bool { 9 | for _, val := range mv { 10 | if strings.HasSuffix(s, val) { 11 | return true 12 | } 13 | } 14 | return false 15 | } 16 | 17 | func MarshalIndent(v interface{}) string { 18 | bs, _ := json.MarshalIndent(v, "", " ") 19 | return string(bs) 20 | } 21 | --------------------------------------------------------------------------------