headscale/db.go

127 lines
2.2 KiB
Go
Raw Normal View History

2020-06-21 12:32:08 +02:00
package headscale
import (
"errors"
2021-06-24 15:44:19 +02:00
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
2020-06-21 12:32:08 +02:00
)
2021-11-15 19:18:14 +00:00
const (
dbVersion = "1"
errValueNotFound = Error("not found")
)
2020-06-21 12:32:08 +02:00
2021-02-23 20:11:30 +01:00
// KV is a key-value store in a psql table. For future use...
2020-06-21 12:32:08 +02:00
type KV struct {
Key string
Value string
}
func (h *Headscale) initDB() error {
2021-07-04 21:40:46 +02:00
db, err := h.openDB()
2020-06-21 12:32:08 +02:00
if err != nil {
return err
}
2021-07-04 21:40:46 +02:00
h.db = db
if h.dbType == Postgres {
2022-01-30 14:53:40 +00:00
db.Exec(`create extension if not exists "uuid-ossp";`)
}
_ = db.Migrator().RenameColumn(&Machine{}, "ip_address", "ip_addresses")
2021-06-24 15:44:19 +02:00
err = db.AutoMigrate(&Machine{})
if err != nil {
return err
}
2021-06-24 15:44:19 +02:00
err = db.AutoMigrate(&KV{})
if err != nil {
return err
}
2021-06-24 15:44:19 +02:00
err = db.AutoMigrate(&Namespace{})
if err != nil {
return err
}
2021-06-24 15:44:19 +02:00
err = db.AutoMigrate(&PreAuthKey{})
if err != nil {
return err
}
2020-06-21 12:32:08 +02:00
2021-09-06 14:43:43 +02:00
err = db.AutoMigrate(&SharedMachine{})
if err != nil {
return err
}
err = h.setValue("db_version", dbVersion)
2021-11-14 16:46:09 +01:00
return err
2020-06-21 12:32:08 +02:00
}
2021-07-04 21:40:46 +02:00
func (h *Headscale) openDB() (*gorm.DB, error) {
2021-06-24 15:44:19 +02:00
var db *gorm.DB
var err error
var log logger.Interface
if h.dbDebug {
log = logger.Default
} else {
log = logger.Default.LogMode(logger.Silent)
}
2021-06-24 15:44:19 +02:00
switch h.dbType {
case Sqlite:
db, err = gorm.Open(sqlite.Open(h.dbString), &gorm.Config{
DisableForeignKeyConstraintWhenMigrating: true,
Logger: log,
})
case Postgres:
db, err = gorm.Open(postgres.Open(h.dbString), &gorm.Config{
DisableForeignKeyConstraintWhenMigrating: true,
Logger: log,
})
2021-06-24 15:44:19 +02:00
}
2020-06-21 12:32:08 +02:00
if err != nil {
return nil, err
}
2020-06-21 12:32:08 +02:00
return db, nil
}
2021-11-13 08:39:04 +00:00
// getValue returns the value for the given key in KV.
2020-06-21 12:32:08 +02:00
func (h *Headscale) getValue(key string) (string, error) {
var row KV
2021-11-13 08:36:45 +00:00
if result := h.db.First(&row, "key = ?", key); errors.Is(
result.Error,
gorm.ErrRecordNotFound,
) {
2021-11-15 19:18:14 +00:00
return "", errValueNotFound
2020-06-21 12:32:08 +02:00
}
2021-11-14 16:46:09 +01:00
2020-06-21 12:32:08 +02:00
return row.Value, nil
}
2021-11-13 08:39:04 +00:00
// setValue sets value for the given key in KV.
2020-06-21 12:32:08 +02:00
func (h *Headscale) setValue(key string, value string) error {
2021-11-15 16:15:50 +00:00
keyValue := KV{
2020-06-21 12:32:08 +02:00
Key: key,
Value: value,
}
2021-07-04 21:40:46 +02:00
2021-11-14 18:09:22 +01:00
if _, err := h.getValue(key); err == nil {
2021-11-15 16:15:50 +00:00
h.db.Model(&keyValue).Where("key = ?", key).Update("value", value)
2021-11-14 16:46:09 +01:00
2020-06-21 12:32:08 +02:00
return nil
}
2021-11-15 16:15:50 +00:00
h.db.Create(keyValue)
2021-11-14 16:46:09 +01:00
2020-06-21 12:32:08 +02:00
return nil
}