diff options
| -rw-r--r-- | cmd/redq/main.go | 14 | ||||
| -rw-r--r-- | usage/main.go | 47 |
2 files changed, 44 insertions, 17 deletions
diff --git a/cmd/redq/main.go b/cmd/redq/main.go index df275d6..abac395 100644 --- a/cmd/redq/main.go +++ b/cmd/redq/main.go @@ -14,10 +14,7 @@ import ( ) func main() { - u := &usage.Usage { - Data : make(usage.UsageMap), - } - + var u usage.Usage iface, err := net.InterfaceByName("wlan0") if err != nil { @@ -32,14 +29,15 @@ func main() { defer conn.Close(ctx) queries := db.New(conn) + err = u.Init(iface) + if err != nil { + os.Exit(0) + } sigs := make(chan os.Signal, 1) signal.Notify(sigs, os.Interrupt, os.Kill, syscall.SIGTERM) go func () { <- sigs - err := u.UpdateDb(queries, ctx, false) - if err != nil { - log.Printf("updating Database: %s", err) - } + u.CleanUp(queries, ctx) os.Exit(0) }() diff --git a/usage/main.go b/usage/main.go index 2293281..8be8ed3 100644 --- a/usage/main.go +++ b/usage/main.go @@ -28,34 +28,63 @@ type Usage struct { Data UsageMap Mutex sync.Mutex objs bpfObjects + egressLink, ingressLink link.Link } -func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) { +func (u *Usage) Init(iface *net.Interface) error { + var err error + if err := loadBpfObjects(&u.objs, nil); err != nil { - log.Fatalf("loading objects: %s", err) + log.Printf("loading objects: %s", err) + return err; } - defer u.objs.Close() + defer func() { + if err != nil { + u.objs.Close() + } + }() - ingressLink, err := link.AttachTCX(link.TCXOptions{ + u.ingressLink, err = link.AttachTCX(link.TCXOptions{ Interface: iface.Index, Program: u.objs.IngressFunc, Attach: ebpf.AttachTCXIngress, }) if err != nil { - log.Fatalf("could not attach TCx program: %s", err) + log.Printf("could not attach TCx program: %s", err) + return err } - defer ingressLink.Close() + defer func() { + if err != nil { + u.ingressLink.Close() + } + }() - egressLink, err := link.AttachTCX(link.TCXOptions{ + u.egressLink, err = link.AttachTCX(link.TCXOptions{ Interface: iface.Index, Program: u.objs.EgressFunc, Attach: ebpf.AttachTCXEgress, }) if err != nil { - log.Fatalf("could not attach TCx program: %s", err) + log.Printf("could not attach TCx program: %s", err) + return err } - defer egressLink.Close() + u.Data = make(UsageMap) + return nil +} + +func (u *Usage) CleanUp(queries *db.Queries, ctxDb context.Context) { + err := u.UpdateDb(queries, ctxDb, false) + if err != nil { + log.Printf("updating Database: %s", err) + } + + u.objs.Close() + u.ingressLink.Close() + u.egressLink.Close() +} + +func (u *Usage) Run(iface *net.Interface, queries *db.Queries, ctxDb context.Context) { bpfTicker := time.NewTicker(time.Second) defer bpfTicker.Stop() dbTicker := time.NewTicker(time.Minute) |