cmd/hi: fix extration of test data

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

cmd/hi/docker.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -17,6 +18,7 @@ import (
 	"github.com/docker/docker/api/types/image"
 	"github.com/docker/docker/api/types/mount"
 	"github.com/docker/docker/client"
+	"github.com/docker/docker/pkg/stdcopy"
 )
 
 var (
@@ -89,6 +91,17 @@ func runTestContainer(ctx context.Context, config *RunConfig) error {
 
 	exitCode, err := streamAndWait(ctx, cli, resp.ID)
 
+	// Give the container a moment to flush any final artifacts
+	time.Sleep(2 * time.Second)
+
+	// Extract artifacts from test containers before cleanup
+	if err := extractArtifactsFromContainers(ctx, resp.ID, logsDir, config.Verbose); err != nil && config.Verbose {
+		log.Printf("Warning: failed to extract artifacts from containers: %v", err)
+	}
+
+	// Always list control files regardless of test outcome
+	listControlFiles(logsDir)
+
 	shouldCleanup := config.CleanAfter && (!config.KeepOnFailure || exitCode == 0)
 	if shouldCleanup {
 		if config.Verbose {
@@ -108,7 +121,6 @@ func runTestContainer(ctx context.Context, config *RunConfig) error {
 	}
 
 	log.Printf("Test completed successfully!")
-	listControlFiles(logsDir)
 
 	return nil
 }
@@ -140,23 +152,36 @@ func createGoTestContainer(ctx context.Context, cli *client.Client, config *RunC
 
 	projectRoot := findProjectRoot(pwd)
 
+	runID := generateRunIDFromContainerName(containerName)
+	
 	env := []string{
 		fmt.Sprintf("HEADSCALE_INTEGRATION_POSTGRES=%d", boolToInt(config.UsePostgres)),
+		fmt.Sprintf("HEADSCALE_INTEGRATION_RUN_ID=%s", runID),
 	}
-
 	containerConfig := &container.Config{
 		Image:      "golang:" + config.GoVersion,
 		Cmd:        goTestCmd,
 		Env:        env,
 		WorkingDir: projectRoot + "/integration",
 		Tty:        true,
+		Labels: map[string]string{
+			"hi.run-id":    runID,
+			"hi.test-type": "test-runner",
+		},
+	}
+
+	// Get the correct Docker socket path from the current context
+	dockerSocketPath := getDockerSocketPath()
+	
+	if config.Verbose {
+		log.Printf("Using Docker socket: %s", dockerSocketPath)
 	}
 
 	hostConfig := &container.HostConfig{
 		AutoRemove: false, // We'll remove manually for better control
 		Binds: []string{
 			fmt.Sprintf("%s:%s", projectRoot, projectRoot),
-			"/var/run/docker.sock:/var/run/docker.sock",
+			fmt.Sprintf("%s:/var/run/docker.sock", dockerSocketPath),
 			logsDir + ":/tmp/control",
 		},
 		Mounts: []mount.Mount{
@@ -207,6 +232,16 @@ func generateRunID() string {
 	return timestamp
 }
 
+// generateRunIDFromContainerName extracts the run ID from container name.
+func generateRunIDFromContainerName(containerName string) string {
+	// Extract run ID from container name like "headscale-test-suite-20250618-143802"
+	parts := strings.Split(containerName, "-")
+	if len(parts) >= 2 {
+		return strings.Join(parts[len(parts)-2:], "-")
+	}
+	return containerName
+}
+
 // findProjectRoot locates the project root by finding the directory containing go.mod.
 func findProjectRoot(startPath string) string {
 	current := startPath
@@ -288,6 +323,13 @@ func getCurrentDockerContext() (*DockerContext, error) {
 	return nil, ErrNoDockerContext
 }
 
+// getDockerSocketPath returns the correct Docker socket path for the current context.
+func getDockerSocketPath() string {
+	// Always use the default socket path for mounting since Docker handles
+	// the translation to the actual socket (e.g., colima socket) internally
+	return "/var/run/docker.sock"
+}
+
 // ensureImageAvailable pulls the specified Docker image to ensure it's available.
 func ensureImageAvailable(ctx context.Context, cli *client.Client, imageName string, verbose bool) error {
 	if verbose {
@@ -341,7 +383,7 @@ func listControlFiles(logsDir string) {
 		switch {
 		case strings.HasSuffix(name, ".stderr.log") || strings.HasSuffix(name, ".stdout.log"):
 			logFiles = append(logFiles, name)
-		case strings.HasSuffix(name, ".pprof.tar") || strings.HasSuffix(name, ".maps.tar") || strings.HasSuffix(name, ".db.tar"):
+		case strings.HasSuffix(name, ".db") || strings.HasSuffix(name, ".pprof") || strings.HasSuffix(name, ".mapresp"):
 			tarFiles = append(tarFiles, name)
 		}
 	}
@@ -356,9 +398,293 @@ func listControlFiles(logsDir string) {
 	}
 
 	if len(tarFiles) > 0 {
-		log.Printf("Headscale archives:")
+		log.Printf("Headscale files:")
 		for _, file := range tarFiles {
 			log.Printf("  %s", file)
 		}
 	}
 }
+
+// extractArtifactsFromContainers collects container logs and files from the specific test run.
+func extractArtifactsFromContainers(ctx context.Context, testContainerID, logsDir string, verbose bool) error {
+	cli, err := createDockerClient()
+	if err != nil {
+		return fmt.Errorf("failed to create Docker client: %w", err)
+	}
+	defer cli.Close()
+
+	// List all containers
+	containers, err := cli.ContainerList(ctx, container.ListOptions{All: true})
+	if err != nil {
+		return fmt.Errorf("failed to list containers: %w", err)
+	}
+
+	// Get containers from the specific test run
+	currentTestContainers := getCurrentTestContainers(containers, testContainerID, verbose)
+
+	extractedCount := 0
+	for _, cont := range currentTestContainers {
+		// Extract container logs and tar files
+		if err := extractContainerArtifacts(ctx, cli, cont.ID, cont.name, logsDir, verbose); err != nil {
+			if verbose {
+				log.Printf("Warning: failed to extract artifacts from container %s (%s): %v", cont.name, cont.ID[:12], err)
+			}
+		} else {
+			if verbose {
+				log.Printf("Extracted artifacts from container %s (%s)", cont.name, cont.ID[:12])
+			}
+			extractedCount++
+		}
+	}
+
+	if verbose && extractedCount > 0 {
+		log.Printf("Extracted artifacts from %d containers", extractedCount)
+	}
+
+	return nil
+}
+
+// testContainer represents a container from the current test run.
+type testContainer struct {
+	ID   string
+	name string
+}
+
+// getCurrentTestContainers filters containers to only include those from the current test run.
+func getCurrentTestContainers(containers []container.Summary, testContainerID string, verbose bool) []testContainer {
+	var testRunContainers []testContainer
+	
+	// Find the test container to get its run ID label
+	var runID string
+	for _, cont := range containers {
+		if cont.ID == testContainerID {
+			if cont.Labels != nil {
+				runID = cont.Labels["hi.run-id"]
+			}
+			break
+		}
+	}
+	
+	if runID == "" {
+		if verbose {
+			log.Printf("Warning: could not find run ID for test container %s, falling back to time-based filtering", testContainerID[:12])
+		}
+		// Fallback to time-based filtering for backward compatibility
+		return getCurrentTestContainersByTime(containers, testContainerID, verbose)
+	}
+	
+	if verbose {
+		log.Printf("Looking for containers with run ID: %s", runID)
+	}
+	
+	// Find all containers with the same run ID
+	for _, cont := range containers {
+		for _, name := range cont.Names {
+			containerName := strings.TrimPrefix(name, "/")
+			if strings.HasPrefix(containerName, "hs-") || strings.HasPrefix(containerName, "ts-") {
+				// Check if container has matching run ID label
+				if cont.Labels != nil && cont.Labels["hi.run-id"] == runID {
+					testRunContainers = append(testRunContainers, testContainer{
+						ID:   cont.ID,
+						name: containerName,
+					})
+					if verbose {
+						log.Printf("Including container %s (run ID: %s)", containerName, runID)
+					}
+				}
+				break
+			}
+		}
+	}
+	
+	return testRunContainers
+}
+
+// extractContainerArtifacts saves logs and tar files from a container.
+func extractContainerArtifacts(ctx context.Context, cli *client.Client, containerID, containerName, logsDir string, verbose bool) error {
+	// Ensure the logs directory exists
+	if err := os.MkdirAll(logsDir, 0755); err != nil {
+		return fmt.Errorf("failed to create logs directory: %w", err)
+	}
+
+	// Extract container logs
+	if err := extractContainerLogs(ctx, cli, containerID, containerName, logsDir, verbose); err != nil {
+		return fmt.Errorf("failed to extract logs: %w", err)
+	}
+
+	// Extract tar files for headscale containers only
+	if strings.HasPrefix(containerName, "hs-") {
+		if err := extractContainerFiles(ctx, cli, containerID, containerName, logsDir, verbose); err != nil {
+			if verbose {
+				log.Printf("Warning: failed to extract files from %s: %v", containerName, err)
+			}
+			// Don't fail the whole extraction if files are missing
+		}
+	}
+
+	return nil
+}
+
+// extractContainerLogs saves the stdout and stderr logs from a container to files.
+func extractContainerLogs(ctx context.Context, cli *client.Client, containerID, containerName, logsDir string, verbose bool) error {
+	// Get container logs
+	logReader, err := cli.ContainerLogs(ctx, containerID, container.LogsOptions{
+		ShowStdout: true,
+		ShowStderr: true,
+		Timestamps: false,
+		Follow:     false,
+		Tail:       "all",
+	})
+	if err != nil {
+		return fmt.Errorf("failed to get container logs: %w", err)
+	}
+	defer logReader.Close()
+
+	// Create log files following the headscale naming convention
+	stdoutPath := filepath.Join(logsDir, containerName+".stdout.log")
+	stderrPath := filepath.Join(logsDir, containerName+".stderr.log")
+
+	// Create buffers to capture stdout and stderr separately
+	var stdoutBuf, stderrBuf bytes.Buffer
+
+	// Demultiplex the Docker logs stream to separate stdout and stderr
+	_, err = stdcopy.StdCopy(&stdoutBuf, &stderrBuf, logReader)
+	if err != nil {
+		return fmt.Errorf("failed to demultiplex container logs: %w", err)
+	}
+
+	// Write stdout logs
+	if err := os.WriteFile(stdoutPath, stdoutBuf.Bytes(), 0644); err != nil {
+		return fmt.Errorf("failed to write stdout log: %w", err)
+	}
+
+	// Write stderr logs
+	if err := os.WriteFile(stderrPath, stderrBuf.Bytes(), 0644); err != nil {
+		return fmt.Errorf("failed to write stderr log: %w", err)
+	}
+
+	if verbose {
+		log.Printf("Saved logs for %s: %s, %s", containerName, stdoutPath, stderrPath)
+	}
+
+	return nil
+}
+
+// getCurrentTestContainersByTime is a fallback method for containers without labels.
+func getCurrentTestContainersByTime(containers []container.Summary, testContainerID string, verbose bool) []testContainer {
+	var testRunContainers []testContainer
+	
+	// Find the test container to get its creation time
+	var testContainerCreated time.Time
+	for _, cont := range containers {
+		if cont.ID == testContainerID {
+			testContainerCreated = time.Unix(cont.Created, 0)
+			break
+		}
+	}
+	
+	if testContainerCreated.IsZero() {
+		if verbose {
+			log.Printf("Warning: could not find test container %s", testContainerID[:12])
+		}
+		return testRunContainers
+	}
+	
+	// Find containers created within a small time window after the test container
+	startTime := testContainerCreated
+	endTime := testContainerCreated.Add(5 * time.Minute)
+	
+	for _, cont := range containers {
+		for _, name := range cont.Names {
+			containerName := strings.TrimPrefix(name, "/")
+			if strings.HasPrefix(containerName, "hs-") || strings.HasPrefix(containerName, "ts-") {
+				createdTime := time.Unix(cont.Created, 0)
+				if createdTime.After(startTime) && createdTime.Before(endTime) {
+					testRunContainers = append(testRunContainers, testContainer{
+						ID:   cont.ID,
+						name: containerName,
+					})
+					if verbose {
+						log.Printf("Including container %s (created %s)", containerName, createdTime.Format("15:04:05"))
+					}
+				}
+				break
+			}
+		}
+	}
+	
+	return testRunContainers
+}
+
+// extractContainerFiles extracts database file and directories from headscale containers.
+func extractContainerFiles(ctx context.Context, cli *client.Client, containerID, containerName, logsDir string, verbose bool) error {
+	// Extract database file
+	if err := extractSingleFile(ctx, cli, containerID, "/tmp/integration_test_db.sqlite3", containerName+".db", logsDir, verbose); err != nil {
+		if verbose {
+			log.Printf("Warning: failed to extract database from %s: %v", containerName, err)
+		}
+	}
+
+	// Extract profile directory
+	if err := extractDirectory(ctx, cli, containerID, "/tmp/profile", containerName+".pprof", logsDir, verbose); err != nil {
+		if verbose {
+			log.Printf("Warning: failed to extract profile from %s: %v", containerName, err)
+		}
+	}
+
+	// Extract map responses directory
+	if err := extractDirectory(ctx, cli, containerID, "/tmp/mapresponses", containerName+".mapresp", logsDir, verbose); err != nil {
+		if verbose {
+			log.Printf("Warning: failed to extract mapresponses from %s: %v", containerName, err)
+		}
+	}
+
+	return nil
+}
+
+// extractSingleFile copies a single file from a container.
+func extractSingleFile(ctx context.Context, cli *client.Client, containerID, sourcePath, fileName, logsDir string, verbose bool) error {
+	tarReader, _, err := cli.CopyFromContainer(ctx, containerID, sourcePath)
+	if err != nil {
+		return fmt.Errorf("failed to copy %s from container: %w", sourcePath, err)
+	}
+	defer tarReader.Close()
+
+	// Extract the single file from the tar
+	filePath := filepath.Join(logsDir, fileName)
+	if err := extractFileFromTar(tarReader, filepath.Base(sourcePath), filePath); err != nil {
+		return fmt.Errorf("failed to extract file from tar: %w", err)
+	}
+
+	if verbose {
+		log.Printf("Extracted %s from %s", fileName, containerID[:12])
+	}
+
+	return nil
+}
+
+// extractDirectory copies a directory from a container and extracts its contents.
+func extractDirectory(ctx context.Context, cli *client.Client, containerID, sourcePath, dirName, logsDir string, verbose bool) error {
+	tarReader, _, err := cli.CopyFromContainer(ctx, containerID, sourcePath)
+	if err != nil {
+		return fmt.Errorf("failed to copy %s from container: %w", sourcePath, err)
+	}
+	defer tarReader.Close()
+
+	// Create target directory
+	targetDir := filepath.Join(logsDir, dirName)
+	if err := os.MkdirAll(targetDir, 0755); err != nil {
+		return fmt.Errorf("failed to create directory %s: %w", targetDir, err)
+	}
+
+	// Extract the directory from the tar
+	if err := extractDirectoryFromTar(tarReader, targetDir); err != nil {
+		return fmt.Errorf("failed to extract directory from tar: %w", err)
+	}
+
+	if verbose {
+		log.Printf("Extracted %s/ from %s", dirName, containerID[:12])
+	}
+
+	return nil
+}

cmd/hi/doctor.go

@@ -7,8 +7,6 @@ import (
 	"log"
 	"os/exec"
 	"strings"
-
-	"github.com/docker/docker/client"
 )
 
 var ErrSystemChecksFailed = errors.New("system checks failed")
@@ -88,7 +86,7 @@ func checkDockerBinary() DoctorResult {
 
 // checkDockerDaemon verifies Docker daemon is running and accessible.
 func checkDockerDaemon(ctx context.Context) DoctorResult {
-	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
+	cli, err := createDockerClient()
 	if err != nil {
 		return DoctorResult{
 			Name:    "Docker Daemon",

cmd/hi/tar_utils.go

@@ -0,0 +1,95 @@
+package main
+
+import (
+	"archive/tar"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+// extractFileFromTar extracts a single file from a tar reader.
+func extractFileFromTar(tarReader io.Reader, fileName, outputPath string) error {
+	tr := tar.NewReader(tarReader)
+
+	for {
+		header, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return fmt.Errorf("failed to read tar header: %w", err)
+		}
+
+		// Check if this is the file we're looking for
+		if filepath.Base(header.Name) == fileName {
+			if header.Typeflag == tar.TypeReg {
+				// Create the output file
+				outFile, err := os.Create(outputPath)
+				if err != nil {
+					return fmt.Errorf("failed to create output file: %w", err)
+				}
+				defer outFile.Close()
+
+				// Copy file contents
+				if _, err := io.Copy(outFile, tr); err != nil {
+					return fmt.Errorf("failed to copy file contents: %w", err)
+				}
+				return nil
+			}
+		}
+	}
+
+	return fmt.Errorf("file %s not found in tar", fileName)
+}
+
+// extractDirectoryFromTar extracts all files from a tar reader to a target directory.
+func extractDirectoryFromTar(tarReader io.Reader, targetDir string) error {
+	tr := tar.NewReader(tarReader)
+
+	for {
+		header, err := tr.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return fmt.Errorf("failed to read tar header: %w", err)
+		}
+
+		// Clean the path to prevent directory traversal
+		cleanName := filepath.Clean(header.Name)
+		if strings.Contains(cleanName, "..") {
+			continue // Skip potentially dangerous paths
+		}
+
+		targetPath := filepath.Join(targetDir, filepath.Base(cleanName))
+
+		switch header.Typeflag {
+		case tar.TypeDir:
+			// Create directory
+			if err := os.MkdirAll(targetPath, os.FileMode(header.Mode)); err != nil {
+				return fmt.Errorf("failed to create directory %s: %w", targetPath, err)
+			}
+		case tar.TypeReg:
+			// Create file
+			outFile, err := os.Create(targetPath)
+			if err != nil {
+				return fmt.Errorf("failed to create file %s: %w", targetPath, err)
+			}
+
+			if _, err := io.Copy(outFile, tr); err != nil {
+				outFile.Close()
+				return fmt.Errorf("failed to copy file contents: %w", err)
+			}
+			outFile.Close()
+
+			// Set file permissions
+			if err := os.Chmod(targetPath, os.FileMode(header.Mode)); err != nil {
+				return fmt.Errorf("failed to set file permissions: %w", err)
+			}
+		}
+	}
+
+	return nil
+}

integration/dockertestutil/config.go

@@ -3,42 +3,37 @@ package dockertestutil
 import (
 	"os"
 
-	"github.com/ory/dockertest/v3/docker"
+	"github.com/ory/dockertest/v3"
 )
 
+// GetIntegrationRunID returns the run ID for the current integration test session.
+// This is set by the hi tool and passed through environment variables.
+func GetIntegrationRunID() string {
+	return os.Getenv("HEADSCALE_INTEGRATION_RUN_ID")
+}
+
+// DockerAddIntegrationLabels adds integration test labels to Docker RunOptions.
+// This allows the hi tool to identify containers belonging to specific test runs.
+// This function should be called before passing RunOptions to dockertest functions.
+func DockerAddIntegrationLabels(opts *dockertest.RunOptions, testType string) {
+	runID := GetIntegrationRunID()
+	if runID == "" {
+		// If no run ID is set, do nothing for backward compatibility
+		return
+	}
+
+	if opts.Labels == nil {
+		opts.Labels = make(map[string]string)
+	}
+	opts.Labels["hi.run-id"] = runID
+	opts.Labels["hi.test-type"] = testType
+}
+
+// IsRunningInContainer checks if the current process is running inside a Docker container.
+// This is used by tests to determine if they should run integration tests.
 func IsRunningInContainer() bool {
-	if _, err := os.Stat("/.dockerenv"); err != nil {
+	// Check for the common indicator that we're in a container
-		return false
+	// This could be improved with more robust detection if needed
-	}
+	_, err := os.Stat("/.dockerenv")
-
+	return err == nil
-	return true
+}
-}
-
-func DockerRestartPolicy(config *docker.HostConfig) {
-	// set AutoRemove to true so that stopped container goes away by itself on error *immediately*.
-	// when set to false, containers remain until the end of the integration test.
-	config.AutoRemove = false
-	config.RestartPolicy = docker.RestartPolicy{
-		Name: "no",
-	}
-}
-
-func DockerAllowLocalIPv6(config *docker.HostConfig) {
-	if config.Sysctls == nil {
-		config.Sysctls = make(map[string]string, 1)
-	}
-	config.Sysctls["net.ipv6.conf.all.disable_ipv6"] = "0"
-}
-
-func DockerAllowNetworkAdministration(config *docker.HostConfig) {
-	// Needed since containerd (1.7.24)
-	// https://github.com/tailscale/tailscale/issues/14256
-	// https://github.com/opencontainers/runc/commit/2ce40b6ad72b4bd4391380cafc5ef1bad1fa0b31
-	config.CapAdd = append(config.CapAdd, "NET_ADMIN")
-	config.CapAdd = append(config.CapAdd, "NET_RAW")
-	config.Devices = append(config.Devices, docker.Device{
-		PathOnHost:        "/dev/net/tun",
-		PathInContainer:   "/dev/net/tun",
-		CgroupPermissions: "rwm",
-	})
-}

integration/dockertestutil/network.go

@@ -126,3 +126,24 @@ func CleanImagesInCI(pool *dockertest.Pool) error {
 
 	return nil
 }
+
+// DockerRestartPolicy sets the restart policy for containers.
+func DockerRestartPolicy(config *docker.HostConfig) {
+	config.RestartPolicy = docker.RestartPolicy{
+		Name: "unless-stopped",
+	}
+}
+
+// DockerAllowLocalIPv6 allows IPv6 traffic within the container.
+func DockerAllowLocalIPv6(config *docker.HostConfig) {
+	config.NetworkMode = "default"
+	config.Sysctls = map[string]string{
+		"net.ipv6.conf.all.disable_ipv6": "0",
+	}
+}
+
+// DockerAllowNetworkAdministration gives the container network administration capabilities.
+func DockerAllowNetworkAdministration(config *docker.HostConfig) {
+	config.CapAdd = append(config.CapAdd, "NET_ADMIN")
+	config.Privileged = true
+}

integration/dsic/dsic.go

@@ -159,6 +159,7 @@ func New(
 		},
 	}
 
+
 	if dsic.workdir != "" {
 		runOptions.WorkingDir = dsic.workdir
 	}
@@ -189,6 +190,9 @@ func New(
 			Value: "v" + version,
 		})
 	}
+	// Add integration test labels if running under hi tool
+	dockertestutil.DockerAddIntegrationLabels(runOptions, "derp")
+	
 	container, err = pool.BuildAndRunWithBuildOptions(
 		buildOptions,
 		runOptions,

integration/hsic/hsic.go

@@ -311,18 +311,22 @@ func New(
 		hsic.env["HEADSCALE_DATABASE_POSTGRES_NAME"] = "headscale"
 		delete(hsic.env, "HEADSCALE_DATABASE_SQLITE_PATH")
 
-		pg, err := pool.RunWithOptions(
+		pgRunOptions := &dockertest.RunOptions{
-			&dockertest.RunOptions{
+			Name:       fmt.Sprintf("postgres-%s", hash),
-				Name:       fmt.Sprintf("postgres-%s", hash),
+			Repository: "postgres",
-				Repository: "postgres",
+			Tag:        "latest",
-				Tag:        "latest",
+			Networks:   networks,
-				Networks:   networks,
+			Env: []string{
-				Env: []string{
+				"POSTGRES_USER=headscale",
-					"POSTGRES_USER=headscale",
+				"POSTGRES_PASSWORD=headscale",
-					"POSTGRES_PASSWORD=headscale",
+				"POSTGRES_DB=headscale",
-					"POSTGRES_DB=headscale",
+			},
-				},
+		}
-			})
+
+		// Add integration test labels if running under hi tool
+		dockertestutil.DockerAddIntegrationLabels(pgRunOptions, "postgres")
+		
+		pg, err := pool.RunWithOptions(pgRunOptions)
 		if err != nil {
 			return nil, fmt.Errorf("starting postgres container: %w", err)
 		}
@@ -366,6 +370,7 @@ func New(
 		Env:        env,
 	}
 
+
 	if len(hsic.hostPortBindings) > 0 {
 		runOptions.PortBindings = map[docker.Port][]docker.PortBinding{}
 		for port, hostPorts := range hsic.hostPortBindings {
@@ -386,6 +391,9 @@ func New(
 		return nil, err
 	}
 
+	// Add integration test labels if running under hi tool
+	dockertestutil.DockerAddIntegrationLabels(runOptions, "headscale")
+	
 	container, err := pool.BuildAndRunWithBuildOptions(
 		headscaleBuildOptions,
 		runOptions,

integration/scenario.go

@@ -1102,6 +1102,7 @@ func (s *Scenario) runMockOIDC(accessTTL time.Duration, users []mockoidc.MockUse
 		},
 	}
 
+
 	headscaleBuildOptions := &dockertest.BuildOptions{
 		Dockerfile: hsic.IntegrationTestDockerFileName,
 		ContextDir: dockerContextPath,
@@ -1114,6 +1115,9 @@ func (s *Scenario) runMockOIDC(accessTTL time.Duration, users []mockoidc.MockUse
 
 	s.mockOIDC = scenarioOIDC{}
 
+	// Add integration test labels if running under hi tool
+	dockertestutil.DockerAddIntegrationLabels(mockOidcOptions, "oidc")
+	
 	if pmockoidc, err := s.pool.BuildAndRunWithBuildOptions(
 		headscaleBuildOptions,
 		mockOidcOptions,
@@ -1198,6 +1202,9 @@ func Webservice(s *Scenario, networkName string) (*dockertest.Resource, error) {
 		Env:      []string{},
 	}
 
+	// Add integration test labels if running under hi tool
+	dockertestutil.DockerAddIntegrationLabels(webOpts, "web")
+
 	webBOpts := &dockertest.BuildOptions{
 		Dockerfile: hsic.IntegrationTestDockerFileName,
 		ContextDir: dockerContextPath,

integration/tsic/tsic.go

@@ -251,6 +251,7 @@ func New(
 		Env:        []string{},
 	}
 
+
 	if tsic.withWebsocketDERP {
 		if version != VersionHead {
 			return tsic, errInvalidClientConfig
@@ -310,6 +311,9 @@ func New(
 			)
 		}
 
+		// Add integration test labels if running under hi tool
+		dockertestutil.DockerAddIntegrationLabels(tailscaleOptions, "tailscale")
+		
 		container, err = pool.BuildAndRunWithBuildOptions(
 			buildOptions,
 			tailscaleOptions,
@@ -321,6 +325,9 @@ func New(
 		tailscaleOptions.Repository = "tailscale/tailscale"
 		tailscaleOptions.Tag = version
 
+		// Add integration test labels if running under hi tool
+		dockertestutil.DockerAddIntegrationLabels(tailscaleOptions, "tailscale")
+		
 		container, err = pool.RunWithOptions(
 			tailscaleOptions,
 			dockertestutil.DockerRestartPolicy,
@@ -331,6 +338,9 @@ func New(
 		tailscaleOptions.Repository = "tailscale/tailscale"
 		tailscaleOptions.Tag = "v" + version
 
+		// Add integration test labels if running under hi tool
+		dockertestutil.DockerAddIntegrationLabels(tailscaleOptions, "tailscale")
+		
 		container, err = pool.RunWithOptions(
 			tailscaleOptions,
 			dockertestutil.DockerRestartPolicy,