From 3da872bb3096c7b9b4a2c07514d4baa7c25b793e Mon Sep 17 00:00:00 2001
From: Juan Font <juanfontalonso@gmail.com>
Date: Thu, 24 Nov 2022 17:14:21 +0000
Subject: [PATCH] Mark as primary the first instance of subnet + tests

In preparation for subnet failover, mark the initial occurrence of a subnet as the primary one.
---
 machine.go     | 12 +++++---
 routes.go      | 17 ++++++++++
 routes_test.go | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 108 insertions(+), 5 deletions(-)

diff --git a/machine.go b/machine.go
index d3d54802..376ba5a1 100644
--- a/machine.go
+++ b/machine.go
@@ -38,11 +38,6 @@ const (
 	maxHostnameLength = 255
 )
 
-var (
-	ExitRouteV4 = netip.MustParsePrefix("0.0.0.0/0")
-	ExitRouteV6 = netip.MustParsePrefix("::/0")
-)
-
 // Machine is a Headscale client.
 type Machine struct {
 	ID          uint64 `gorm:"primary_key"`
@@ -1025,6 +1020,13 @@ func (h *Headscale) EnableRoutes(machine *Machine, routeStrs ...string) error {
 			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 prefix != ExitRouteV4 && prefix != ExitRouteV6 {
+				route.IsPrimary = h.isUniquePrefix(route)
+			}
+
 			err = h.db.Save(&route).Error
 			if err != nil {
 				return fmt.Errorf("failed to enable route: %w", err)
diff --git a/routes.go b/routes.go
index 3ba710be..36d67a90 100644
--- a/routes.go
+++ b/routes.go
@@ -11,6 +11,11 @@ const (
 	ErrRouteIsNotAvailable = Error("route is not available")
 )
 
+var (
+	ExitRouteV4 = netip.MustParsePrefix("0.0.0.0/0")
+	ExitRouteV6 = netip.MustParsePrefix("::/0")
+)
+
 type Route struct {
 	gorm.Model
 
@@ -37,6 +42,18 @@ func (rs Routes) toPrefixes() []netip.Prefix {
 	return prefixes
 }
 
+// isUniquePrefix returns if there is another machine providing the same route already
+func (h *Headscale) isUniquePrefix(route Route) bool {
+	var count int64
+	h.db.
+		Model(&Route{}).
+		Where("prefix = ? AND machine_id != ? AND advertised = ? AND enabled = ?",
+			route.Prefix,
+			route.MachineID,
+			true, true).Count(&count)
+	return count == 0
+}
+
 // 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) {
diff --git a/routes_test.go b/routes_test.go
index 2f034942..2560898e 100644
--- a/routes_test.go
+++ b/routes_test.go
@@ -125,3 +125,87 @@ func (s *Suite) TestGetEnableRoutes(c *check.C) {
 	c.Assert(err, check.IsNil)
 	c.Assert(len(enabledRoutesWithAdditionalRoute), check.Equals, 2)
 }
+
+func (s *Suite) TestIsUniquePrefix(c *check.C) {
+	namespace, err := app.CreateNamespace("test")
+	c.Assert(err, check.IsNil)
+
+	pak, err := app.CreatePreAuthKey(namespace.Name, false, false, nil, nil)
+	c.Assert(err, check.IsNil)
+
+	_, err = app.GetMachine("test", "test_enable_route_machine")
+	c.Assert(err, check.NotNil)
+
+	route, err := netip.ParsePrefix(
+		"10.0.0.0/24",
+	)
+	c.Assert(err, check.IsNil)
+
+	route2, err := netip.ParsePrefix(
+		"150.0.10.0/25",
+	)
+	c.Assert(err, check.IsNil)
+
+	hostInfo1 := tailcfg.Hostinfo{
+		RoutableIPs: []netip.Prefix{route, route2},
+	}
+	machine1 := Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Hostname:       "test_enable_route_machine",
+		NamespaceID:    namespace.ID,
+		RegisterMethod: RegisterMethodAuthKey,
+		AuthKeyID:      uint(pak.ID),
+		HostInfo:       HostInfo(hostInfo1),
+	}
+	app.db.Save(&machine1)
+
+	err = app.processMachineRoutes(&machine1)
+	c.Assert(err, check.IsNil)
+
+	err = app.EnableRoutes(&machine1, route.String())
+	c.Assert(err, check.IsNil)
+
+	err = app.EnableRoutes(&machine1, route2.String())
+	c.Assert(err, check.IsNil)
+
+	hostInfo2 := tailcfg.Hostinfo{
+		RoutableIPs: []netip.Prefix{route2},
+	}
+	machine2 := Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Hostname:       "test_enable_route_machine",
+		NamespaceID:    namespace.ID,
+		RegisterMethod: RegisterMethodAuthKey,
+		AuthKeyID:      uint(pak.ID),
+		HostInfo:       HostInfo(hostInfo2),
+	}
+	app.db.Save(&machine2)
+
+	err = app.processMachineRoutes(&machine2)
+	c.Assert(err, check.IsNil)
+
+	err = app.EnableRoutes(&machine2, route2.String())
+	c.Assert(err, check.IsNil)
+
+	enabledRoutes1, err := app.GetEnabledRoutes(&machine1)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(enabledRoutes1), check.Equals, 2)
+
+	enabledRoutes2, err := app.GetEnabledRoutes(&machine2)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(enabledRoutes2), check.Equals, 1)
+
+	routes, err := app.getMachinePrimaryRoutes(&machine1)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(routes), check.Equals, 2)
+
+	routes, err = app.getMachinePrimaryRoutes(&machine2)
+	c.Assert(err, check.IsNil)
+	c.Assert(len(routes), check.Equals, 0)
+}