From 541d676b986f01e00791b820f19402f64a9f1ca5 Mon Sep 17 00:00:00 2001 From: Juan Font Alonso Date: Thu, 8 Apr 2021 23:57:31 +0200 Subject: [PATCH] Minor code reorg --- app.go | 114 ----------------------------------------------------- cli.go | 121 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 114 deletions(-) create mode 100644 cli.go diff --git a/app.go b/app.go index 75fd1f7d..3466a396 100644 --- a/app.go +++ b/app.go @@ -1,15 +1,11 @@ package headscale import ( - "encoding/json" "fmt" - "log" "os" "sync" "github.com/gin-gonic/gin" - "github.com/jinzhu/gorm/dialects/postgres" - "inet.af/netaddr" "tailscale.com/tailcfg" "tailscale.com/wgengine/wgcfg" ) @@ -75,113 +71,3 @@ func (h *Headscale) Serve() error { err := r.Run(h.cfg.Addr) return err } - -// RegisterMachine is executed from the CLI to register a new Machine using its MachineKey -func (h *Headscale) RegisterMachine(key string, namespace string) error { - ns, err := h.GetNamespace(namespace) - if err != nil { - return err - } - mKey, err := wgcfg.ParseHexKey(key) - if err != nil { - log.Printf("Cannot parse client key: %s", err) - return err - } - db, err := h.db() - if err != nil { - log.Printf("Cannot open DB: %s", err) - return err - } - defer db.Close() - m := Machine{} - if db.First(&m, "machine_key = ?", mKey.HexString()).RecordNotFound() { - log.Printf("Cannot find machine with machine key: %s", mKey.Base64()) - return err - } - - if m.isAlreadyRegistered() { - fmt.Println("This machine already registered") - return nil - } - - ip, err := h.getAvailableIP() - if err != nil { - log.Println(err) - return err - } - m.IPAddress = ip.String() - m.NamespaceID = ns.ID - m.Registered = true - db.Save(&m) - fmt.Println("Machine registered 🎉") - return nil -} - -func (h *Headscale) ListNodeRoutes(namespace string, nodeName string) error { - m, err := h.GetMachine(namespace, nodeName) - if err != nil { - return err - } - - hi, err := m.GetHostInfo() - if err != nil { - return err - } - fmt.Println(hi.RoutableIPs) - return nil -} - -func (h *Headscale) EnableNodeRoute(namespace string, nodeName string, routeStr string) error { - m, err := h.GetMachine(namespace, nodeName) - if err != nil { - return err - } - hi, err := m.GetHostInfo() - if err != nil { - return err - } - route, err := netaddr.ParseIPPrefix(routeStr) - if err != nil { - return err - } - - for _, rIP := range hi.RoutableIPs { - if rIP == route { - db, err := h.db() - if err != nil { - log.Printf("Cannot open DB: %s", err) - return err - } - defer db.Close() - routes, _ := json.Marshal([]string{routeStr}) // TODO: only one for the time being, so overwriting the rest - m.EnabledRoutes = postgres.Jsonb{RawMessage: json.RawMessage(routes)} - db.Save(&m) - db.Close() - - peers, _ := h.getPeers(*m) - h.pollMu.Lock() - for _, p := range *peers { - if pUp, ok := h.clientsPolling[uint64(p.ID)]; ok { - pUp <- []byte{} - } else { - } - } - h.pollMu.Unlock() - return nil - } - } - return fmt.Errorf("Could not find routable range") - -} - -func eqCIDRs(a, b []netaddr.IPPrefix) bool { - if len(a) != len(b) || ((a == nil) != (b == nil)) { - return false - } - for i, v := range a { - if v != b[i] { - return false - } - } - return true -} diff --git a/cli.go b/cli.go new file mode 100644 index 00000000..90acb427 --- /dev/null +++ b/cli.go @@ -0,0 +1,121 @@ +package headscale + +import ( + "encoding/json" + "fmt" + "log" + + "github.com/jinzhu/gorm/dialects/postgres" + "inet.af/netaddr" + "tailscale.com/wgengine/wgcfg" +) + +// RegisterMachine is executed from the CLI to register a new Machine using its MachineKey +func (h *Headscale) RegisterMachine(key string, namespace string) error { + ns, err := h.GetNamespace(namespace) + if err != nil { + return err + } + mKey, err := wgcfg.ParseHexKey(key) + if err != nil { + log.Printf("Cannot parse client key: %s", err) + return err + } + db, err := h.db() + if err != nil { + log.Printf("Cannot open DB: %s", err) + return err + } + defer db.Close() + m := Machine{} + if db.First(&m, "machine_key = ?", mKey.HexString()).RecordNotFound() { + log.Printf("Cannot find machine with machine key: %s", mKey.Base64()) + return err + } + + if m.isAlreadyRegistered() { + fmt.Println("This machine already registered") + return nil + } + + ip, err := h.getAvailableIP() + if err != nil { + log.Println(err) + return err + } + m.IPAddress = ip.String() + m.NamespaceID = ns.ID + m.Registered = true + db.Save(&m) + fmt.Println("Machine registered 🎉") + return nil +} + +func (h *Headscale) ListNodeRoutes(namespace string, nodeName string) error { + m, err := h.GetMachine(namespace, nodeName) + if err != nil { + return err + } + + hi, err := m.GetHostInfo() + if err != nil { + return err + } + fmt.Println(hi.RoutableIPs) + return nil +} + +func (h *Headscale) EnableNodeRoute(namespace string, nodeName string, routeStr string) error { + m, err := h.GetMachine(namespace, nodeName) + if err != nil { + return err + } + hi, err := m.GetHostInfo() + if err != nil { + return err + } + route, err := netaddr.ParseIPPrefix(routeStr) + if err != nil { + return err + } + + for _, rIP := range hi.RoutableIPs { + if rIP == route { + db, err := h.db() + if err != nil { + log.Printf("Cannot open DB: %s", err) + return err + } + defer db.Close() + routes, _ := json.Marshal([]string{routeStr}) // TODO: only one for the time being, so overwriting the rest + m.EnabledRoutes = postgres.Jsonb{RawMessage: json.RawMessage(routes)} + db.Save(&m) + db.Close() + + peers, _ := h.getPeers(*m) + h.pollMu.Lock() + for _, p := range *peers { + if pUp, ok := h.clientsPolling[uint64(p.ID)]; ok { + pUp <- []byte{} + } else { + } + } + h.pollMu.Unlock() + return nil + } + } + return fmt.Errorf("Could not find routable range") + +} + +func eqCIDRs(a, b []netaddr.IPPrefix) bool { + if len(a) != len(b) || ((a == nil) != (b == nil)) { + return false + } + for i, v := range a { + if v != b[i] { + return false + } + } + return true +}