diff --git a/api.go b/api.go index dfda0c2a..e78bdb9f 100644 --- a/api.go +++ b/api.go @@ -18,10 +18,44 @@ import ( "tailscale.com/wgengine/wgcfg" ) +// KeyHandler provides the Headscale pub key +// Listens in /key func (h *Headscale) KeyHandler(c *gin.Context) { c.Data(200, "text/plain; charset=utf-8", []byte(h.publicKey.HexString())) } +// RegisterWebAPI shows a simple message in the browser to point to the CLI +// Listens in /register +func (h *Headscale) RegisterWebAPI(c *gin.Context) { + mKeyStr := c.Query("key") + if mKeyStr == "" { + c.String(http.StatusBadRequest, "Wrong params") + return + } + + c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(fmt.Sprintf(` + + +

headscale

+

+ Run the command below in the headscale server to add this machine to your network: +

+ +

+ + headscale register %s + +

+ + + + + `, mKeyStr))) + return +} + +// RegistrationHandler handles the actual registration process of a machine +// Endpoint /machine/:id func (h *Headscale) RegistrationHandler(c *gin.Context) { body, _ := io.ReadAll(c.Request.Body) mKeyStr := c.Param("id") @@ -59,6 +93,7 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) { if m.Registered { log.Println("Client is registered and we have the current key. All clear to /map") resp.AuthURL = "" + resp.User = *m.Namespace.toUser() respBody, err := encode(resp, &mKey, h.privateKey) if err != nil { log.Printf("Cannot encode message: %s", err) @@ -89,6 +124,7 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) { m.NodeKey = wgcfg.Key(req.NodeKey).HexString() db.Save(&m) resp.AuthURL = "" + resp.User = *m.Namespace.toUser() respBody, err := encode(resp, &mKey, h.privateKey) if err != nil { log.Printf("Cannot encode message: %s", err) @@ -318,35 +354,6 @@ func (h *Headscale) getMapKeepAliveResponse(mKey wgcfg.Key, req tailcfg.MapReque return &data, nil } -// RegisterWebAPI shows a simple message in the browser to point to the CLI -func (h *Headscale) RegisterWebAPI(c *gin.Context) { - mKeyStr := c.Query("key") - if mKeyStr == "" { - c.String(http.StatusBadRequest, "Wrong params") - return - } - - c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(fmt.Sprintf(` - - -

headscale

-

- Run the command below in the headscale server to add this machine to your network: -

- -

- - headscale register %s - -

- - - - - `, mKeyStr))) - return -} - func (h *Headscale) handleNewServer(c *gin.Context, db *gorm.DB, idKey wgcfg.Key, req tailcfg.RegisterRequest) { mNew := Machine{ MachineKey: idKey.HexString(), diff --git a/app.go b/app.go index a6c10ec0..b29ff985 100644 --- a/app.go +++ b/app.go @@ -74,7 +74,11 @@ func (h *Headscale) Serve() error { } // RegisterMachine is executed from the CLI to register a new Machine using its MachineKey -func (h *Headscale) RegisterMachine(key string) error { +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) @@ -103,6 +107,7 @@ func (h *Headscale) RegisterMachine(key string) error { return err } m.IPAddress = ip.String() + m.NamespaceID = ns.ID m.Registered = true db.Save(&m) fmt.Println("Machine registered 🎉") diff --git a/cmd/headscale/headscale.go b/cmd/headscale/headscale.go index 3b3ea123..b300c421 100644 --- a/cmd/headscale/headscale.go +++ b/cmd/headscale/headscale.go @@ -50,8 +50,31 @@ var serveCmd = &cobra.Command{ } var registerCmd = &cobra.Command{ - Use: "register machineID", + Use: "register machineID namespace", Short: "Registers a machine to your network", + Args: func(cmd *cobra.Command, args []string) error { + if len(args) < 2 { + return fmt.Errorf("Missing parameters") + } + return nil + }, + Run: func(cmd *cobra.Command, args []string) { + h, err := getHeadscaleApp() + if err != nil { + log.Fatalf("Error initializing: %s", err) + } + h.RegisterMachine(args[0], args[1]) + }, +} + +var namespaceCmd = &cobra.Command{ + Use: "namespace", + Short: "Manage the namespaces of Headscale", +} + +var createNamespaceCmd = &cobra.Command{ + Use: "create NAME", + Short: "Creates a new namespace", Args: func(cmd *cobra.Command, args []string) error { if len(args) < 1 { return fmt.Errorf("Missing parameters") @@ -63,7 +86,32 @@ var registerCmd = &cobra.Command{ if err != nil { log.Fatalf("Error initializing: %s", err) } - h.RegisterMachine(args[0]) + _, err = h.CreateNamespace(args[0]) + if err != nil { + fmt.Println(err) + return + } + fmt.Printf("Ook.\n") + }, +} + +var listNamespacesCmd = &cobra.Command{ + Use: "list", + Short: "Creates a new namespace", + Run: func(cmd *cobra.Command, args []string) { + h, err := getHeadscaleApp() + if err != nil { + log.Fatalf("Error initializing: %s", err) + } + ns, err := h.ListNamespaces() + if err != nil { + fmt.Println(err) + return + } + fmt.Printf("ID\tName\n") + for _, n := range *ns { + fmt.Printf("%d\t%s\n", n.ID, n.Name) + } }, } @@ -79,6 +127,9 @@ func main() { headscaleCmd.AddCommand(versionCmd) headscaleCmd.AddCommand(serveCmd) headscaleCmd.AddCommand(registerCmd) + headscaleCmd.AddCommand(namespaceCmd) + namespaceCmd.AddCommand(createNamespaceCmd) + namespaceCmd.AddCommand(listNamespacesCmd) if err := headscaleCmd.Execute(); err != nil { fmt.Println(err) diff --git a/db.go b/db.go index b9057829..7e1e0cc9 100644 --- a/db.go +++ b/db.go @@ -23,6 +23,7 @@ func (h *Headscale) initDB() error { db.Exec("create extension if not exists \"uuid-ossp\";") db.AutoMigrate(&Machine{}) db.AutoMigrate(&KV{}) + db.AutoMigrate(&Namespace{}) db.Close() h.setValue("db_version", dbVersion) diff --git a/machine.go b/machine.go index 648478ad..ea15cb27 100644 --- a/machine.go +++ b/machine.go @@ -16,12 +16,14 @@ import ( // Machine is a Headscale client type Machine struct { - ID uint64 `gorm:"primary_key"` - MachineKey string `gorm:"type:varchar(64);unique_index"` - NodeKey string - DiscoKey string - IPAddress string - Name string + ID uint64 `gorm:"primary_key"` + MachineKey string `gorm:"type:varchar(64);unique_index"` + NodeKey string + DiscoKey string + IPAddress string + Name string + NamespaceID uint + Namespace Namespace Registered bool // temp LastSeen *time.Time @@ -106,7 +108,7 @@ func (m Machine) toNode() (*tailcfg.Node, error) { ID: tailcfg.NodeID(m.ID), // this is the actual ID StableID: tailcfg.StableNodeID(strconv.FormatUint(m.ID, 10)), // in headscale, unlike tailcontrol server, IDs are permantent Name: hostinfo.Hostname, - User: 1, + User: tailcfg.UserID(m.NamespaceID), Key: tailcfg.NodeKey(nKey), KeyExpiry: *m.Expiry, Machine: tailcfg.MachineKey(mKey), @@ -136,9 +138,9 @@ func (h *Headscale) getPeers(m Machine) (*[]*tailcfg.Node, error) { } defer db.Close() - // Add user management here machines := []Machine{} - if err = db.Where("machine_key <> ? AND registered", m.MachineKey).Find(&machines).Error; err != nil { + if err = db.Where("namespace_id = ? AND machine_key <> ? AND registered", + m.NamespaceID, m.MachineKey).Find(&machines).Error; err != nil { log.Printf("Error accessing db: %s", err) return nil, err } diff --git a/namespaces.go b/namespaces.go new file mode 100644 index 00000000..4c9132b9 --- /dev/null +++ b/namespaces.go @@ -0,0 +1,122 @@ +package headscale + +import ( + "fmt" + "log" + "time" + + "github.com/jinzhu/gorm" + "tailscale.com/tailcfg" +) + +// Namespace is the way Headscale implements the concept of users in Tailscale +// +// At the end of the day, users in Tailscale are some kind of 'bubbles' or namespaces +// that contain our machines. +type Namespace struct { + gorm.Model + Name string `gorm:"unique"` +} + +// CreateNamespace creates a new Namespace. Returns error if could not be created +// or another namespace already exists +func (h *Headscale) CreateNamespace(name string) (*Namespace, error) { + db, err := h.db() + if err != nil { + log.Printf("Cannot open DB: %s", err) + return nil, err + } + defer db.Close() + + n := Namespace{} + if err := db.Where("name = ?", name).First(&n).Error; err == nil { + return nil, fmt.Errorf("Namespace already exists") + } + n.Name = name + if err := db.Create(&n).Error; err != nil { + log.Printf("Could not create row: %s", err) + return nil, err + } + return &n, nil +} + +// GetNamespace fetches a namespace by name +func (h *Headscale) GetNamespace(name string) (*Namespace, error) { + db, err := h.db() + if err != nil { + log.Printf("Cannot open DB: %s", err) + return nil, err + } + defer db.Close() + + n := Namespace{} + if db.First(&n, "name = ?", name).RecordNotFound() { + return nil, fmt.Errorf("Namespace not found") + } + return &n, nil +} + +// ListNamespaces gets all the existing namespaces +func (h *Headscale) ListNamespaces() (*[]Namespace, error) { + db, err := h.db() + if err != nil { + log.Printf("Cannot open DB: %s", err) + return nil, err + } + defer db.Close() + namespaces := []Namespace{} + if err := db.Find(&namespaces).Error; err != nil { + return nil, err + } + return &namespaces, nil +} + +// ListMachinesInNamespace gets all the nodes in a given namespace +func (h *Headscale) ListMachinesInNamespace(name string) (*[]Machine, error) { + n, err := h.GetNamespace(name) + if err != nil { + return nil, err + } + db, err := h.db() + if err != nil { + log.Printf("Cannot open DB: %s", err) + return nil, err + } + defer db.Close() + + machines := []Machine{} + if err := db.Where(&Machine{NamespaceID: n.ID}).Find(&machines).Error; err != nil { + return nil, err + } + return &machines, nil +} + +func (h *Headscale) SetMachineNamespace(m *Machine, namespaceName string) error { + n, err := h.GetNamespace(namespaceName) + if err != nil { + return err + } + db, err := h.db() + if err != nil { + log.Printf("Cannot open DB: %s", err) + return err + } + defer db.Close() + m.NamespaceID = n.ID + db.Save(&m) + return nil +} + +func (n *Namespace) toUser() *tailcfg.User { + u := tailcfg.User{ + ID: tailcfg.UserID(n.ID), + LoginName: "", + DisplayName: n.Name, + ProfilePicURL: "", + Domain: "", + Logins: []tailcfg.LoginID{}, + Roles: []tailcfg.RoleID{}, + Created: time.Time{}, + } + return &u +}