Skip to content

Commit

Permalink
delete all, rebuild
Browse files Browse the repository at this point in the history 
Signed-off-by: Kristoffer Dalby <[email protected]>
  • Loading branch information
@kradalby
kradalby committed Feb 13, 2025
1 parent c61fbe9 commit 175f5be
Show file tree
Hide file tree
Showing 12 changed files with 464 additions and 2,351 deletions.
6 changes: 6 additions & 0 deletions hscontrol/app.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import (
"github.com/juanfont/headscale/hscontrol/mapper"
"github.com/juanfont/headscale/hscontrol/notifier"
"github.com/juanfont/headscale/hscontrol/policy"
"github.com/juanfont/headscale/hscontrol/routes"
"github.com/juanfont/headscale/hscontrol/types"
"github.com/juanfont/headscale/hscontrol/util"
zerolog "github.com/philip-bui/grpc-zerolog"
Expand Down Expand Up @@ -92,6 +93,7 @@ type Headscale struct {
polManOnce sync.Once
polMan policy.PolicyManager
extraRecordMan *dns.ExtraRecordsMan
primaryRoutes *routes.PrimaryRoutes

mapper *mapper.Mapper
nodeNotifier *notifier.Notifier
Expand Down Expand Up @@ -134,6 +136,7 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
registrationCache: registrationCache,
pollNetMapStreamWG: sync.WaitGroup{},
nodeNotifier: notifier.NewNotifier(cfg),
primaryRoutes: routes.New(),
}

app.db, err = db.NewHeadscaleDatabase(
Expand Down Expand Up @@ -495,6 +498,8 @@ func (h *Headscale) createRouter(grpcMux *grpcRuntime.ServeMux) *mux.Router {

// TODO(kradalby): Do a variant of this, and polman which only updates the node that has changed.
// Maybe we should attempt a new in memory state and not go via the DB?
// Maybe this should be implemented as an event bus?
// A bool is returned indicating if a full update was sent to all nodes
func usersChangedHook(db *db.HSDatabase, polMan policy.PolicyManager, notif *notifier.Notifier) error {
users, err := db.ListUsers()
if err != nil {
Expand All @@ -516,6 +521,7 @@ func usersChangedHook(db *db.HSDatabase, polMan policy.PolicyManager, notif *not

// TODO(kradalby): Do a variant of this, and polman which only updates the node that has changed.
// Maybe we should attempt a new in memory state and not go via the DB?
// Maybe this should be implemented as an event bus?
// A bool is returned indicating if a full update was sent to all nodes
func nodesChangedHook(db *db.HSDatabase, polMan policy.PolicyManager, notif *notifier.Notifier) (bool, error) {
nodes, err := db.ListNodes()
Expand Down
155 changes: 6 additions & 149 deletions hscontrol/db/node.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,10 @@ import (

"github.com/juanfont/headscale/hscontrol/types"
"github.com/juanfont/headscale/hscontrol/util"
"github.com/puzpuzpuz/xsync/v3"
"github.com/rs/zerolog/log"
"gorm.io/gorm"
"tailscale.com/tailcfg"
"tailscale.com/types/key"
"tailscale.com/types/ptr"
)

const (
Expand Down Expand Up @@ -267,29 +265,23 @@ func NodeSetExpiry(tx *gorm.DB,
return tx.Model(&types.Node{}).Where("id = ?", nodeID).Update("expiry", expiry).Error
}

func (hsdb *HSDatabase) DeleteNode(node *types.Node, isLikelyConnected *xsync.MapOf[types.NodeID, bool]) ([]types.NodeID, error) {
return Write(hsdb.DB, func(tx *gorm.DB) ([]types.NodeID, error) {
return DeleteNode(tx, node, isLikelyConnected)
func (hsdb *HSDatabase) DeleteNode(node *types.Node) error {
return hsdb.Write(func(tx *gorm.DB) error {
return DeleteNode(tx, node)
})
}

// DeleteNode deletes a Node from the database.
// Caller is responsible for notifying all of change.
func DeleteNode(tx *gorm.DB,
node *types.Node,
isLikelyConnected *xsync.MapOf[types.NodeID, bool],
) ([]types.NodeID, error) {
changed, err := deleteNodeRoutes(tx, node, isLikelyConnected)
if err != nil {
return changed, err
}

) error {
// Unscoped causes the node to be fully removed from the database.
if err := tx.Unscoped().Delete(&types.Node{}, node.ID).Error; err != nil {
return changed, err
return err
}

return changed, nil
return nil
}

// DeleteEphemeralNode deletes a Node from the database, note that this method
Expand Down Expand Up @@ -495,141 +487,6 @@ func NodeSave(tx *gorm.DB, node *types.Node) error {
return tx.Save(node).Error
}

func (hsdb *HSDatabase) GetAdvertisedRoutes(node *types.Node) ([]netip.Prefix, error) {
return Read(hsdb.DB, func(rx *gorm.DB) ([]netip.Prefix, error) {
return GetAdvertisedRoutes(rx, node)
})
}

// GetAdvertisedRoutes returns the routes that are be advertised by the given node.
func GetAdvertisedRoutes(tx *gorm.DB, node *types.Node) ([]netip.Prefix, error) {
routes := types.Routes{}

err := tx.
Preload("Node").
Where("node_id = ? AND advertised = ?", node.ID, true).Find(&routes).Error
if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
return nil, fmt.Errorf("getting advertised routes for node(%d): %w", node.ID, err)
}

var prefixes []netip.Prefix
for _, route := range routes {
prefixes = append(prefixes, netip.Prefix(route.Prefix))
}

return prefixes, nil
}

func (hsdb *HSDatabase) GetEnabledRoutes(node *types.Node) ([]netip.Prefix, error) {
return Read(hsdb.DB, func(rx *gorm.DB) ([]netip.Prefix, error) {
return GetEnabledRoutes(rx, node)
})
}

// GetEnabledRoutes returns the routes that are enabled for the node.
func GetEnabledRoutes(tx *gorm.DB, node *types.Node) ([]netip.Prefix, error) {
routes := types.Routes{}

err := tx.
Preload("Node").
Where("node_id = ? AND advertised = ? AND enabled = ?", node.ID, true, true).
Find(&routes).Error
if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
return nil, fmt.Errorf("getting enabled routes for node(%d): %w", node.ID, err)
}

var prefixes []netip.Prefix
for _, route := range routes {
prefixes = append(prefixes, netip.Prefix(route.Prefix))
}

return prefixes, nil
}

func IsRoutesEnabled(tx *gorm.DB, node *types.Node, routeStr string) bool {
route, err := netip.ParsePrefix(routeStr)
if err != nil {
return false
}

enabledRoutes, err := GetEnabledRoutes(tx, node)
if err != nil {
return false
}

for _, enabledRoute := range enabledRoutes {
if route == enabledRoute {
return true
}
}

return false
}

func (hsdb *HSDatabase) enableRoutes(
node *types.Node,
newRoutes ...netip.Prefix,
) (*types.StateUpdate, error) {
return Write(hsdb.DB, func(tx *gorm.DB) (*types.StateUpdate, error) {
return enableRoutes(tx, node, newRoutes...)
})
}

// enableRoutes enables new routes based on a list of new routes.
func enableRoutes(tx *gorm.DB,
node *types.Node, newRoutes ...netip.Prefix,
) (*types.StateUpdate, error) {
advertisedRoutes, err := GetAdvertisedRoutes(tx, node)
if err != nil {
return nil, err
}

for _, newRoute := range newRoutes {
if !slices.Contains(advertisedRoutes, newRoute) {
return nil, fmt.Errorf(
"route (%s) is not available on node %s: %w",
node.Hostname,
newRoute, ErrNodeRouteIsNotAvailable,
)
}
}

// Separate loop so we don't leave things in a half-updated state
for _, prefix := range newRoutes {
route := types.Route{}
err := tx.Preload("Node").
Where("node_id = ? AND prefix = ?", node.ID, prefix.String()).
First(&route).Error
if err == nil {
route.Enabled = true

// Mark already as primary if there is only this node offering this subnet
// (and is not an exit route)
if !route.IsExitRoute() {
route.IsPrimary = isUniquePrefix(tx, route)
}

err = tx.Save(&route).Error
if err != nil {
return nil, fmt.Errorf("failed to enable route: %w", err)
}
} else {
return nil, fmt.Errorf("failed to find route: %w", err)
}
}

// Ensure the node has the latest routes when notifying the other
// nodes
nRoutes, err := GetNodeRoutes(tx, node)
if err != nil {
return nil, fmt.Errorf("failed to read back routes: %w", err)
}

node.Routes = nRoutes

return ptr.To(types.UpdatePeerChanged(node.ID)), nil
}

func generateGivenName(suppliedName string, randomSuffix bool) (string, error) {
suppliedName = util.ConvertWithFQDNRules(suppliedName)
if len(suppliedName) > util.LabelHostnameLength {
Expand Down
3 changes: 1 addition & 2 deletions hscontrol/db/node_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ import (
"github.com/juanfont/headscale/hscontrol/policy"
"github.com/juanfont/headscale/hscontrol/types"
"github.com/juanfont/headscale/hscontrol/util"
"github.com/puzpuzpuz/xsync/v3"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gopkg.in/check.v1"
Expand Down Expand Up @@ -102,7 +101,7 @@ func (s *Suite) TestHardDeleteNode(c *check.C) {
trx := db.DB.Save(&node)
c.Assert(trx.Error, check.IsNil)

_, err = db.DeleteNode(&node, xsync.NewMapOf[types.NodeID, bool]())
err = db.DeleteNode(&node)
c.Assert(err, check.IsNil)

_, err = db.getNode(types.UserID(user.ID), "testnode3")
Expand Down
Loading

0 comments on commit 175f5be

Please sign in to comment.