headscale/integration/scenario.go

package integration

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/netip"
	"net/url"
	"os"
	"slices"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
	"github.com/juanfont/headscale/hscontrol/capver"
	"github.com/juanfont/headscale/hscontrol/types"
	"github.com/juanfont/headscale/hscontrol/util"
	"github.com/juanfont/headscale/integration/dockertestutil"
	"github.com/juanfont/headscale/integration/dsic"
	"github.com/juanfont/headscale/integration/hsic"
	"github.com/juanfont/headscale/integration/integrationutil"
	"github.com/juanfont/headscale/integration/tsic"
	"github.com/oauth2-proxy/mockoidc"
	"github.com/ory/dockertest/v3"
	"github.com/ory/dockertest/v3/docker"
	"github.com/puzpuzpuz/xsync/v4"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	xmaps "golang.org/x/exp/maps"
	"golang.org/x/sync/errgroup"
	"tailscale.com/envknob"
	"tailscale.com/util/mak"
	"tailscale.com/util/multierr"
)

const (
	scenarioHashLength = 6
)

var usePostgresForTest = envknob.Bool("HEADSCALE_INTEGRATION_POSTGRES")

var (
	errNoHeadscaleAvailable = errors.New("no headscale available")
	errNoUserAvailable      = errors.New("no user available")
	errNoClientFound        = errors.New("client not found")

	// AllVersions represents a list of Tailscale versions the suite
	// uses to test compatibility with the ControlServer.
	//
	// The list contains two special cases, "head" and "unstable" which
	// points to the current tip of Tailscale's main branch and the latest
	// released unstable version.
	//
	// The rest of the version represents Tailscale versions that can be
	// found in Tailscale's apt repository.
	AllVersions = append([]string{"head", "unstable"}, capver.TailscaleLatestMajorMinor(capver.SupportedMajorMinorVersions, true)...)

	// MustTestVersions is the minimum set of versions we should test.
	// At the moment, this is arbitrarily chosen as:
	//
	// - Two unstable (HEAD and unstable)
	// - Two latest versions
	// - Two oldest supported version.
	MustTestVersions = append(
		AllVersions[0:4],
		AllVersions[len(AllVersions)-2:]...,
	)
)

// User represents a User in the ControlServer and a map of TailscaleClient's
// associated with the User.
type User struct {
	Clients map[string]TailscaleClient

	createWaitGroup errgroup.Group
	joinWaitGroup   errgroup.Group
	syncWaitGroup   errgroup.Group
}

// Scenario is a representation of an environment with one ControlServer and
// one or more User's and its associated TailscaleClients.
// A Scenario is intended to simplify setting up a new testcase for testing
// a ControlServer with TailscaleClients.
// TODO(kradalby): make control server configurable, test correctness with Tailscale SaaS.
type Scenario struct {
	// TODO(kradalby): support multiple headcales for later, currently only
	// use one.
	controlServers *xsync.Map[string, ControlServer]
	derpServers    []*dsic.DERPServerInContainer

	users map[string]*User

	pool          *dockertest.Pool
	networks      map[string]*dockertest.Network
	mockOIDC      scenarioOIDC
	extraServices map[string][]*dockertest.Resource

	mu sync.Mutex

	spec          ScenarioSpec
	userToNetwork map[string]*dockertest.Network

	testHashPrefix     string
	testDefaultNetwork string
}

// ScenarioSpec describes the users, nodes, and network topology to
// set up for a given scenario.
type ScenarioSpec struct {
	// Users is a list of usernames that will be created.
	// Each created user will get nodes equivalent to NodesPerUser
	Users []string

	// NodesPerUser is how many nodes should be attached to each user.
	NodesPerUser int

	// Networks, if set, is the separate Docker networks that should be
	// created and a list of the users that should be placed in those networks.
	// If not set, a single network will be created and all users+nodes will be
	// added there.
	// Please note that Docker networks are not necessarily routable and
	// connections between them might fall back to DERP.
	Networks map[string][]string

	// ExtraService, if set, is additional a map of network to additional
	// container services that should be set up. These container services
	// typically dont run Tailscale, e.g. web service to test subnet router.
	ExtraService map[string][]extraServiceFunc

	// Versions is specific list of versions to use for the test.
	Versions []string

	// OIDCSkipUserCreation, if true, skips creating users via headscale CLI
	// during environment setup. Useful for OIDC tests where the SSH policy
	// references users by name, since OIDC login creates users automatically
	// and pre-creating them via CLI causes duplicate user records.
	OIDCSkipUserCreation bool

	// OIDCUsers, if populated, will start a Mock OIDC server and populate
	// the user login stack with the given users.
	// If the NodesPerUser is set, it should align with this list to ensure
	// the correct users are logged in.
	// This is because the MockOIDC server can only serve login
	// requests based on a queue it has been given on startup.
	// We currently only populates it with one login request per user.
	OIDCUsers     []mockoidc.MockUser
	OIDCAccessTTL time.Duration

	MaxWait time.Duration
}

func (s *Scenario) prefixedNetworkName(name string) string {
	return s.testHashPrefix + "-" + name
}

// NewScenario creates a test Scenario which can be used to bootstraps a ControlServer with
// a set of Users and TailscaleClients.
func NewScenario(spec ScenarioSpec) (*Scenario, error) {
	pool, err := dockertest.NewPool("")
	if err != nil {
		return nil, fmt.Errorf("connecting to docker: %w", err)
	}

	// Opportunity to clean up unreferenced networks.
	// This might be a no op, but it is worth a try as we sometime
	// dont clean up nicely after ourselves.
	_ = dockertestutil.CleanUnreferencedNetworks(pool)
	_ = dockertestutil.CleanImagesInCI(pool)

	if spec.MaxWait == 0 {
		pool.MaxWait = dockertestMaxWait()
	} else {
		pool.MaxWait = spec.MaxWait
	}

	testHashPrefix := "hs-" + util.MustGenerateRandomStringDNSSafe(scenarioHashLength)
	s := &Scenario{
		controlServers: xsync.NewMap[string, ControlServer](),
		users:          make(map[string]*User),

		pool: pool,
		spec: spec,

		testHashPrefix:     testHashPrefix,
		testDefaultNetwork: testHashPrefix + "-default",
	}

	var userToNetwork map[string]*dockertest.Network

	if spec.Networks != nil || len(spec.Networks) != 0 {
		for name, users := range s.spec.Networks {
			networkName := testHashPrefix + "-" + name

			network, err := s.AddNetwork(networkName)
			if err != nil {
				return nil, err
			}

			for _, user := range users {
				if n2, ok := userToNetwork[user]; ok {
					return nil, fmt.Errorf("users can only have nodes placed in one network: %s into %s but already in %s", user, network.Network.Name, n2.Network.Name) //nolint:err113
				}

				mak.Set(&userToNetwork, user, network)
			}
		}
	} else {
		_, err := s.AddNetwork(s.testDefaultNetwork)
		if err != nil {
			return nil, err
		}
	}

	for network, extras := range spec.ExtraService {
		for _, extra := range extras {
			svc, err := extra(s, network)
			if err != nil {
				return nil, err
			}

			mak.Set(&s.extraServices, s.prefixedNetworkName(network), append(s.extraServices[s.prefixedNetworkName(network)], svc))
		}
	}

	s.userToNetwork = userToNetwork

	if len(spec.OIDCUsers) != 0 {
		ttl := defaultAccessTTL
		if spec.OIDCAccessTTL != 0 {
			ttl = spec.OIDCAccessTTL
		}

		err = s.runMockOIDC(ttl, spec.OIDCUsers)
		if err != nil {
			return nil, err
		}
	}

	return s, nil
}

func (s *Scenario) AddNetwork(name string) (*dockertest.Network, error) {
	network, err := dockertestutil.GetFirstOrCreateNetwork(s.pool, name)
	if err != nil {
		return nil, fmt.Errorf("creating or getting network: %w", err)
	}

	// We run the test suite in a docker container that calls a couple of endpoints for
	// readiness checks, this ensures that we can run the tests with individual networks
	// and have the client reach the different containers.
	// The container name includes the run ID to support multiple concurrent test runs.
	testSuiteName := "headscale-test-suite"
	if runID := dockertestutil.GetIntegrationRunID(); runID != "" {
		testSuiteName = "headscale-test-suite-" + runID
	}

	err = dockertestutil.AddContainerToNetwork(s.pool, network, testSuiteName)
	if err != nil {
		return nil, fmt.Errorf("adding test suite container to network: %w", err)
	}

	mak.Set(&s.networks, name, network)

	return network, nil
}

func (s *Scenario) Networks() []*dockertest.Network {
	if len(s.networks) == 0 {
		panic("Scenario.Networks called with empty network list")
	}

	return xmaps.Values(s.networks)
}

func (s *Scenario) Network(name string) (*dockertest.Network, error) {
	net, ok := s.networks[s.prefixedNetworkName(name)]
	if !ok {
		return nil, fmt.Errorf("no network named: %s", name) //nolint:err113
	}

	return net, nil
}

func (s *Scenario) SubnetOfNetwork(name string) (*netip.Prefix, error) {
	net, ok := s.networks[s.prefixedNetworkName(name)]
	if !ok {
		return nil, fmt.Errorf("no network named: %s", name) //nolint:err113
	}

	if len(net.Network.IPAM.Config) == 0 {
		return nil, fmt.Errorf("no IPAM config found in network: %s", name) //nolint:err113
	}

	pref, err := netip.ParsePrefix(net.Network.IPAM.Config[0].Subnet)
	if err != nil {
		return nil, err
	}

	return &pref, nil
}

func (s *Scenario) Services(name string) ([]*dockertest.Resource, error) {
	res, ok := s.extraServices[s.prefixedNetworkName(name)]
	if !ok {
		return nil, fmt.Errorf("no network named: %s", name) //nolint:err113
	}

	return res, nil
}

func (s *Scenario) ShutdownAssertNoPanics(t *testing.T) {
	t.Helper()

	defer func() { _ = dockertestutil.CleanUnreferencedNetworks(s.pool) }()
	defer func() { _ = dockertestutil.CleanImagesInCI(s.pool) }()

	s.controlServers.Range(func(_ string, control ControlServer) bool {
		stdoutPath, stderrPath, err := control.Shutdown()
		if err != nil {
			log.Printf(
				"shutting down control: %s",
				fmt.Errorf("tearing down control: %w", err),
			)
		}

		if t != nil {
			stdout, err := os.ReadFile(stdoutPath)
			require.NoError(t, err)
			assert.NotContains(t, string(stdout), "panic")

			stderr, err := os.ReadFile(stderrPath)
			require.NoError(t, err)
			assert.NotContains(t, string(stderr), "panic")
		}

		return true
	})

	s.mu.Lock()

	for userName, user := range s.users {
		for _, client := range user.Clients {
			log.Printf("removing client %s in user %s", client.Hostname(), userName)

			stdoutPath, stderrPath, err := client.Shutdown()
			if err != nil {
				log.Printf("tearing down client: %s", err)
			}

			if t != nil {
				stdout, err := os.ReadFile(stdoutPath)
				require.NoError(t, err)
				assert.NotContains(t, string(stdout), "panic")

				stderr, err := os.ReadFile(stderrPath)
				require.NoError(t, err)
				assert.NotContains(t, string(stderr), "panic")
			}
		}
	}

	s.mu.Unlock()

	for _, derp := range s.derpServers {
		err := derp.Shutdown()
		if err != nil {
			log.Printf("tearing down derp server: %s", err)
		}
	}

	for _, svcs := range s.extraServices {
		for _, svc := range svcs {
			err := svc.Close()
			if err != nil {
				log.Printf("tearing down service %q: %s", svc.Container.Name, err)
			}
		}
	}

	if s.mockOIDC.r != nil {
		s.mockOIDC.r.Close()

		err := s.mockOIDC.r.Close()
		if err != nil {
			log.Printf("tearing down oidc server: %s", err)
		}
	}

	for _, network := range s.networks {
		err := network.Close()
		if err != nil {
			log.Printf("tearing down network: %s", err)
		}
	}
}

// Shutdown shuts down and cleans up all the containers (ControlServer, TailscaleClient)
// and networks associated with it.
// In addition, it will save the logs of the ControlServer to `/tmp/control` in the
// environment running the tests.
func (s *Scenario) Shutdown() {
	s.ShutdownAssertNoPanics(nil)
}

// Users returns the name of all users associated with the Scenario.
func (s *Scenario) Users() []string {
	users := make([]string, 0, len(s.users))
	for user := range s.users {
		users = append(users, user)
	}

	return users
}

/// Headscale related stuff
// Note: These functions assume that there is a _single_ headscale instance for now

// Headscale returns a ControlServer instance based on hsic (HeadscaleInContainer)
// If the Scenario already has an instance, the pointer to the running container
// will be return, otherwise a new instance will be created.
// TODO(kradalby): make port and headscale configurable, multiple instances support?
func (s *Scenario) Headscale(opts ...hsic.Option) (ControlServer, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if headscale, ok := s.controlServers.Load("headscale"); ok {
		return headscale, nil
	}

	if usePostgresForTest {
		opts = append(opts, hsic.WithPostgres())
	}

	headscale, err := hsic.New(s.pool, s.Networks(), opts...)
	if err != nil {
		return nil, fmt.Errorf("creating headscale container: %w", err)
	}

	err = headscale.WaitForRunning()
	if err != nil {
		return nil, fmt.Errorf("reaching headscale container: %w", err)
	}

	s.controlServers.Store("headscale", headscale)

	return headscale, nil
}

// Pool returns the dockertest pool for the scenario.
func (s *Scenario) Pool() *dockertest.Pool {
	return s.pool
}

// GetOrCreateUser gets or creates a user in the scenario.
func (s *Scenario) GetOrCreateUser(userStr string) *User {
	s.mu.Lock()
	defer s.mu.Unlock()

	if user, ok := s.users[userStr]; ok {
		return user
	}

	user := &User{
		Clients: make(map[string]TailscaleClient),
	}
	s.users[userStr] = user

	return user
}

// CreatePreAuthKey creates a "pre authentorised key" to be created in the
// Headscale instance on behalf of the Scenario.
func (s *Scenario) CreatePreAuthKey(
	user uint64,
	reusable bool,
	ephemeral bool,
) (*v1.PreAuthKey, error) {
	if headscale, err := s.Headscale(); err == nil { //nolint:noinlineerr
		key, err := headscale.CreateAuthKey(user, reusable, ephemeral)
		if err != nil {
			return nil, fmt.Errorf("creating user: %w", err)
		}

		return key, nil
	}

	return nil, fmt.Errorf("creating user: %w", errNoHeadscaleAvailable)
}

// CreatePreAuthKeyWithOptions creates a "pre authorised key" with the specified options
// to be created in the Headscale instance on behalf of the Scenario.
func (s *Scenario) CreatePreAuthKeyWithOptions(opts hsic.AuthKeyOptions) (*v1.PreAuthKey, error) {
	headscale, err := s.Headscale()
	if err != nil {
		return nil, fmt.Errorf("creating preauth key with options: %w", errNoHeadscaleAvailable)
	}

	key, err := headscale.CreateAuthKeyWithOptions(opts)
	if err != nil {
		return nil, fmt.Errorf("creating preauth key with options: %w", err)
	}

	return key, nil
}

// CreatePreAuthKeyWithTags creates a "pre authorised key" with the specified tags
// to be created in the Headscale instance on behalf of the Scenario.
func (s *Scenario) CreatePreAuthKeyWithTags(
	user uint64,
	reusable bool,
	ephemeral bool,
	tags []string,
) (*v1.PreAuthKey, error) {
	headscale, err := s.Headscale()
	if err != nil {
		return nil, fmt.Errorf("creating preauth key with tags: %w", errNoHeadscaleAvailable)
	}

	key, err := headscale.CreateAuthKeyWithTags(user, reusable, ephemeral, tags)
	if err != nil {
		return nil, fmt.Errorf("creating preauth key with tags: %w", err)
	}

	return key, nil
}

// CreateUser creates a User to be created in the
// Headscale instance on behalf of the Scenario.
func (s *Scenario) CreateUser(user string) (*v1.User, error) {
	if headscale, err := s.Headscale(); err == nil { //nolint:noinlineerr
		u, err := headscale.CreateUser(user)
		if err != nil {
			return nil, fmt.Errorf("creating user: %w", err)
		}

		s.mu.Lock()
		s.users[user] = &User{
			Clients: make(map[string]TailscaleClient),
		}
		s.mu.Unlock()

		return u, nil
	}

	return nil, fmt.Errorf("creating user: %w", errNoHeadscaleAvailable)
}

/// Client related stuff

func (s *Scenario) CreateTailscaleNode(
	version string,
	opts ...tsic.Option,
) (TailscaleClient, error) {
	headscale, err := s.Headscale()
	if err != nil {
		return nil, fmt.Errorf("creating tailscale node (version: %s): %w", version, err)
	}

	cert := headscale.GetCert()
	hostname := headscale.GetHostname()

	s.mu.Lock()
	defer s.mu.Unlock()

	opts = append(opts,
		tsic.WithCACert(cert),
		tsic.WithHeadscaleName(hostname),
	)

	tsClient, err := tsic.New(
		s.pool,
		version,
		opts...,
	)
	if err != nil {
		return nil, fmt.Errorf(
			"creating tailscale node: %w",
			err,
		)
	}

	err = tsClient.WaitForNeedsLogin(integrationutil.PeerSyncTimeout())
	if err != nil {
		return nil, fmt.Errorf(
			"waiting for tailscaled (%s) to need login: %w",
			tsClient.Hostname(),
			err,
		)
	}

	return tsClient, nil
}

// CreateTailscaleNodesInUser creates and adds a new TailscaleClient to a
// User in the Scenario.
func (s *Scenario) CreateTailscaleNodesInUser(
	userStr string,
	requestedVersion string,
	count int,
	opts ...tsic.Option,
) error {
	if user, ok := s.users[userStr]; ok {
		var versions []string

		for i := range count {
			version := requestedVersion
			if requestedVersion == "all" {
				if s.spec.Versions != nil {
					version = s.spec.Versions[i%len(s.spec.Versions)]
				} else {
					version = MustTestVersions[i%len(MustTestVersions)]
				}
			}

			versions = append(versions, version)

			headscale, err := s.Headscale()
			if err != nil {
				return fmt.Errorf("creating tailscale node (version: %s): %w", version, err)
			}

			cert := headscale.GetCert()
			hostname := headscale.GetHostname()

			// Determine which network this tailscale client will be in
			var network *dockertest.Network
			if s.userToNetwork != nil && s.userToNetwork[userStr] != nil {
				network = s.userToNetwork[userStr]
			} else {
				network = s.networks[s.testDefaultNetwork]
			}

			// Get headscale IP in this network for /etc/hosts fallback DNS
			headscaleIP := headscale.GetIPInNetwork(network)
			extraHosts := []string{hostname + ":" + headscaleIP}

			s.mu.Lock()

			opts = append(opts,
				tsic.WithCACert(cert),
				tsic.WithHeadscaleName(hostname),
				tsic.WithExtraHosts(extraHosts),
			)

			s.mu.Unlock()

			user.createWaitGroup.Go(func() error {
				s.mu.Lock()
				tsClient, err := tsic.New(
					s.pool,
					version,
					opts...,
				)
				s.mu.Unlock()

				if err != nil {
					return fmt.Errorf(
						"creating tailscale node: %w",
						err,
					)
				}

				err = tsClient.WaitForNeedsLogin(integrationutil.PeerSyncTimeout())
				if err != nil {
					return fmt.Errorf(
						"waiting for tailscaled (%s) to need login: %w",
						tsClient.Hostname(),
						err,
					)
				}

				s.mu.Lock()

				user.Clients[tsClient.Hostname()] = tsClient

				s.mu.Unlock()

				return nil
			})
		}

		err := user.createWaitGroup.Wait()
		if err != nil {
			return err
		}

		log.Printf("testing versions %v, MustTestVersions %v", lo.Uniq(versions), MustTestVersions)

		return nil
	}

	return fmt.Errorf("adding tailscale node: %w", errNoUserAvailable)
}

// RunTailscaleUp will log in all of the TailscaleClients associated with a
// User to the given ControlServer (by URL).
func (s *Scenario) RunTailscaleUp(
	userStr, loginServer, authKey string,
) error {
	if user, ok := s.users[userStr]; ok {
		for _, client := range user.Clients {
			c := client

			user.joinWaitGroup.Go(func() error {
				return c.Login(loginServer, authKey)
			})
		}

		err := user.joinWaitGroup.Wait()
		if err != nil {
			return err
		}

		for _, client := range user.Clients {
			err := client.WaitForRunning(integrationutil.PeerSyncTimeout())
			if err != nil {
				return fmt.Errorf("%s bringing up tailscale node: %w", client.Hostname(), err)
			}
		}

		return nil
	}

	return fmt.Errorf("bringing up tailscale node: %w", errNoUserAvailable)
}

// CountTailscale returns the total number of TailscaleClients in a Scenario.
// This is the sum of Users x TailscaleClients.
func (s *Scenario) CountTailscale() int {
	count := 0

	for _, user := range s.users {
		count += len(user.Clients)
	}

	return count
}

// WaitForTailscaleSync blocks execution until all the TailscaleClient reports
// to have all other TailscaleClients present in their netmap.NetworkMap.
func (s *Scenario) WaitForTailscaleSync() error {
	tsCount := s.CountTailscale()

	err := s.WaitForTailscaleSyncWithPeerCount(tsCount-1, integrationutil.PeerSyncTimeout(), integrationutil.PeerSyncRetryInterval())
	if err != nil {
		for _, user := range s.users {
			for _, client := range user.Clients {
				peers, allOnline, _ := client.FailingPeersAsString()
				if !allOnline {
					log.Println(peers)
				}
			}
		}
	}

	return err
}

// WaitForTailscaleSyncPerUser blocks execution until each TailscaleClient has the expected
// number of peers for its user. This is useful for policies like autogroup:self where nodes
// only see same-user peers, not all nodes in the network.
func (s *Scenario) WaitForTailscaleSyncPerUser(timeout, retryInterval time.Duration) error {
	var allErrors []error

	for _, user := range s.users {
		// Calculate expected peer count: number of nodes in this user minus 1 (self)
		expectedPeers := len(user.Clients) - 1

		for _, client := range user.Clients {
			c := client
			expectedCount := expectedPeers

			user.syncWaitGroup.Go(func() error {
				return c.WaitForPeers(expectedCount, timeout, retryInterval)
			})
		}

		err := user.syncWaitGroup.Wait()
		if err != nil {
			allErrors = append(allErrors, err)
		}
	}

	if len(allErrors) > 0 {
		return multierr.New(allErrors...)
	}

	return nil
}

// WaitForTailscaleSyncWithPeerCount blocks execution until all the TailscaleClient reports
// to have all other TailscaleClients present in their netmap.NetworkMap.
func (s *Scenario) WaitForTailscaleSyncWithPeerCount(peerCount int, timeout, retryInterval time.Duration) error {
	var allErrors []error

	for _, user := range s.users {
		for _, client := range user.Clients {
			c := client

			user.syncWaitGroup.Go(func() error {
				return c.WaitForPeers(peerCount, timeout, retryInterval)
			})
		}

		err := user.syncWaitGroup.Wait()
		if err != nil {
			allErrors = append(allErrors, err)
		}
	}

	if len(allErrors) > 0 {
		return multierr.New(allErrors...)
	}

	return nil
}

func (s *Scenario) CreateHeadscaleEnvWithLoginURL(
	tsOpts []tsic.Option,
	opts ...hsic.Option,
) error {
	return s.createHeadscaleEnv(true, tsOpts, opts...)
}

func (s *Scenario) CreateHeadscaleEnv(
	tsOpts []tsic.Option,
	opts ...hsic.Option,
) error {
	return s.createHeadscaleEnv(false, tsOpts, opts...)
}

// CreateHeadscaleEnv starts the headscale environment and the clients
// according to the ScenarioSpec passed to the Scenario.
func (s *Scenario) createHeadscaleEnv(
	withURL bool,
	tsOpts []tsic.Option,
	opts ...hsic.Option,
) error {
	return s.createHeadscaleEnvWithTags(withURL, tsOpts, nil, "", opts...)
}

// createHeadscaleEnvWithTags starts the headscale environment and the clients
// according to the ScenarioSpec passed to the Scenario. If preAuthKeyTags is
// non-empty and withURL is false, the tags will be applied to the PreAuthKey
// (tags-as-identity model).
//
// For webauth (withURL=true), if webauthTagUser is non-empty and preAuthKeyTags
// is non-empty, only nodes belonging to that user will request tags via
// --advertise-tags. This is necessary because tagOwners ACL controls which
// users can request specific tags.
func (s *Scenario) createHeadscaleEnvWithTags(
	withURL bool,
	tsOpts []tsic.Option,
	preAuthKeyTags []string,
	webauthTagUser string,
	opts ...hsic.Option,
) error {
	headscale, err := s.Headscale(opts...)
	if err != nil {
		return err
	}

	for _, user := range s.spec.Users {
		var u *v1.User

		if s.spec.OIDCSkipUserCreation {
			// Only register locally — OIDC login will create the headscale user.
			s.mu.Lock()
			s.users[user] = &User{Clients: make(map[string]TailscaleClient)}
			s.mu.Unlock()
		} else {
			u, err = s.CreateUser(user)
			if err != nil {
				return err
			}
		}

		var userOpts []tsic.Option
		if s.userToNetwork != nil {
			userOpts = append(tsOpts, tsic.WithNetwork(s.userToNetwork[user]))
		} else {
			userOpts = append(tsOpts, tsic.WithNetwork(s.networks[s.testDefaultNetwork]))
		}

		// For webauth with tags, only apply tags to the specified webauthTagUser
		// (other users may not be authorized via tagOwners)
		if withURL && webauthTagUser != "" && len(preAuthKeyTags) > 0 && user == webauthTagUser {
			userOpts = append(userOpts, tsic.WithTags(preAuthKeyTags))
		}

		err = s.CreateTailscaleNodesInUser(user, "all", s.spec.NodesPerUser, userOpts...)
		if err != nil {
			return err
		}

		if withURL {
			err = s.RunTailscaleUpWithURL(user, headscale.GetEndpoint())
			if err != nil {
				return err
			}
		} else {
			// Use tagged PreAuthKey if tags are provided (tags-as-identity model)
			var key *v1.PreAuthKey
			if len(preAuthKeyTags) > 0 {
				key, err = s.CreatePreAuthKeyWithTags(u.GetId(), true, false, preAuthKeyTags)
			} else {
				key, err = s.CreatePreAuthKey(u.GetId(), true, false)
			}

			if err != nil {
				return err
			}

			err = s.RunTailscaleUp(user, headscale.GetEndpoint(), key.GetKey())
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (s *Scenario) RunTailscaleUpWithURL(userStr, loginServer string) error {
	log.Printf("running tailscale up for user %s", userStr)

	if user, ok := s.users[userStr]; ok {
		for _, client := range user.Clients {
			tsc := client

			user.joinWaitGroup.Go(func() error {
				loginURL, err := tsc.LoginWithURL(loginServer)
				if err != nil {
					log.Printf("%s running tailscale up: %s", tsc.Hostname(), err)
				}

				body, err := doLoginURL(tsc.Hostname(), loginURL)
				if err != nil {
					return err
				}

				// If the URL is not a OIDC URL, then we need to
				// run the register command to fully log in the client.
				if !strings.Contains(loginURL.String(), "/oidc/") {
					_ = s.runHeadscaleRegister(userStr, body)
				}

				return nil
			})

			log.Printf("client %s is ready", client.Hostname())
		}

		err := user.joinWaitGroup.Wait()
		if err != nil {
			return err
		}

		for _, client := range user.Clients {
			err := client.WaitForRunning(integrationutil.PeerSyncTimeout())
			if err != nil {
				return fmt.Errorf(
					"%s tailscale node has not reached running: %w",
					client.Hostname(),
					err,
				)
			}
		}

		return nil
	}

	return fmt.Errorf("bringing up tailscale node: %w", errNoUserAvailable)
}

type debugJar struct {
	inner *cookiejar.Jar
	mu    sync.RWMutex
	store map[string]map[string]map[string]*http.Cookie // domain -> path -> name -> cookie
}

func newDebugJar() (*debugJar, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, err
	}

	return &debugJar{
		inner: jar,
		store: make(map[string]map[string]map[string]*http.Cookie),
	}, nil
}

func (j *debugJar) SetCookies(u *url.URL, cookies []*http.Cookie) {
	j.inner.SetCookies(u, cookies)

	j.mu.Lock()
	defer j.mu.Unlock()

	for _, c := range cookies {
		if c == nil || c.Name == "" {
			continue
		}

		domain := c.Domain
		if domain == "" {
			domain = u.Hostname()
		}

		path := c.Path
		if path == "" {
			path = "/"
		}

		if _, ok := j.store[domain]; !ok {
			j.store[domain] = make(map[string]map[string]*http.Cookie)
		}

		if _, ok := j.store[domain][path]; !ok {
			j.store[domain][path] = make(map[string]*http.Cookie)
		}

		j.store[domain][path][c.Name] = copyCookie(c)
	}
}

func (j *debugJar) Cookies(u *url.URL) []*http.Cookie {
	return j.inner.Cookies(u)
}

func (j *debugJar) Dump(w io.Writer) {
	j.mu.RLock()
	defer j.mu.RUnlock()

	for domain, paths := range j.store {
		fmt.Fprintf(w, "Domain: %s\n", domain)

		for path, byName := range paths {
			fmt.Fprintf(w, "  Path: %s\n", path)

			for _, c := range byName {
				fmt.Fprintf(
					w, "    %s=%s; Expires=%v; Secure=%v; HttpOnly=%v; SameSite=%v\n",
					c.Name, c.Value, c.Expires, c.Secure, c.HttpOnly, c.SameSite,
				)
			}
		}
	}
}

func copyCookie(c *http.Cookie) *http.Cookie {
	cc := *c
	return &cc
}

func newLoginHTTPClient(hostname string) (*http.Client, error) {
	hc := &http.Client{
		Transport: LoggingRoundTripper{Hostname: hostname},
	}

	jar, err := newDebugJar()
	if err != nil {
		return nil, fmt.Errorf("%s creating cookiejar: %w", hostname, err)
	}

	hc.Jar = jar

	return hc, nil
}

// doLoginURL visits the given login URL and returns the body as a string.
func doLoginURL(hostname string, loginURL *url.URL) (string, error) {
	log.Printf("%s login url: %s\n", hostname, loginURL.String())

	hc, err := newLoginHTTPClient(hostname)
	if err != nil {
		return "", err
	}

	body, _, err := doLoginURLWithClient(hostname, loginURL, hc, true)
	if err != nil {
		return "", err
	}

	return body, nil
}

// doLoginURLWithClient performs the login request using the provided HTTP client.
// When followRedirects is false, it will return the first redirect without following it.
func doLoginURLWithClient(hostname string, loginURL *url.URL, hc *http.Client, followRedirects bool) (
	string,
	*url.URL,
	error,
) {
	if hc == nil {
		return "", nil, fmt.Errorf("%s http client is nil", hostname) //nolint:err113
	}

	if loginURL == nil {
		return "", nil, fmt.Errorf("%s login url is nil", hostname) //nolint:err113
	}

	log.Printf("%s logging in with url: %s", hostname, loginURL.String())

	ctx := context.Background()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, loginURL.String(), nil)
	if err != nil {
		return "", nil, fmt.Errorf("%s creating http request: %w", hostname, err)
	}

	originalRedirect := hc.CheckRedirect
	if !followRedirects {
		hc.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		}
	}

	defer func() {
		hc.CheckRedirect = originalRedirect
	}()

	resp, err := hc.Do(req)
	if err != nil {
		return "", nil, fmt.Errorf("%s sending http request: %w", hostname, err)
	}
	defer resp.Body.Close()

	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", nil, fmt.Errorf("%s reading response body: %w", hostname, err)
	}

	body := string(bodyBytes)

	var redirectURL *url.URL
	if resp.StatusCode >= http.StatusMultipleChoices && resp.StatusCode < http.StatusBadRequest {
		redirectURL, err = resp.Location()
		if err != nil {
			return body, nil, fmt.Errorf("%s resolving redirect location: %w", hostname, err)
		}
	}

	if followRedirects && resp.StatusCode != http.StatusOK {
		log.Printf("body: %s", body)

		return body, redirectURL, fmt.Errorf("%s unexpected status code %d", hostname, resp.StatusCode) //nolint:err113
	}

	if resp.StatusCode >= http.StatusBadRequest {
		log.Printf("body: %s", body)

		return body, redirectURL, fmt.Errorf("%s unexpected status code %d", hostname, resp.StatusCode) //nolint:err113
	}

	if hc.Jar != nil {
		if jar, ok := hc.Jar.(*debugJar); ok {
			jar.Dump(os.Stdout)
		} else {
			log.Printf("cookies: %+v", hc.Jar.Cookies(loginURL))
		}
	}

	return body, redirectURL, nil
}

var errParseAuthPage = errors.New("parsing auth page")

func (s *Scenario) runHeadscaleRegister(userStr string, body string) error {
	// see api.go HTML template
	codeSep := strings.Split(body, "</code>")
	if len(codeSep) != 2 {
		return errParseAuthPage
	}

	keySep := strings.Split(codeSep[0], "--auth-id ")
	if len(keySep) != 2 {
		return errParseAuthPage
	}

	key := keySep[1]
	key = strings.SplitN(key, " ", 2)[0]
	log.Printf("registering node %s", key)

	if headscale, err := s.Headscale(); err == nil { //nolint:noinlineerr
		_, err = headscale.Execute(
			[]string{"headscale", "auth", "register", "--user", userStr, "--auth-id", key},
		)
		if err != nil {
			log.Printf("registering node: %s", err)

			return err
		}

		return nil
	}

	return fmt.Errorf("finding headscale: %w", errNoHeadscaleAvailable)
}

type LoggingRoundTripper struct {
	Hostname string
}

func (t LoggingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	noTls := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, // nolint
	}

	resp, err := noTls.RoundTrip(req)
	if err != nil {
		return nil, err
	}

	log.Printf(`
---
%s - method: %s | url: %s
%s - status: %d | cookies: %+v
---
`, t.Hostname, req.Method, req.URL.String(), t.Hostname, resp.StatusCode, resp.Cookies())

	return resp, nil
}

// GetIPs returns all netip.Addr of TailscaleClients associated with a User
// in a Scenario.
func (s *Scenario) GetIPs(user string) ([]netip.Addr, error) {
	var ips []netip.Addr

	if ns, ok := s.users[user]; ok {
		for _, client := range ns.Clients {
			clientIps, err := client.IPs()
			if err != nil {
				return ips, fmt.Errorf("getting IPs: %w", err)
			}

			ips = append(ips, clientIps...)
		}

		return ips, nil
	}

	return ips, fmt.Errorf("getting IPs: %w", errNoUserAvailable)
}

// GetClients returns all TailscaleClients associated with a User in a Scenario.
func (s *Scenario) GetClients(user string) ([]TailscaleClient, error) {
	var clients []TailscaleClient

	if ns, ok := s.users[user]; ok {
		for _, client := range ns.Clients {
			clients = append(clients, client)
		}

		return clients, nil
	}

	return clients, fmt.Errorf("getting clients: %w", errNoUserAvailable)
}

// ListTailscaleClients returns a list of TailscaleClients given the Users
// passed as parameters.
func (s *Scenario) ListTailscaleClients(users ...string) ([]TailscaleClient, error) {
	var allClients []TailscaleClient

	if len(users) == 0 {
		users = s.Users()
	}

	for _, user := range users {
		clients, err := s.GetClients(user)
		if err != nil {
			return nil, err
		}

		allClients = append(allClients, clients...)
	}

	return allClients, nil
}

// FindTailscaleClientByIP returns a TailscaleClient associated with an IP address
// if it exists.
func (s *Scenario) FindTailscaleClientByIP(ip netip.Addr) (TailscaleClient, error) {
	clients, err := s.ListTailscaleClients()
	if err != nil {
		return nil, err
	}

	for _, client := range clients {
		ips, _ := client.IPs()
		if slices.Contains(ips, ip) {
			return client, nil
		}
	}

	return nil, errNoClientFound
}

// ListTailscaleClientsIPs returns a list of netip.Addr based on Users
// passed as parameters.
func (s *Scenario) ListTailscaleClientsIPs(users ...string) ([]netip.Addr, error) {
	var allIps []netip.Addr

	if len(users) == 0 {
		users = s.Users()
	}

	for _, user := range users {
		ips, err := s.GetIPs(user)
		if err != nil {
			return nil, err
		}

		allIps = append(allIps, ips...)
	}

	return allIps, nil
}

// ListTailscaleClientsFQDNs returns a list of FQDN based on Users
// passed as parameters.
func (s *Scenario) ListTailscaleClientsFQDNs(users ...string) ([]string, error) {
	allFQDNs := make([]string, 0)

	clients, err := s.ListTailscaleClients(users...)
	if err != nil {
		return nil, err
	}

	for _, client := range clients {
		fqdn, err := client.FQDN()
		if err != nil {
			return nil, err
		}

		allFQDNs = append(allFQDNs, fqdn)
	}

	return allFQDNs, nil
}

// WaitForTailscaleLogout blocks execution until all TailscaleClients have
// logged out of the ControlServer.
func (s *Scenario) WaitForTailscaleLogout() error {
	for _, user := range s.users {
		for _, client := range user.Clients {
			c := client

			user.syncWaitGroup.Go(func() error {
				return c.WaitForNeedsLogin(integrationutil.PeerSyncTimeout())
			})
		}

		err := user.syncWaitGroup.Wait()
		if err != nil {
			return err
		}
	}

	return nil
}

// CreateDERPServer creates a new DERP server in a container.
func (s *Scenario) CreateDERPServer(version string, opts ...dsic.Option) (*dsic.DERPServerInContainer, error) {
	derp, err := dsic.New(s.pool, version, s.Networks(), opts...)
	if err != nil {
		return nil, fmt.Errorf("creating DERP server: %w", err)
	}

	err = derp.WaitForRunning()
	if err != nil {
		return nil, fmt.Errorf("reaching DERP server: %w", err)
	}

	s.derpServers = append(s.derpServers, derp)

	return derp, nil
}

type scenarioOIDC struct {
	r   *dockertest.Resource
	cfg *types.OIDCConfig
}

func (o *scenarioOIDC) Issuer() string {
	if o.cfg == nil {
		panic("OIDC has not been created")
	}

	return o.cfg.Issuer
}

func (o *scenarioOIDC) ClientSecret() string {
	if o.cfg == nil {
		panic("OIDC has not been created")
	}

	return o.cfg.ClientSecret
}

func (o *scenarioOIDC) ClientID() string {
	if o.cfg == nil {
		panic("OIDC has not been created")
	}

	return o.cfg.ClientID
}

const (
	dockerContextPath      = "../."
	hsicOIDCMockHashLength = 6
	defaultAccessTTL       = 10 * time.Minute
)

var errStatusCodeNotOK = errors.New("status code not OK")

func (s *Scenario) runMockOIDC(accessTTL time.Duration, users []mockoidc.MockUser) error {
	port, err := dockertestutil.RandomFreeHostPort()
	if err != nil {
		log.Fatalf("finding open port: %s", err)
	}

	portNotation := fmt.Sprintf("%d/tcp", port)

	hash, _ := util.GenerateRandomStringDNSSafe(hsicOIDCMockHashLength)

	hostname := "hs-oidcmock-" + hash

	usersJSON, err := json.Marshal(users)
	if err != nil {
		return err
	}

	mockOidcOptions := &dockertest.RunOptions{
		Name:         hostname,
		Cmd:          []string{"headscale", "mockoidc"},
		ExposedPorts: []string{portNotation},
		PortBindings: map[docker.Port][]docker.PortBinding{
			docker.Port(portNotation): {{HostPort: strconv.Itoa(port)}},
		},
		Networks: s.Networks(),
		Env: []string{
			"MOCKOIDC_ADDR=" + hostname,
			fmt.Sprintf("MOCKOIDC_PORT=%d", port),
			"MOCKOIDC_CLIENT_ID=superclient",
			"MOCKOIDC_CLIENT_SECRET=supersecret",
			"MOCKOIDC_ACCESS_TTL=" + accessTTL.String(),
			"MOCKOIDC_USERS=" + string(usersJSON),
		},
	}

	headscaleBuildOptions := &dockertest.BuildOptions{
		Dockerfile: hsic.IntegrationTestDockerFileName,
		ContextDir: dockerContextPath,
	}

	err = s.pool.RemoveContainerByName(hostname)
	if err != nil {
		return err
	}

	s.mockOIDC = scenarioOIDC{}

	// Add integration test labels if running under hi tool
	dockertestutil.DockerAddIntegrationLabels(mockOidcOptions, "oidc")

	if pmockoidc, err := s.pool.BuildAndRunWithBuildOptions( //nolint:noinlineerr
		headscaleBuildOptions,
		mockOidcOptions,
		dockertestutil.DockerRestartPolicy); err == nil {
		s.mockOIDC.r = pmockoidc
	} else {
		return err
	}

	// headscale needs to set up the provider with a specific
	// IP addr to ensure we get the correct config from the well-known
	// endpoint.
	network := s.Networks()[0]
	ipAddr := s.mockOIDC.r.GetIPInNetwork(network)

	log.Println("Waiting for headscale mock oidc to be ready for tests")

	hostEndpoint := net.JoinHostPort(ipAddr, strconv.Itoa(port))

	if err := s.pool.Retry(func() error { //nolint:noinlineerr
		oidcConfigURL := fmt.Sprintf("http://%s/oidc/.well-known/openid-configuration", hostEndpoint)
		httpClient := &http.Client{}
		ctx := context.Background()
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, oidcConfigURL, nil)

		resp, err := httpClient.Do(req)
		if err != nil {
			log.Printf("headscale mock OIDC tests is not ready: %s\n", err)

			return err
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			return errStatusCodeNotOK
		}

		return nil
	}); err != nil {
		return err
	}

	s.mockOIDC.cfg = &types.OIDCConfig{
		Issuer: fmt.Sprintf(
			"http://%s/oidc",
			hostEndpoint,
		),
		ClientID:                   "superclient",
		ClientSecret:               "supersecret",
		OnlyStartIfOIDCIsAvailable: true,
	}

	log.Printf("headscale mock oidc is ready for tests at %s", hostEndpoint)

	return nil
}

type extraServiceFunc func(*Scenario, string) (*dockertest.Resource, error)

func Webservice(s *Scenario, networkName string) (*dockertest.Resource, error) {
	// port, err := dockertestutil.RandomFreeHostPort()
	// if err != nil {
	// 	log.Fatalf("finding open port: %s", err)
	// }
	// portNotation := fmt.Sprintf("%d/tcp", port)
	hash := util.MustGenerateRandomStringDNSSafe(hsicOIDCMockHashLength)

	hostname := "hs-webservice-" + hash

	network, ok := s.networks[s.prefixedNetworkName(networkName)]
	if !ok {
		return nil, fmt.Errorf("network does not exist: %s", networkName) //nolint:err113
	}

	webOpts := &dockertest.RunOptions{
		Name: hostname,
		Cmd:  []string{"/bin/sh", "-c", "cd / ; python3 -m http.server --bind :: 80"},
		// ExposedPorts: []string{portNotation},
		// PortBindings: map[docker.Port][]docker.PortBinding{
		// 	docker.Port(portNotation): {{HostPort: strconv.Itoa(port)}},
		// },
		Networks: []*dockertest.Network{network},
		Env:      []string{},
	}

	// Add integration test labels if running under hi tool
	dockertestutil.DockerAddIntegrationLabels(webOpts, "web")

	webBOpts := &dockertest.BuildOptions{
		Dockerfile: hsic.IntegrationTestDockerFileName,
		ContextDir: dockerContextPath,
	}

	web, err := s.pool.BuildAndRunWithBuildOptions(
		webBOpts,
		webOpts,
		dockertestutil.DockerRestartPolicy)
	if err != nil {
		return nil, err
	}

	// headscale needs to set up the provider with a specific
	// IP addr to ensure we get the correct config from the well-known
	// endpoint.
	// ipAddr := web.GetIPInNetwork(network)

	// log.Println("Waiting for headscale mock oidc to be ready for tests")
	// hostEndpoint := net.JoinHostPort(ipAddr, strconv.Itoa(port))

	// if err := s.pool.Retry(func() error {
	// 	oidcConfigURL := fmt.Sprintf("http://%s/etc/hostname", hostEndpoint)
	// 	httpClient := &http.Client{}
	// 	ctx := context.Background()
	// 	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, oidcConfigURL, nil)
	// 	resp, err := httpClient.Do(req)
	// 	if err != nil {
	// 		log.Printf("headscale mock OIDC tests is not ready: %s\n", err)

	// 		return err
	// 	}
	// 	defer resp.Body.Close()

	// 	if resp.StatusCode != http.StatusOK {
	// 		return errStatusCodeNotOK
	// 	}

	// 	return nil
	// }); err != nil {
	// 	return err
	// }

	return web, nil
}