Add Route DB model and migration from existing field

Add migration from Machine.EnabledRoutes to the new Route table

Cleanup route.go and add helper methods to process HostInfo
This commit is contained in:
Juan Font 2022-11-23 18:50:30 +00:00 committed by Kristoffer Dalby
parent d1bca105ef
commit fba77de4eb
2 changed files with 159 additions and 89 deletions

87
db.go
View file

@ -18,8 +18,10 @@ import (
) )
const ( const (
dbVersion = "1" dbVersion = "1"
errValueNotFound = Error("not found")
errValueNotFound = Error("not found")
ErrCannotParsePrefix = Error("cannot parse prefix")
) )
// KV is a key-value store in a psql table. For future use... // KV is a key-value store in a psql table. For future use...
@ -79,6 +81,65 @@ func (h *Headscale) initDB() error {
} }
} }
err = db.AutoMigrate(&Route{})
if err != nil {
return err
}
if db.Migrator().HasColumn(&Machine{}, "enabled_routes") {
log.Info().Msgf("Database has legacy enabled_routes column in machine, migrating...")
type MachineAux struct {
ID uint64
EnabledRoutes IPPrefixes
}
machinesAux := []MachineAux{}
err := db.Table("machines").Select("id, enabled_routes").Scan(&machinesAux).Error
if err != nil {
log.Fatal().Err(err).Msg("Error accessing db")
}
for _, machine := range machinesAux {
for _, prefix := range machine.EnabledRoutes {
if err != nil {
log.Error().
Err(err).
Str("enabled_route", prefix.String()).
Msg("Error parsing enabled_route")
continue
}
err = db.Preload("Machine").Where("machine_id = ? AND prefix = ?", machine.ID, IPPrefix(prefix)).First(&Route{}).Error
if err == nil {
log.Info().
Str("enabled_route", prefix.String()).
Msg("Route already migrated to new table, skipping")
continue
}
route := Route{
MachineID: machine.ID,
Advertised: true,
Enabled: true,
Prefix: IPPrefix(prefix),
}
if err := h.db.Create(&route).Error; err != nil {
log.Error().Err(err).Msg("Error creating route")
} else {
log.Info().
Uint64("machine_id", route.MachineID).
Str("prefix", prefix.String()).
Msg("Route migrated")
}
}
}
err = db.Migrator().DropColumn(&Machine{}, "enabled_routes")
if err != nil {
log.Error().Err(err).Msg("Error dropping enabled_routes column")
}
}
err = db.AutoMigrate(&Machine{}) err = db.AutoMigrate(&Machine{})
if err != nil { if err != nil {
return err return err
@ -264,6 +325,28 @@ func (hi HostInfo) Value() (driver.Value, error) {
return string(bytes), err return string(bytes), err
} }
type IPPrefix netip.Prefix
func (i *IPPrefix) Scan(destination interface{}) error {
switch value := destination.(type) {
case string:
prefix, err := netip.ParsePrefix(value)
if err != nil {
return err
}
*i = IPPrefix(prefix)
return nil
default:
return fmt.Errorf("%w: unexpected data type %T", ErrCannotParsePrefix, destination)
}
}
// Value return json value, implement driver.Valuer interface.
func (i IPPrefix) Value() (driver.Value, error) {
prefixStr := netip.Prefix(i).String()
return prefixStr, nil
}
type IPPrefixes []netip.Prefix type IPPrefixes []netip.Prefix
func (i *IPPrefixes) Scan(destination interface{}) error { func (i *IPPrefixes) Scan(destination interface{}) error {

161
routes.go
View file

@ -3,115 +3,102 @@ package headscale
import ( import (
"fmt" "fmt"
"net/netip" "net/netip"
"gorm.io/gorm"
) )
const ( const (
ErrRouteIsNotAvailable = Error("route is not available") ErrRouteIsNotAvailable = Error("route is not available")
) )
// Deprecated: use machine function instead type Route struct {
// GetAdvertisedNodeRoutes returns the subnet routes advertised by a node (identified by gorm.Model
// namespace and node name).
func (h *Headscale) GetAdvertisedNodeRoutes( MachineID uint64
namespace string, Machine Machine
nodeName string, Prefix IPPrefix
) (*[]netip.Prefix, error) {
machine, err := h.GetMachine(namespace, nodeName) Advertised bool
Enabled bool
IsPrimary bool
}
type Routes []Route
func (r *Route) String() string {
return fmt.Sprintf("%s:%s", r.Machine, netip.Prefix(r.Prefix).String())
}
func (rs Routes) toPrefixes() []netip.Prefix {
prefixes := make([]netip.Prefix, len(rs))
for i, r := range rs {
prefixes[i] = netip.Prefix(r.Prefix)
}
return prefixes
}
// getMachinePrimaryRoutes returns the routes that are enabled and marked as primary (for subnet failover)
// Exit nodes are not considered for this, as they are never marked as Primary
func (h *Headscale) getMachinePrimaryRoutes(m *Machine) ([]Route, error) {
var routes []Route
err := h.db.
Preload("Machine").
Where("machine_id = ? AND advertised = ? AND enabled = ? AND is_primary = ?", m.ID, true, true, true).
Find(&routes).Error
if err != nil { if err != nil {
return nil, err return nil, err
} }
return &machine.HostInfo.RoutableIPs, nil return routes, nil
} }
// Deprecated: use machine function instead func (h *Headscale) processMachineRoutes(machine *Machine) error {
// GetEnabledNodeRoutes returns the subnet routes enabled by a node (identified by currentRoutes := []Route{}
// namespace and node name). err := h.db.Where("machine_id = ?", machine.ID).Find(&currentRoutes).Error
func (h *Headscale) GetEnabledNodeRoutes(
namespace string,
nodeName string,
) ([]netip.Prefix, error) {
machine, err := h.GetMachine(namespace, nodeName)
if err != nil {
return nil, err
}
return machine.EnabledRoutes, nil
}
// Deprecated: use machine function instead
// IsNodeRouteEnabled checks if a certain route has been enabled.
func (h *Headscale) IsNodeRouteEnabled(
namespace string,
nodeName string,
routeStr string,
) bool {
route, err := netip.ParsePrefix(routeStr)
if err != nil {
return false
}
enabledRoutes, err := h.GetEnabledNodeRoutes(namespace, nodeName)
if err != nil {
return false
}
for _, enabledRoute := range enabledRoutes {
if route == enabledRoute {
return true
}
}
return false
}
// Deprecated: use EnableRoute in machine.go
// EnableNodeRoute enables a subnet route advertised by a node (identified by
// namespace and node name).
func (h *Headscale) EnableNodeRoute(
namespace string,
nodeName string,
routeStr string,
) error {
machine, err := h.GetMachine(namespace, nodeName)
if err != nil { if err != nil {
return err return err
} }
route, err := netip.ParsePrefix(routeStr) advertisedRoutes := map[netip.Prefix]bool{}
if err != nil { for _, prefix := range machine.HostInfo.RoutableIPs {
return err advertisedRoutes[prefix] = false
} }
availableRoutes, err := h.GetAdvertisedNodeRoutes(namespace, nodeName) for _, route := range currentRoutes {
if err != nil { if _, ok := advertisedRoutes[netip.Prefix(route.Prefix)]; ok {
return err if !route.Advertised {
} route.Advertised = true
err := h.db.Save(&route).Error
enabledRoutes, err := h.GetEnabledNodeRoutes(namespace, nodeName) if err != nil {
if err != nil { return err
return err }
} }
advertisedRoutes[netip.Prefix(route.Prefix)] = true
available := false } else {
for _, availableRoute := range *availableRoutes { if route.Advertised {
// If the route is available, and not yet enabled, add it to the new routing table route.Advertised = false
if route == availableRoute { route.Enabled = false
available = true err := h.db.Save(&route).Error
if !h.IsNodeRouteEnabled(namespace, nodeName, routeStr) { if err != nil {
enabledRoutes = append(enabledRoutes, route) return err
}
} }
} }
} }
if !available { for prefix, exists := range advertisedRoutes {
return ErrRouteIsNotAvailable if !exists {
} route := Route{
MachineID: machine.ID,
machine.EnabledRoutes = enabledRoutes Prefix: IPPrefix(prefix),
Advertised: true,
if err := h.db.Save(&machine).Error; err != nil { Enabled: false,
return fmt.Errorf("failed to update node routes in the database: %w", err) }
err := h.db.Create(&route).Error
if err != nil {
return err
}
}
} }
return nil return nil