myriation/headscale
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages 1 Wiki Activity

Merge pull request #45 from juanfont/reuse-gorm-connection

Browse source 
Use gorm connection pool
This commit is contained in:
Juan Font 2021-07-04 21:51:43 +02:00 committed by GitHub
commit 90e9ad9a0e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 50 additions and 167 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
api.go
View file

@ -75,15 +75,8 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
return
}
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
c.String(http.StatusInternalServerError, ":(")
return
}
var m Machine
if result := db.First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
if result := h.db.First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
log.Println("New Machine!")
m = Machine{
Expiry: &req.Expiry,
@ -91,14 +84,14 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
Name: req.Hostinfo.Hostname,
NodeKey: wgkey.Key(req.NodeKey).HexString(),
}
if err := db.Create(&m).Error; err != nil {
if err := h.db.Create(&m).Error; err != nil {
log.Printf("Could not create row: %s", err)
return
}
}
if !m.Registered && req.Auth.AuthKey != "" {
h.handleAuthKey(c, db, mKey, req, m)
h.handleAuthKey(c, h.db, mKey, req, m)
return
}
@ -138,7 +131,7 @@ func (h *Headscale) RegistrationHandler(c *gin.Context) {
if m.NodeKey == wgkey.Key(req.OldNodeKey).HexString() {
log.Printf("[%s] We have the OldNodeKey in the database. This is a key refresh", m.Name)
m.NodeKey = wgkey.Key(req.NodeKey).HexString()
db.Save(&m)
h.db.Save(&m)
resp.AuthURL = ""
resp.User = *m.Namespace.toUser()
@ -204,13 +197,8 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
return
}
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return
}
var m Machine
if result := db.First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
if result := h.db.First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
log.Printf("Ignoring request, cannot find machine with key %s", mKey.HexString())
return
}
@ -234,7 +222,7 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
m.Endpoints = datatypes.JSON(endpoints)
m.LastSeen = &now
}
db.Save(&m)
h.db.Save(&m)
pollData := make(chan []byte, 1)
update := make(chan []byte, 1)
@ -303,7 +291,7 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
}
now := time.Now().UTC()
m.LastSeen = &now
db.Save(&m)
h.db.Save(&m)
return true
case <-update:
@ -322,7 +310,7 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
log.Printf("[%s] The client has closed the connection", m.Name)
now := time.Now().UTC()
m.LastSeen = &now
db.Save(&m)
h.db.Save(&m)
h.pollMu.Lock()
cancelKeepAlive <- []byte{}
delete(h.clientsPolling, m.ID)

11
app.go
View file

@ -12,6 +12,7 @@ import (
"github.com/gin-gonic/gin"
"golang.org/x/crypto/acme/autocert"
"gorm.io/gorm"
"tailscale.com/tailcfg"
"tailscale.com/types/wgkey"
)
@ -43,6 +44,7 @@ type Config struct {
// Headscale represents the base app of the service
type Headscale struct {
cfg Config
db *gorm.DB
dbString string
dbType string
dbDebug bool
@ -87,6 +89,7 @@ func NewHeadscale(cfg Config) (*Headscale, error) {
if err != nil {
return nil, err
}
h.clientsPolling = make(map[uint64]chan []byte)
return &h, nil
}
@ -107,12 +110,6 @@ func (h *Headscale) ExpireEphemeralNodes(milliSeconds int64) {
}
func (h *Headscale) expireEphemeralNodesWorker() {
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return
}
namespaces, err := h.ListNamespaces()
if err != nil {
log.Printf("Error listing namespaces: %s", err)
@ -127,7 +124,7 @@ func (h *Headscale) expireEphemeralNodesWorker() {
for _, m := range *machines {
if m.AuthKey != nil && m.LastSeen != nil && m.AuthKey.Ephemeral && time.Now().After(m.LastSeen.Add(h.cfg.EphemeralNodeInactivityTimeout)) {
log.Printf("[%s] Ephemeral client removed from database\n", m.Name)
err = db.Unscoped().Delete(m).Error
err = h.db.Unscoped().Delete(m).Error
if err != nil {
log.Printf("[%s] 🤮 Cannot delete ephemeral machine from the database: %s", m.Name, err)
}

5
app_test.go
View file

@ -47,4 +47,9 @@ func (s *Suite) ResetDB(c *check.C) {
if err != nil {
c.Fatal(err)
}
db, err := h.openDB()
if err != nil {
c.Fatal(err)
}
h.db = db
}

11
cli.go
View file

@ -2,7 +2,6 @@ package headscale
import (
"errors"
"log"
"gorm.io/gorm"
"tailscale.com/types/wgkey"
@ -18,13 +17,9 @@ func (h *Headscale) RegisterMachine(key string, namespace string) (*Machine, err
if err != nil {
return nil, err
}
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return nil, err
}
m := Machine{}
if result := db.First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
if result := h.db.First(&m, "machine_key = ?", mKey.HexString()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, errors.New("Machine not found")
}
@ -40,6 +35,6 @@ func (h *Headscale) RegisterMachine(key string, namespace string) (*Machine, err
m.NamespaceID = ns.ID
m.Registered = true
m.RegisterMethod = "cli"
db.Save(&m)
h.db.Save(&m)
return &m, nil
}

7
cli_test.go
View file

@ -8,11 +8,6 @@ func (s *Suite) TestRegisterMachine(c *check.C) {
n, err := h.CreateNamespace("test")
c.Assert(err, check.IsNil)
db, err := h.db()
if err != nil {
c.Fatal(err)
}
m := Machine{
ID: 0,
MachineKey: "8ce002a935f8c394e55e78fbbb410576575ff8ec5cfa2e627e4b807f1be15b0e",
@ -21,7 +16,7 @@ func (s *Suite) TestRegisterMachine(c *check.C) {
Name: "testmachine",
NamespaceID: n.ID,
}
db.Save(&m)
h.db.Save(&m)
_, err = h.GetMachine("test", "testmachine")
c.Assert(err, check.IsNil)

23
db.go
View file

@ -17,10 +17,12 @@ type KV struct {
}
func (h *Headscale) initDB() error {
db, err := h.db()
db, err := h.openDB()
if err != nil {
return err
}
h.db = db
if h.dbType == "postgres" {
db.Exec("create extension if not exists \"uuid-ossp\";")
}
@ -45,7 +47,7 @@ func (h *Headscale) initDB() error {
return err
}
func (h *Headscale) db() (*gorm.DB, error) {
func (h *Headscale) openDB() (*gorm.DB, error) {
var db *gorm.DB
var err error
switch h.dbType {
@ -69,12 +71,8 @@ func (h *Headscale) db() (*gorm.DB, error) {
}
func (h *Headscale) getValue(key string) (string, error) {
db, err := h.db()
if err != nil {
return "", err
}
var row KV
if result := db.First(&row, "key = ?", key); errors.Is(result.Error, gorm.ErrRecordNotFound) {
if result := h.db.First(&row, "key = ?", key); errors.Is(result.Error, gorm.ErrRecordNotFound) {
return "", errors.New("not found")
}
return row.Value, nil
@ -85,16 +83,13 @@ func (h *Headscale) setValue(key string, value string) error {
Key: key,
Value: value,
}
db, err := h.db()
if err != nil {
return err
}
_, err = h.getValue(key)
_, err := h.getValue(key)
if err == nil {
db.Model(&kv).Where("key = ?", key).Update("value", value)
h.db.Model(&kv).Where("key = ?", key).Update("value", value)
return nil
}
db.Create(kv)
h.db.Create(kv)
return nil
}

7
machine.go
View file

@ -154,14 +154,9 @@ func (m Machine) toNode() (*tailcfg.Node, error) {
}
func (h *Headscale) getPeers(m Machine) (*[]*tailcfg.Node, error) {
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return nil, err
}
machines := []Machine{}
if err = db.Where("namespace_id = ? AND machine_key <> ? AND registered",
if err := h.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

7
machine_test.go
View file

@ -11,11 +11,6 @@ func (s *Suite) TestGetMachine(c *check.C) {
pak, err := h.CreatePreAuthKey(n.Name, false, false, nil)
c.Assert(err, check.IsNil)
db, err := h.db()
if err != nil {
c.Fatal(err)
}
_, err = h.GetMachine("test", "testmachine")
c.Assert(err, check.NotNil)
@ -30,7 +25,7 @@ func (s *Suite) TestGetMachine(c *check.C) {
RegisterMethod: "authKey",
AuthKeyID: uint(pak.ID),
}
db.Save(&m)
h.db.Save(&m)
m1, err := h.GetMachine("test", "testmachine")
c.Assert(err, check.IsNil)

47
namespaces.go
View file

@ -25,18 +25,12 @@ type Namespace struct {
// 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
}
n := Namespace{}
if err := db.Where("name = ?", name).First(&n).Error; err == nil {
if err := h.db.Where("name = ?", name).First(&n).Error; err == nil {
return nil, errorNamespaceExists
}
n.Name = name
if err := db.Create(&n).Error; err != nil {
if err := h.db.Create(&n).Error; err != nil {
log.Printf("Could not create row: %s", err)
return nil, err
}
@ -46,12 +40,6 @@ func (h *Headscale) CreateNamespace(name string) (*Namespace, error) {
// DestroyNamespace destroys a Namespace. Returns error if the Namespace does
// not exist or if there are machines associated with it.
func (h *Headscale) DestroyNamespace(name string) error {
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return err
}
n, err := h.GetNamespace(name)
if err != nil {
return errorNamespaceNotFound
@ -65,7 +53,7 @@ func (h *Headscale) DestroyNamespace(name string) error {
return errorNamespaceNotEmpty
}
if result := db.Unscoped().Delete(&n); result.Error != nil {
if result := h.db.Unscoped().Delete(&n); result.Error != nil {
return err
}
@ -74,14 +62,8 @@ func (h *Headscale) DestroyNamespace(name string) error {
// 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
}
n := Namespace{}
if result := db.First(&n, "name = ?", name); errors.Is(result.Error, gorm.ErrRecordNotFound) {
if result := h.db.First(&n, "name = ?", name); errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, errorNamespaceNotFound
}
return &n, nil
@ -89,13 +71,8 @@ func (h *Headscale) GetNamespace(name string) (*Namespace, error) {
// 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
}
namespaces := []Namespace{}
if err := db.Find(&namespaces).Error; err != nil {
if err := h.db.Find(&namespaces).Error; err != nil {
return nil, err
}
return &namespaces, nil
@ -107,14 +84,9 @@ func (h *Headscale) ListMachinesInNamespace(name string) (*[]Machine, error) {
if err != nil {
return nil, err
}
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return nil, err
}
machines := []Machine{}
if err := db.Preload("AuthKey").Where(&Machine{NamespaceID: n.ID}).Find(&machines).Error; err != nil {
if err := h.db.Preload("AuthKey").Where(&Machine{NamespaceID: n.ID}).Find(&machines).Error; err != nil {
return nil, err
}
return &machines, nil
@ -126,13 +98,8 @@ func (h *Headscale) SetMachineNamespace(m *Machine, namespaceName string) error
if err != nil {
return err
}
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return err
}
m.NamespaceID = n.ID
db.Save(&m)
h.db.Save(&m)
return nil
}

6
namespaces_test.go
View file

@ -30,10 +30,6 @@ func (s *Suite) TestDestroyNamespaceErrors(c *check.C) {
pak, err := h.CreatePreAuthKey(n.Name, false, false, nil)
c.Assert(err, check.IsNil)
db, err := h.db()
if err != nil {
c.Fatal(err)
}
m := Machine{
ID: 0,
MachineKey: "foo",
@ -45,7 +41,7 @@ func (s *Suite) TestDestroyNamespaceErrors(c *check.C) {
RegisterMethod: "authKey",
AuthKeyID: uint(pak.ID),
}
db.Save(&m)
h.db.Save(&m)
err = h.DestroyNamespace("test")
c.Assert(err, check.Equals, errorNamespaceNotEmpty)

25
preauth_keys.go
View file

@ -4,7 +4,6 @@ import (
"crypto/rand"
"encoding/hex"
"errors"
"log"
"time"
"gorm.io/gorm"
@ -34,12 +33,6 @@ func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, epheme
return nil, err
}
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return nil, err
}
now := time.Now().UTC()
kstr, err := h.generateKey()
if err != nil {
@ -55,7 +48,7 @@ func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, epheme
CreatedAt: &now,
Expiration: expiration,
}
db.Save(&k)
h.db.Save(&k)
return &k, nil
}
@ -66,14 +59,9 @@ func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error)
if err != nil {
return nil, err
}
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return nil, err
}
keys := []PreAuthKey{}
if err := db.Preload("Namespace").Where(&PreAuthKey{NamespaceID: n.ID}).Find(&keys).Error; err != nil {
if err := h.db.Preload("Namespace").Where(&PreAuthKey{NamespaceID: n.ID}).Find(&keys).Error; err != nil {
return nil, err
}
return &keys, nil
@ -82,13 +70,8 @@ func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error)
// checkKeyValidity does the heavy lifting for validation of the PreAuthKey coming from a node
// If returns no error and a PreAuthKey, it can be used
func (h *Headscale) checkKeyValidity(k string) (*PreAuthKey, error) {
db, err := h.db()
if err != nil {
return nil, err
}
pak := PreAuthKey{}
if result := db.Preload("Namespace").First(&pak, "key = ?", k); errors.Is(result.Error, gorm.ErrRecordNotFound) {
if result := h.db.Preload("Namespace").First(&pak, "key = ?", k); errors.Is(result.Error, gorm.ErrRecordNotFound) {
return nil, errorAuthKeyNotFound
}
@ -101,7 +84,7 @@ func (h *Headscale) checkKeyValidity(k string) (*PreAuthKey, error) {
}
machines := []Machine{}
if err := db.Preload("AuthKey").Where(&Machine{AuthKeyID: uint(pak.ID)}).Find(&machines).Error; err != nil {
if err := h.db.Preload("AuthKey").Where(&Machine{AuthKeyID: uint(pak.ID)}).Find(&machines).Error; err != nil {
return nil, err
}

18
preauth_keys_test.go
View file

@ -73,10 +73,6 @@ func (*Suite) TestAlreadyUsedKey(c *check.C) {
pak, err := h.CreatePreAuthKey(n.Name, false, false, nil)
c.Assert(err, check.IsNil)
db, err := h.db()
if err != nil {
c.Fatal(err)
}
m := Machine{
ID: 0,
MachineKey: "foo",
@ -88,7 +84,7 @@ func (*Suite) TestAlreadyUsedKey(c *check.C) {
RegisterMethod: "authKey",
AuthKeyID: uint(pak.ID),
}
db.Save(&m)
h.db.Save(&m)
p, err := h.checkKeyValidity(pak.Key)
c.Assert(err, check.Equals, errorAuthKeyNotReusableAlreadyUsed)
@ -102,10 +98,6 @@ func (*Suite) TestReusableBeingUsedKey(c *check.C) {
pak, err := h.CreatePreAuthKey(n.Name, true, false, nil)
c.Assert(err, check.IsNil)
db, err := h.db()
if err != nil {
c.Fatal(err)
}
m := Machine{
ID: 1,
MachineKey: "foo",
@ -117,7 +109,7 @@ func (*Suite) TestReusableBeingUsedKey(c *check.C) {
RegisterMethod: "authKey",
AuthKeyID: uint(pak.ID),
}
db.Save(&m)
h.db.Save(&m)
p, err := h.checkKeyValidity(pak.Key)
c.Assert(err, check.IsNil)
@ -143,10 +135,6 @@ func (*Suite) TestEphemeralKey(c *check.C) {
pak, err := h.CreatePreAuthKey(n.Name, false, true, nil)
c.Assert(err, check.IsNil)
db, err := h.db()
if err != nil {
c.Fatal(err)
}
now := time.Now()
m := Machine{
ID: 0,
@ -160,7 +148,7 @@ func (*Suite) TestEphemeralKey(c *check.C) {
LastSeen: &now,
AuthKeyID: uint(pak.ID),
}
db.Save(&m)
h.db.Save(&m)
_, err = h.checkKeyValidity(pak.Key)
// Ephemeral keys are by definition reusable

9
routes.go
View file

@ -3,7 +3,6 @@ package headscale
import (
"encoding/json"
"errors"
"log"
"gorm.io/datatypes"
"inet.af/netaddr"
@ -42,15 +41,9 @@ func (h *Headscale) EnableNodeRoute(namespace string, nodeName string, routeStr
for _, rIP := range hi.RoutableIPs {
if rIP == route {
db, err := h.db()
if err != nil {
log.Printf("Cannot open DB: %s", err)
return nil, err
}
routes, _ := json.Marshal([]string{routeStr}) // TODO: only one for the time being, so overwriting the rest
m.EnabledRoutes = datatypes.JSON(routes)
db.Save(&m)
h.db.Save(&m)
// THIS IS COMPLETELY USELESS.
// The peers map is stored in memory in the server process.

7
routes_test.go
View file

@ -16,11 +16,6 @@ func (s *Suite) TestGetRoutes(c *check.C) {
pak, err := h.CreatePreAuthKey(n.Name, false, false, nil)
c.Assert(err, check.IsNil)
db, err := h.db()
if err != nil {
c.Fatal(err)
}
_, err = h.GetMachine("test", "testmachine")
c.Assert(err, check.NotNil)
@ -45,7 +40,7 @@ func (s *Suite) TestGetRoutes(c *check.C) {
AuthKeyID: uint(pak.ID),
HostInfo: datatypes.JSON(hostinfo),
}
db.Save(&m)
h.db.Save(&m)
r, err := h.GetNodeRoutes("test", "testmachine")
c.Assert(err, check.IsNil)

6
utils.go
View file

@ -78,10 +78,6 @@ func encodeMsg(b []byte, pubKey *wgkey.Key, privKey *wgkey.Private) ([]byte, err
}
func (h *Headscale) getAvailableIP() (*net.IP, error) {
db, err := h.db()
if err != nil {
return nil, err
}
i := 0
for {
ip, err := getRandomIP()
@ -89,7 +85,7 @@ func (h *Headscale) getAvailableIP() (*net.IP, error) {
return nil, err
}
m := Machine{}
if result := db.First(&m, "ip_address = ?", ip.String()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
if result := h.db.First(&m, "ip_address = ?", ip.String()); errors.Is(result.Error, gorm.ErrRecordNotFound) {
return ip, nil
}
i++