new IP allocator and add postgres to integration tests. (#1756)

integration/acl_test.go

@@ -415,7 +415,7 @@ func TestACLAllowUserDst(t *testing.T) {
 		},
 		2,
 	)
-	defer scenario.Shutdown()
+	// defer scenario.Shutdown()
 
 	user1Clients, err := scenario.ListTailscaleClients("user1")
 	assertNoErr(t, err)

integration/cli_test.go

@@ -7,12 +7,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
-
 	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/juanfont/headscale/hscontrol/policy"
 	"github.com/juanfont/headscale/integration/hsic"
 	"github.com/juanfont/headscale/integration/tsic"
+	"github.com/stretchr/testify/assert"
 )
 
 func executeAndUnmarshal[T any](headscale ControlServer, command []string, result T) error {

integration/hsic/config.go

@@ -72,9 +72,9 @@ database:
   sqlite.path: /tmp/integration_test_db.sqlite3
 ephemeral_node_inactivity_timeout: 30m
 node_update_check_interval: 10s
-ip_prefixes:
-  - fd7a:115c:a1e0::/48
-  - 100.64.0.0/10
+prefixes:
+  v6: fd7a:115c:a1e0::/48
+  v4: 100.64.0.0/10
 dns_config:
   base_domain: headscale.net
   magic_dns: true
@@ -115,7 +115,8 @@ func DefaultConfigEnv() map[string]string {
 		"HEADSCALE_DATABASE_SQLITE_PATH":              "/tmp/integration_test_db.sqlite3",
 		"HEADSCALE_EPHEMERAL_NODE_INACTIVITY_TIMEOUT": "30m",
 		"HEADSCALE_NODE_UPDATE_CHECK_INTERVAL":        "10s",
-		"HEADSCALE_IP_PREFIXES":                       "fd7a:115c:a1e0::/48 100.64.0.0/10",
+		"HEADSCALE_PREFIXES_V4":                       "100.64.0.0/10",
+		"HEADSCALE_PREFIXES_V6":                       "fd7a:115c:a1e0::/48",
 		"HEADSCALE_DNS_CONFIG_BASE_DOMAIN":            "headscale.net",
 		"HEADSCALE_DNS_CONFIG_MAGIC_DNS":              "true",
 		"HEADSCALE_DNS_CONFIG_DOMAINS":                "",

integration/hsic/hsic.go

@@ -56,6 +56,8 @@ type HeadscaleInContainer struct {
 	container *dockertest.Resource
 	network   *dockertest.Network
 
+	pgContainer *dockertest.Resource
+
 	// optional config
 	port             int
 	extraPorts       []string
@@ -65,6 +67,7 @@ type HeadscaleInContainer struct {
 	tlsCert          []byte
 	tlsKey           []byte
 	filesInContainer []fileInContainer
+	postgres         bool
 }
 
 // Option represent optional settings that can be given to a
@@ -162,6 +165,14 @@ func WithFileInContainer(path string, contents []byte) Option {
 	}
 }
 
+// WithPostgres spins up a Postgres container and
+// sets it as the main database.
+func WithPostgres() Option {
+	return func(hsic *HeadscaleInContainer) {
+		hsic.postgres = true
+	}
+}
+
 // New returns a new HeadscaleInContainer instance.
 func New(
 	pool *dockertest.Pool,
@@ -209,6 +220,33 @@ func New(
 		ContextDir: dockerContextPath,
 	}
 
+	if hsic.postgres {
+		hsic.env["HEADSCALE_DATABASE_TYPE"] = "postgres"
+		hsic.env["HEADSCALE_DATABASE_POSTGRES_HOST"] = fmt.Sprintf("postgres-%s", hash)
+		hsic.env["HEADSCALE_DATABASE_POSTGRES_USER"] = "headscale"
+		hsic.env["HEADSCALE_DATABASE_POSTGRES_PASS"] = "headscale"
+		hsic.env["HEADSCALE_DATABASE_POSTGRES_NAME"] = "headscale"
+		delete(hsic.env, "HEADSCALE_DATABASE_SQLITE_PATH")
+
+		pg, err := pool.RunWithOptions(
+			&dockertest.RunOptions{
+				Name:       fmt.Sprintf("postgres-%s", hash),
+				Repository: "postgres",
+				Tag:        "latest",
+				Networks:   []*dockertest.Network{network},
+				Env: []string{
+					"POSTGRES_USER=headscale",
+					"POSTGRES_PASSWORD=headscale",
+					"POSTGRES_DB=headscale",
+				},
+			})
+		if err != nil {
+			return nil, fmt.Errorf("starting postgres container: %w", err)
+		}
+
+		hsic.pgContainer = pg
+	}
+
 	env := []string{
 		"HEADSCALE_PROFILING_ENABLED=1",
 		"HEADSCALE_PROFILING_PATH=/tmp/profile",
@@ -348,12 +386,20 @@ func (t *HeadscaleInContainer) Shutdown() error {
 		)
 	}
 
-	err = t.SaveDatabase("/tmp/control")
-	if err != nil {
-		log.Printf(
-			"Failed to save database from control: %s",
-			fmt.Errorf("failed to save database from control: %w", err),
-		)
+	// We dont have a database to save if we use postgres
+	if !t.postgres {
+		err = t.SaveDatabase("/tmp/control")
+		if err != nil {
+			log.Printf(
+				"Failed to save database from control: %s",
+				fmt.Errorf("failed to save database from control: %w", err),
+			)
+		}
+	}
+
+	// Cleanup postgres container if enabled.
+	if t.postgres {
+		t.pool.Purge(t.pgContainer)
 	}
 
 	return t.pool.Purge(t.container)

integration/scenario.go

@@ -18,12 +18,15 @@ import (
 	"github.com/puzpuzpuz/xsync/v3"
 	"github.com/samber/lo"
 	"golang.org/x/sync/errgroup"
+	"tailscale.com/envknob"
 )
 
 const (
 	scenarioHashLength = 6
 )
 
+var usePostgresForTest = envknob.Bool("HEADSCALE_INTEGRATION_POSTGRES")
+
 func enabledVersions(vs map[string]bool) []string {
 	var ret []string
 	for version, enabled := range vs {
@@ -452,6 +455,10 @@ func (s *Scenario) CreateHeadscaleEnv(
 	tsOpts []tsic.Option,
 	opts ...hsic.Option,
 ) error {
+	if usePostgresForTest {
+		opts = append(opts, hsic.WithPostgres())
+	}
+
 	headscale, err := s.Headscale(opts...)
 	if err != nil {
 		return err

integration/utils.go

@@ -181,7 +181,6 @@ func assertValidNetmap(t *testing.T, client TailscaleClient) {
 			if ni := hi.NetInfo(); ni.Valid() {
 				assert.NotEqualf(t, 0, ni.PreferredDERP(), "peer (%s) has no home DERP in %q's netmap, got: %s", peer.ComputedName(), client.Hostname(), peer.Hostinfo().NetInfo().PreferredDERP())
 			}
-
 		}
 
 		assert.NotEmptyf(t, peer.Endpoints(), "peer (%s) of %q does not have any endpoints", peer.ComputedName(), client.Hostname())