Redesign unit tests

go-tests.sh

@@ -1,6 +1,5 @@
 #!/bin/bash
 
-
 # Sliver Implant Framework
 # Copyright (C) 2019  Bishop Fox
 
@@ -17,6 +16,8 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+set -euo pipefail
+
 SKIP_GENERATE=0
 for arg in "$@"; do
     if [ "$arg" = "--skip-generate" ]; then
@@ -28,231 +29,59 @@ echo "----------------------------------------------------------------"
 echo "WARNING: Running unit tests on slow systems can take a LONG time"
 echo "         Recommended to only run on 16+ CPU cores and 32Gb+ RAM"
 echo "----------------------------------------------------------------"
-TAGS=osusergo,netgo,go_sqlite
 
-## Client
+BASE_TAGS="osusergo,netgo,go_sqlite"
+ALL_TAGS="server,client,${BASE_TAGS}"
 
-# client / command / alias
-if go test -tags=client,$TAGS ./client/command/alias ; then
-    :
-else
+SLIVER_ROOT_DIR="$(mktemp -d)"
+SLIVER_CLIENT_ROOT_DIR="$(mktemp -d)"
+export SLIVER_ROOT_DIR
+export SLIVER_CLIENT_ROOT_DIR
+
+cleanup() {
+    rm -rf "${SLIVER_ROOT_DIR:-}" "${SLIVER_CLIENT_ROOT_DIR:-}" "${SLIVER_ROOT_DIR_E2E:-}" "${SLIVER_CLIENT_ROOT_DIR_E2E:-}"
+}
+trap cleanup EXIT
+
+TEST_PKGS=()
+while IFS= read -r pkg; do
+    TEST_PKGS+=("$pkg")
+done < <(
+    go list -e -tags="$ALL_TAGS" -f '{{if not .Error}}{{.ImportPath}}{{end}}' ./... \
+        | rg -v '^github.com/bishopfox/sliver/implant/sliver$' \
+        | rg -v '^github.com/bishopfox/sliver/implant/sliver/runner$' \
+        | rg -v '^github.com/bishopfox/sliver/implant/sliver/proxy$' \
+        | rg -v '^github.com/bishopfox/sliver/implant/sliver/transports/httpclient/drivers/win/wininet$' \
+        | rg -v '^github.com/bishopfox/sliver/client/command/websites$' \
+        | rg -v '^github.com/bishopfox/sliver/server/generate$'
+)
+
+if [ "${#TEST_PKGS[@]}" -eq 0 ]; then
+    echo "No test packages selected"
     exit 1
 fi
 
-# client / command / extensions
-if go test -tags=client,$TAGS ./client/command/extensions ; then
-    :
-else
-    exit 1
-fi
+go test -vet=off -count=1 -timeout 45m -tags="$ALL_TAGS" "${TEST_PKGS[@]}"
 
-# client / command / generate
-if go test -tags=client,$TAGS ./client/command/generate ; then
-    :
-else
-    exit 1
-fi
-
-# client / credentials
-if go test -tags=client,$TAGS ./client/credentials ; then
-    :
-else
-    exit 1
-fi
-
-## Util
-
-# util 
-if go test -tags=server,$TAGS ./util ; then
-    :
-else
-    exit 1
-fi
-
-# util / encoders
-if go test -tags=server,$TAGS ./util/encoders/basex ; then
-    :
-else
-    exit 1
-fi
-if go test -tags=server,$TAGS ./util/encoders ; then
-    :
-else
-    exit 1
-fi
-
-# util / minisign
-if go test -tags=server,$TAGS ./util/minisign ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-## Implant
-
-# implant / sliver / extension
-if go test ./implant/sliver/extension ; then
-    :
-else
-    exit 1
-fi
-
-# implant / sliver / transports / dnsclient
-if go test ./implant/sliver/transports/dnsclient ; then
-    :
-else
-    exit 1
-fi
-
-# implant / sliver / transports / wireguard
-if go test ./implant/sliver/transports/wireguard ; then
-    :
-else
-    exit 1
-fi
-
-## Server
-
-# server / assets / traffic encoders
-if go test -timeout 10m -tags=server,$TAGS ./server/assets/traffic-encoders ; then
-    :
-else
-    exit 1
-fi
-
-# server / encoders
-if go test -timeout 10m -tags=server,$TAGS ./server/encoders ; then
-    :
-else
-    exit 1
-fi
-
-# server / website
-if go test -tags=server,$TAGS ./server/website ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / loot
-if go test -tags=server,$TAGS ./server/loot ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / certs
-if go test -tags=server,$TAGS ./server/certs ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / cryptography
-if go test -tags=server,$TAGS ./server/cryptography ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / encoders / shellcode / sgn
-if go test -tags=server,$TAGS ./server/encoders/shellcode/sgn ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / encoders / shellcode / arm64
-if go test -tags=server,$TAGS ./server/encoders/shellcode/arm64 ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / gogo / goname
-if go test -tags=server,$TAGS ./server/gogo/goname ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / gogo
-if go test -tags=server,$TAGS ./server/gogo ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / c2
-if go test -tags=server,$TAGS ./server/c2 ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / c2 / e2e (mtls + wg + yamux)
-SLIVER_ROOT_DIR_E2E="$(mktemp -d)"
-if SLIVER_ROOT_DIR="$SLIVER_ROOT_DIR_E2E" go test -tags=server,$TAGS,sliver_e2e ./server/c2 -run 'Test(MTLS|WG)Yamux_' -count=1 ; then
-    :
-else
-    cat "$SLIVER_ROOT_DIR_E2E/logs/sliver.log" 2>/dev/null || true
-    rm -rf "$SLIVER_ROOT_DIR_E2E"
-    exit 1
-fi
-
-# server / c2 / e2e (dns)
-if SLIVER_ROOT_DIR="$SLIVER_ROOT_DIR_E2E" go test -tags=server,$TAGS,sliver_e2e ./server/c2 -run 'TestDNS_' -count=1 ; then
-    :
-else
-    cat "$SLIVER_ROOT_DIR_E2E/logs/sliver.log" 2>/dev/null || true
-    rm -rf "$SLIVER_ROOT_DIR_E2E"
-    exit 1
-fi
-rm -rf "$SLIVER_ROOT_DIR_E2E"
-
-# server / configs
-if go test -tags=server,$TAGS ./server/configs ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / console
-if go test -tags=server,$TAGS ./server/console ; then
-    :
-else
-    cat ~/.sliver/logs/sliver.log
-    exit 1
-fi
-
-# server / generate
 if [ "$SKIP_GENERATE" -eq 0 ]; then
-    export GOPROXY=off
-    if go test -tags=server,$TAGS ./server/generate -timeout 6h ; then
-        :
-    else
-        cat ~/.sliver/logs/sliver.log
-        exit 1
-    fi
+    GOPROXY=off go test -vet=off -count=1 -timeout 6h -tags="server,${BASE_TAGS}" ./server/generate
 else
     echo "Skipping ./server/generate tests (--skip-generate)"
 fi
 
-# server / rpc / spoof metadata
-if go test -vet=off -tags=server,$TAGS ./server/rpc -run '^TestGenerateSpoofMetadataAppliesPETimestampOverBufnet$' -count=1 -timeout 30m ; then
+SLIVER_ROOT_DIR_E2E="$(mktemp -d)"
+SLIVER_CLIENT_ROOT_DIR_E2E="$(mktemp -d)"
+
+if SLIVER_ROOT_DIR="$SLIVER_ROOT_DIR_E2E" SLIVER_CLIENT_ROOT_DIR="$SLIVER_CLIENT_ROOT_DIR_E2E" go test -vet=off -tags="server,${BASE_TAGS},sliver_e2e" ./server/c2 -run 'Test(MTLS|WG)Yamux_' -count=1 -timeout 30m ; then
     :
 else
-    cat ~/.sliver/logs/sliver.log 2>/dev/null || true
+    cat "$SLIVER_ROOT_DIR_E2E/logs/sliver.log" 2>/dev/null || true
+    exit 1
+fi
+
+if SLIVER_ROOT_DIR="$SLIVER_ROOT_DIR_E2E" SLIVER_CLIENT_ROOT_DIR="$SLIVER_CLIENT_ROOT_DIR_E2E" go test -vet=off -tags="server,${BASE_TAGS},sliver_e2e" ./server/c2 -run 'TestDNS_' -count=1 -timeout 30m ; then
+    :
+else
+    cat "$SLIVER_ROOT_DIR_E2E/logs/sliver.log" 2>/dev/null || true
     exit 1
 fi

server/c2/c2_test.go

@@ -21,6 +21,7 @@ package c2
 import (
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
 	"os"
 	"testing"
 
@@ -29,6 +30,7 @@ import (
 	"github.com/bishopfox/sliver/server/cryptography"
 	"github.com/bishopfox/sliver/server/db"
 	"github.com/bishopfox/sliver/server/db/models"
+	"github.com/bishopfox/sliver/server/testutils"
 )
 
 var (
@@ -36,10 +38,25 @@ var (
 )
 
 func TestMain(m *testing.M) {
+	envState, err := testutils.SetupTestEnv(testutils.EnvOptions{
+		ForceFreshServerRoot: true,
+		ForceFreshClientRoot: true,
+	})
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "setup test env: %v\n", err)
+		os.Exit(1)
+	}
+
 	implantBuild := setup()
-	code1 := m.Run()
+	exitCode := m.Run()
 	cleanup(implantBuild)
-	os.Exit(code1)
+	if cleanupErr := envState.Cleanup(); cleanupErr != nil {
+		fmt.Fprintf(os.Stderr, "%v\n", cleanupErr)
+		if exitCode == 0 {
+			exitCode = 1
+		}
+	}
+	os.Exit(exitCode)
 }
 
 func setup() *models.ImplantBuild {

server/cryptography/cryptography_test.go

@@ -21,12 +21,14 @@ package cryptography
 import (
 	"bytes"
 	"crypto/rand"
+	"fmt"
 	insecureRand "math/rand"
 	"os"
 	"sync"
 	"testing"
 
 	implantCrypto "github.com/bishopfox/sliver/implant/sliver/cryptography"
+	"github.com/bishopfox/sliver/server/testutils"
 	"github.com/bishopfox/sliver/util/minisign"
 )
 
@@ -45,8 +47,24 @@ func randomData() []byte {
 }
 
 func TestMain(m *testing.M) {
+	envState, err := testutils.SetupTestEnv(testutils.EnvOptions{
+		ForceFreshServerRoot: true,
+		ForceFreshClientRoot: true,
+	})
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "setup test env: %v\n", err)
+		os.Exit(1)
+	}
+
 	setup()
-	os.Exit(m.Run())
+	exitCode := m.Run()
+	if cleanupErr := envState.Cleanup(); cleanupErr != nil {
+		fmt.Fprintf(os.Stderr, "%v\n", cleanupErr)
+		if exitCode == 0 {
+			exitCode = 1
+		}
+	}
+	os.Exit(exitCode)
 }
 
 func setup() {

server/generate/linkflags_test.go

@@ -81,7 +81,7 @@ func TestApplyZigStaticLinking_NotLinux(t *testing.T) {
 }
 
 func TestSanitizeZigForBuild_LinuxSharedObject(t *testing.T) {
-	t.Setenv("SLIVER_ROOT_DIR", "/tmp/sliver-test-root")
+	t.Setenv("SLIVER_ROOT_DIR", t.TempDir())
 	cfg := &gogo.GoConfig{
 		GOOS: "linux",
 		CGO:  "1",
@@ -108,7 +108,7 @@ func TestSanitizeZigForBuild_LinuxSharedObject(t *testing.T) {
 }
 
 func TestSanitizeZigForBuild_Windows(t *testing.T) {
-	t.Setenv("SLIVER_ROOT_DIR", "/tmp/sliver-test-root")
+	t.Setenv("SLIVER_ROOT_DIR", t.TempDir())
 	cfg := &gogo.GoConfig{
 		GOOS: "windows",
 		CGO:  "1",

server/generate/testmain_test.go

@@ -0,0 +1,30 @@
+package generate
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/bishopfox/sliver/server/testutils"
+)
+
+func TestMain(m *testing.M) {
+	envState, err := testutils.SetupTestEnv(testutils.EnvOptions{
+		ForceFreshServerRoot: true,
+		ForceFreshClientRoot: true,
+		UnpackServerAssets:   true,
+	})
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "setup test env: %v\n", err)
+		os.Exit(1)
+	}
+
+	exitCode := m.Run()
+	if cleanupErr := envState.Cleanup(); cleanupErr != nil {
+		fmt.Fprintf(os.Stderr, "%v\n", cleanupErr)
+		if exitCode == 0 {
+			exitCode = 1
+		}
+	}
+	os.Exit(exitCode)
+}

server/gogo/testmain_test.go

@@ -0,0 +1,30 @@
+package gogo
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/bishopfox/sliver/server/testutils"
+)
+
+func TestMain(m *testing.M) {
+	envState, err := testutils.SetupTestEnv(testutils.EnvOptions{
+		ForceFreshServerRoot: true,
+		ForceFreshClientRoot: true,
+		UnpackServerAssets:   true,
+	})
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "setup test env: %v\n", err)
+		os.Exit(1)
+	}
+
+	exitCode := m.Run()
+	if cleanupErr := envState.Cleanup(); cleanupErr != nil {
+		fmt.Fprintf(os.Stderr, "%v\n", cleanupErr)
+		if exitCode == 0 {
+			exitCode = 1
+		}
+	}
+	os.Exit(exitCode)
+}

server/rpc/testmain_test.go

@@ -0,0 +1,29 @@
+package rpc
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/bishopfox/sliver/server/testutils"
+)
+
+func TestMain(m *testing.M) {
+	envState, err := testutils.SetupTestEnv(testutils.EnvOptions{
+		ForceFreshServerRoot: true,
+		ForceFreshClientRoot: true,
+	})
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "setup test env: %v\n", err)
+		os.Exit(1)
+	}
+
+	exitCode := m.Run()
+	if cleanupErr := envState.Cleanup(); cleanupErr != nil {
+		fmt.Fprintf(os.Stderr, "%v\n", cleanupErr)
+		if exitCode == 0 {
+			exitCode = 1
+		}
+	}
+	os.Exit(exitCode)
+}

server/testutils/testenv.go

@@ -0,0 +1,101 @@
+//go:build server
+
+package testutils
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/bishopfox/sliver/server/assets"
+)
+
+const (
+	ServerRootEnv = "SLIVER_ROOT_DIR"
+	ClientRootEnv = "SLIVER_CLIENT_ROOT_DIR"
+)
+
+// EnvOptions configures how SetupTestEnv prepares test roots.
+type EnvOptions struct {
+	ForceFreshServerRoot bool
+	ForceFreshClientRoot bool
+	UnpackServerAssets   bool
+}
+
+// EnvState tracks test root directories and what must be cleaned up.
+type EnvState struct {
+	ServerRoot  string
+	ClientRoot  string
+	cleanupDirs []string
+}
+
+// SetupTestEnv ensures SLIVER_ROOT_DIR/SLIVER_CLIENT_ROOT_DIR are set to usable
+// directories for tests. It only removes directories it creates.
+func SetupTestEnv(opts EnvOptions) (*EnvState, error) {
+	state := &EnvState{}
+
+	serverRoot, serverCreated, err := ensureRootEnv(ServerRootEnv, "sliver-test-root-*", opts.ForceFreshServerRoot)
+	if err != nil {
+		return nil, err
+	}
+	state.ServerRoot = serverRoot
+	if serverCreated {
+		state.cleanupDirs = append(state.cleanupDirs, serverRoot)
+	}
+
+	clientRoot, clientCreated, err := ensureRootEnv(ClientRootEnv, "sliver-test-client-root-*", opts.ForceFreshClientRoot)
+	if err != nil {
+		_ = state.Cleanup()
+		return nil, err
+	}
+	state.ClientRoot = clientRoot
+	if clientCreated {
+		state.cleanupDirs = append(state.cleanupDirs, clientRoot)
+	}
+
+	if opts.UnpackServerAssets {
+		assets.Setup(true, false)
+	}
+
+	return state, nil
+}
+
+// Cleanup removes any temporary directories created by SetupTestEnv.
+func (s *EnvState) Cleanup() error {
+	if s == nil {
+		return nil
+	}
+
+	var errs []string
+	for i := len(s.cleanupDirs) - 1; i >= 0; i-- {
+		dir := s.cleanupDirs[i]
+		if err := os.RemoveAll(dir); err != nil {
+			errs = append(errs, fmt.Sprintf("%s: %v", dir, err))
+		}
+	}
+	if len(errs) == 0 {
+		return nil
+	}
+	return fmt.Errorf("test env cleanup failed: %s", strings.Join(errs, "; "))
+}
+
+func ensureRootEnv(envVar string, pattern string, forceFresh bool) (string, bool, error) {
+	if !forceFresh {
+		if current := os.Getenv(envVar); current != "" {
+			if err := os.MkdirAll(current, 0o700); err != nil {
+				return "", false, fmt.Errorf("create %s root %q: %w", envVar, current, err)
+			}
+			return current, false, nil
+		}
+	}
+
+	root, err := os.MkdirTemp("", pattern)
+	if err != nil {
+		return "", false, fmt.Errorf("create %s temp root: %w", envVar, err)
+	}
+	if err := os.Setenv(envVar, root); err != nil {
+		_ = os.RemoveAll(root)
+		return "", false, fmt.Errorf("set %s: %w", envVar, err)
+	}
+	return root, true, nil
+}