Simplify unit tests and strengthen deterministic coverage (#2433)

* Preserve DNS canaries in obfuscated implants

* Refactor canary regression assertions for lint

* Fix DNS canary detection and trigger persistence

.golangci.yml

@@ -9,6 +9,7 @@ run:
     - osusergo
     - netgo
     - go_sqlite
+    - sliver_generate
     # Raw implant entry points are Go templates; lint their supporting packages.
     - sliver_lint
 

client/command/ai/model_test.go

@@ -156,13 +156,6 @@ func TestMessageBlockLabelIgnoresUnknownOperatorPlaceholders(t *testing.T) {
 	}
 }
 
-func TestTranscriptSpeakerStyleUsesPrimaryThemeForUserMessages(t *testing.T) {
-	styles := transcriptSpeakerStyle("alice", "user")
-
-	assertColorEqual(t, styles.border.GetForeground(), clienttheme.Primary())
-	assertColorEqual(t, styles.label.GetBackground(), clienttheme.Primary())
-}
-
 func TestThinkingLevelOptionFocusedStyleUsesDarkerPrimaryShade(t *testing.T) {
 	model := newAIModel(nil, aiContext{}, nil)
 
@@ -544,19 +537,12 @@ func TestRenderHeaderUsesSingleCompactRow(t *testing.T) {
 	}
 }
 
-func TestTranscriptSpeakerPaletteVariesAcrossUsers(t *testing.T) {
-	seen := map[int]struct{}{}
-	for _, label := range []string{"alice", "bob", "charlie", "dana", "erin", "frank"} {
-		seen[transcriptSpeakerPaletteIndex(label, "user")] = struct{}{}
-	}
-	if len(seen) < 2 {
-		t.Fatalf("expected distinct users to map to more than one transcript color, got %d palette entries", len(seen))
-	}
-}
-
 func TestTranscriptSpeakerStyleUsesPrimaryBorderForUserMessages(t *testing.T) {
 	userStyles := transcriptSpeakerStyle("alice", "user")
-	userBorder := transcriptSpeakerStyle("alice", "user").border.Render("│")
+	assertColorEqual(t, userStyles.border.GetForeground(), clienttheme.Primary())
+	assertColorEqual(t, userStyles.label.GetBackground(), clienttheme.Primary())
+
+	userBorder := userStyles.border.Render("│")
 	wantUserBorder := lipgloss.NewStyle().Foreground(clienttheme.Primary()).Render("│")
 	if userBorder != wantUserBorder {
 		t.Fatalf("expected user message border to use theme primary color")

client/command/armory/armory_test.go

@@ -1 +0,0 @@
-package armory

client/command/exec/execute-shellcode_test.go

@@ -3,6 +3,8 @@ package exec
 import (
 	"testing"
 
+	"github.com/bishopfox/sliver/client/console"
+	consts "github.com/bishopfox/sliver/client/constants"
 	"github.com/spf13/cobra"
 )
 
@@ -87,14 +89,11 @@ func TestShouldConvertExecuteShellcodePE(t *testing.T) {
 func newExecuteShellcodeFlagsTestCommand(t *testing.T) *cobra.Command {
 	t.Helper()
 
-	cmd := &cobra.Command{Use: "execute-shellcode"}
-	cmd.Flags().Uint32("shellcode-entropy", 1, "")
-	cmd.Flags().Bool("shellcode-compress", false, "")
-	cmd.Flags().Uint32("shellcode-exitopt", 1, "")
-	cmd.Flags().Uint32("shellcode-bypass", 3, "")
-	cmd.Flags().Uint32("shellcode-headers", 1, "")
-	cmd.Flags().Bool("shellcode-thread", false, "")
-	cmd.Flags().Bool("shellcode-unicode", false, "")
-	cmd.Flags().Uint32("shellcode-oep", 0, "")
-	return cmd
+	for _, cmd := range Commands(&console.SliverClient{}) {
+		if cmd.Name() == consts.ExecuteShellcodeStr {
+			return cmd
+		}
+	}
+	t.Fatalf("missing %q command", consts.ExecuteShellcodeStr)
+	return nil
 }

client/command/generate/canaries_test.go

@@ -0,0 +1,54 @@
+package generate
+
+import (
+	"strconv"
+	"testing"
+
+	"github.com/bishopfox/sliver/client/command/flags"
+	"github.com/bishopfox/sliver/client/console"
+	consts "github.com/bishopfox/sliver/client/constants"
+	"github.com/spf13/cobra"
+)
+
+func TestCanariesCommandRegistered(t *testing.T) {
+	root := &cobra.Command{Use: "sliver"}
+	root.AddCommand(Commands(&console.SliverClient{})...)
+
+	cmd, _, err := root.Find([]string{consts.CanariesStr})
+	if err != nil {
+		t.Fatalf("find canaries command: %v", err)
+	}
+	if cmd.Parent() != root {
+		t.Fatal("canaries command is not registered at the root")
+	}
+	if cmd.Run == nil {
+		t.Fatal("canaries command has no handler")
+	}
+	if cmd.GroupID != consts.PayloadsHelpGroup {
+		t.Errorf("canaries group = %q, want %q", cmd.GroupID, consts.PayloadsHelpGroup)
+	}
+
+	for _, want := range []struct {
+		name         string
+		shorthand    string
+		valueType    string
+		defaultValue string
+	}{
+		{name: "burned", shorthand: "b", valueType: "bool", defaultValue: "false"},
+		{name: "timeout", shorthand: "t", valueType: "int64", defaultValue: strconv.Itoa(flags.DefaultTimeout)},
+	} {
+		flag := cmd.Flags().Lookup(want.name)
+		if flag == nil {
+			t.Fatalf("canaries command is missing --%s", want.name)
+		}
+		if flag.Shorthand != want.shorthand {
+			t.Errorf("--%s shorthand = %q, want %q", want.name, flag.Shorthand, want.shorthand)
+		}
+		if flag.Value.Type() != want.valueType {
+			t.Errorf("--%s type = %q, want %q", want.name, flag.Value.Type(), want.valueType)
+		}
+		if flag.DefValue != want.defaultValue {
+			t.Errorf("--%s default = %q, want %q", want.name, flag.DefValue, want.defaultValue)
+		}
+	}
+}

client/command/generate/commands.go

@@ -313,7 +313,7 @@ func Commands(con *console.SliverClient) []*cobra.Command {
 		f.Int64P("timeout", "t", flags.DefaultTimeout, "grpc timeout in seconds")
 	})
 
-	return []*cobra.Command{generateCmd, regenerateCmd, profilesCmd, implantBuildsCmd}
+	return []*cobra.Command{generateCmd, regenerateCmd, profilesCmd, implantBuildsCmd, canariesCmd}
 }
 
 // coreImplantFlags binds all flags common to all sliver implant types.

client/console/event_listener_test.go

@@ -3,57 +3,77 @@ package console
 import (
 	"sync"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/bishopfox/sliver/protobuf/clientpb"
 )
 
 func TestRemoveEventListenerUnblocksInFlightDispatchAndClosesReceiver(t *testing.T) {
-	con := &SliverClient{EventListeners: &sync.Map{}}
-	listenerID, events := con.CreateEventListener()
-	value, ok := con.EventListeners.Load(listenerID)
-	if !ok {
-		t.Fatal("created listener is not registered")
-	}
-	listener := value.(*consoleEventListener)
-	for index := 0; index < cap(listener.events); index++ {
-		listener.events <- &clientpb.Event{EventType: "buffered"}
-	}
-
-	dispatchDone := make(chan struct{})
-	go func() {
-		listener.send(&clientpb.Event{EventType: "in-flight"})
-		close(dispatchDone)
-	}()
-
-	con.RemoveEventListener(listenerID)
-	select {
-	case <-dispatchDone:
-	case <-time.After(time.Second):
-		t.Fatal("listener removal did not unblock in-flight dispatch")
-	}
-	if _, stillRegistered := con.EventListeners.Load(listenerID); stillRegistered {
-		t.Fatal("removed listener remains registered")
-	}
-
-	for event := range events {
-		_ = event
-	}
-}
-
-func TestConcurrentEventDispatchAndRemovalDoesNotPanic(_ *testing.T) {
-	for iteration := 0; iteration < 1_000; iteration++ {
+	synctest.Test(t, func(t *testing.T) {
 		con := &SliverClient{EventListeners: &sync.Map{}}
 		listenerID, events := con.CreateEventListener()
+		value, ok := con.EventListeners.Load(listenerID)
+		if !ok {
+			t.Fatal("created listener is not registered")
+		}
+		listener := value.(*consoleEventListener)
+		buffered := &clientpb.Event{EventType: "buffered"}
+		for index := 0; index < cap(listener.events); index++ {
+			listener.events <- buffered
+		}
+
 		dispatchDone := make(chan struct{})
 		go func() {
-			con.triggerEventListeners(&clientpb.Event{EventType: "race"})
+			event := &clientpb.Event{EventType: "in-flight"}
+			con.triggerEventListeners(event)
 			close(dispatchDone)
 		}()
-		con.RemoveEventListener(listenerID)
-		<-dispatchDone
-		for event := range events {
-			_ = event
+
+		// The full buffer and Wait establish that send is blocked before removal.
+		synctest.Wait()
+		select {
+		case <-dispatchDone:
+			t.Fatal("dispatch completed before listener removal")
+		default:
 		}
-	}
+
+		removeDone := make(chan struct{})
+		go func() {
+			con.RemoveEventListener(listenerID)
+			// Dispatch and repeated removal after closure must remain harmless.
+			con.triggerEventListeners(&clientpb.Event{EventType: "removed"})
+			listener.send(&clientpb.Event{EventType: "closed"})
+			con.RemoveEventListener(listenerID)
+			close(removeDone)
+		}()
+		for _, done := range []<-chan struct{}{dispatchDone, removeDone} {
+			select {
+			case <-done:
+			case <-time.After(time.Second):
+				t.Fatal("listener removal and in-flight dispatch did not complete")
+			}
+		}
+		if _, stillRegistered := con.EventListeners.Load(listenerID); stillRegistered {
+			t.Fatal("removed listener remains registered")
+		}
+
+		deadline := time.After(time.Second)
+		for received := 0; ; received++ {
+			select {
+			case event, open := <-events:
+				if !open {
+					if received != cap(listener.events) {
+						t.Fatalf("received %d buffered events, want %d", received, cap(listener.events))
+					}
+					return
+				}
+				if event != buffered {
+					t.Fatalf("unexpected event after removal: %v", event)
+				}
+			case <-deadline:
+				t.Fatal("removed listener receiver did not close")
+			}
+		}
+	})
 }

client/transport/mtls_test.go

@@ -16,113 +16,103 @@ func TestMultiplayerConnectModeDefaultsToAuto(t *testing.T) {
 	}
 }
 
-func TestSelectMultiplayerDialStrategyAutoWithoutWGConfigUsesDirectMTLS(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectAuto)
-
-	strategy, err := selectMultiplayerDialStrategy(&assets.ClientConfig{})
-	if err != nil {
-		t.Fatalf("select dial strategy: %v", err)
-	}
-	if strategy != multiplayerDialDirect {
-		t.Fatalf("expected direct mTLS strategy, got %v", strategy)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyAutoLegacyWGConfigUsesDirectMTLS(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectAuto)
-
-	strategy, err := selectMultiplayerDialStrategy(completeWGClientConfig(false))
-	if err != nil {
-		t.Fatalf("select dial strategy: %v", err)
-	}
-	if strategy != multiplayerDialDirect {
-		t.Fatalf("expected direct mTLS strategy, got %v", strategy)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyAutoEnabledWGConfigRequiresCompleteConfig(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectAuto)
-
-	_, err := selectMultiplayerDialStrategy(&assets.ClientConfig{
-		WG: &assets.ClientWGConfig{
-			Enabled:      true,
-			ServerPubKey: "server-pub",
-			ClientIP:     "100.64.0.2",
+func TestSelectMultiplayerDialStrategy(t *testing.T) {
+	for _, test := range []struct {
+		name        string
+		mode        MultiplayerConnectMode
+		config      *assets.ClientConfig
+		want        multiplayerDialStrategy
+		wantErr     error
+		wantErrText string
+	}{
+		{
+			name:   "auto without WireGuard config",
+			mode:   MultiplayerConnectAuto,
+			config: &assets.ClientConfig{},
+			want:   multiplayerDialDirect,
 		},
-	})
-	if !errors.Is(err, ErrIncompleteWireGuardConfig) {
-		t.Fatalf("expected incomplete WG config error, got %v", err)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyAutoEnabledWGConfigUsesWireGuard(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectAuto)
-
-	strategy, err := selectMultiplayerDialStrategy(completeWGClientConfig(true))
-	if err != nil {
-		t.Fatalf("select dial strategy: %v", err)
-	}
-	if strategy != multiplayerDialWireGuard {
-		t.Fatalf("expected WireGuard strategy, got %v", strategy)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyEnableWGRejectsMissingWGConfig(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectEnableWG)
-
-	_, err := selectMultiplayerDialStrategy(&assets.ClientConfig{})
-	if !errors.Is(err, ErrMissingWireGuardConfig) {
-		t.Fatalf("expected missing WG config error, got %v", err)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyEnableWGRejectsIncompleteWGConfig(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectEnableWG)
-
-	_, err := selectMultiplayerDialStrategy(&assets.ClientConfig{
-		WG: &assets.ClientWGConfig{
-			ServerPubKey: "server-pub",
-			ClientIP:     "100.64.0.2",
+		{
+			name:   "auto with legacy WireGuard config",
+			mode:   MultiplayerConnectAuto,
+			config: completeWGClientConfig(false),
+			want:   multiplayerDialDirect,
 		},
-	})
-	if !errors.Is(err, ErrIncompleteWireGuardConfig) {
-		t.Fatalf("expected incomplete WG config error, got %v", err)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyEnableWGUsesCompleteLegacyConfig(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectEnableWG)
-
-	strategy, err := selectMultiplayerDialStrategy(completeWGClientConfig(false))
-	if err != nil {
-		t.Fatalf("select dial strategy: %v", err)
-	}
-	if strategy != multiplayerDialWireGuard {
-		t.Fatalf("expected WireGuard strategy, got %v", strategy)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyDisableWGOverridesEnabledConfig(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectDisableWG)
-
-	strategy, err := selectMultiplayerDialStrategy(completeWGClientConfig(true))
-	if err != nil {
-		t.Fatalf("select dial strategy: %v", err)
-	}
-	if strategy != multiplayerDialDirect {
-		t.Fatalf("expected direct strategy, got %v", strategy)
-	}
-}
-
-func TestSelectMultiplayerDialStrategyRejectsInvalidMode(t *testing.T) {
-	setTestMultiplayerConnectMode(t, MultiplayerConnectMode(255))
-
-	strategy, err := selectMultiplayerDialStrategy(completeWGClientConfig(true))
-	if strategy != multiplayerDialDirect {
-		t.Fatalf("expected invalid mode to fall back to direct strategy, got %v", strategy)
-	}
-	if err == nil || err.Error() != "invalid multiplayer connect mode 255" {
-		t.Fatalf("expected invalid mode error, got %v", err)
+		{
+			name: "auto with incomplete enabled config",
+			mode: MultiplayerConnectAuto,
+			config: &assets.ClientConfig{WG: &assets.ClientWGConfig{
+				Enabled: true, ServerPubKey: "server-pub", ClientIP: "100.64.0.2",
+			}},
+			want:    multiplayerDialDirect,
+			wantErr: ErrIncompleteWireGuardConfig,
+		},
+		{
+			name:   "auto with enabled WireGuard config",
+			mode:   MultiplayerConnectAuto,
+			config: completeWGClientConfig(true),
+			want:   multiplayerDialWireGuard,
+		},
+		{
+			name:    "enable with missing WireGuard config",
+			mode:    MultiplayerConnectEnableWG,
+			config:  &assets.ClientConfig{},
+			want:    multiplayerDialDirect,
+			wantErr: ErrMissingWireGuardConfig,
+		},
+		{
+			name: "enable with incomplete WireGuard config",
+			mode: MultiplayerConnectEnableWG,
+			config: &assets.ClientConfig{WG: &assets.ClientWGConfig{
+				ServerPubKey: "server-pub", ClientIP: "100.64.0.2",
+			}},
+			want:    multiplayerDialDirect,
+			wantErr: ErrIncompleteWireGuardConfig,
+		},
+		{
+			name:   "enable with complete legacy config",
+			mode:   MultiplayerConnectEnableWG,
+			config: completeWGClientConfig(false),
+			want:   multiplayerDialWireGuard,
+		},
+		{
+			name:   "disable overrides enabled config",
+			mode:   MultiplayerConnectDisableWG,
+			config: completeWGClientConfig(true),
+			want:   multiplayerDialDirect,
+		},
+		{
+			name:        "invalid mode",
+			mode:        MultiplayerConnectMode(255),
+			config:      completeWGClientConfig(true),
+			want:        multiplayerDialDirect,
+			wantErrText: "invalid multiplayer connect mode 255",
+		},
+		{
+			name:        "nil config",
+			mode:        MultiplayerConnectAuto,
+			want:        multiplayerDialDirect,
+			wantErrText: "client config is required",
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			setTestMultiplayerConnectMode(t, test.mode)
+			strategy, err := selectMultiplayerDialStrategy(test.config)
+			if strategy != test.want {
+				t.Errorf("dial strategy = %v, want %v", strategy, test.want)
+			}
+			switch {
+			case test.wantErr != nil:
+				if !errors.Is(err, test.wantErr) {
+					t.Fatalf("select dial strategy error = %v, want %v", err, test.wantErr)
+				}
+			case test.wantErrText != "":
+				if err == nil || err.Error() != test.wantErrText {
+					t.Fatalf("select dial strategy error = %v, want %q", err, test.wantErrText)
+				}
+			case err != nil:
+				t.Fatalf("select dial strategy: %v", err)
+			}
+		})
 	}
 }
 

go-tests.ps1

@@ -280,7 +280,6 @@ try {
         $pkg = if ($testDir.StartsWith(".")) { $testDir } else { "./$testDir" }
         switch -Wildcard ($pkg) {
             "./server/c2" { continue }
-            "./server/generate" { continue }
             "./client/*" {
                 $clientTestPkgs += $pkg
                 continue
@@ -359,6 +358,7 @@ try {
         }
     }
 
+    # Includes generation helpers without the sliver_generate compilation tests.
     foreach ($pkg in $serverUtilTestPkgs) {
         if (Test-SkipPackage -Package $pkg -Tags "server,$tags") {
             continue
@@ -421,7 +421,7 @@ try {
 
         Invoke-TestStep -Name "./server/generate" -Action {
             Invoke-ExternalCommand -FilePath "go" `
-                -Arguments @("test", "-timeout", "6h", "-p", $generateGoP, "-parallel", $generateTestParallel, "-tags=server,$tags", "./server/generate") `
+                -Arguments @("test", "-timeout", "6h", "-p", $generateGoP, "-parallel", $generateTestParallel, "-tags=server,$tags,sliver_generate", "./server/generate") `
                 -Environment @{
                     GOMAXPROCS = $generateGoMaxProcs
                     GOPROXY    = "off"
@@ -430,9 +430,9 @@ try {
     } else {
         Write-Host ""
         if ($UnitOnly) {
-            Write-Host "Skipping ./server/generate tests (--unit-only)"
+            Write-Host "Skipping ./server/generate compilation tests (--unit-only)"
         } else {
-            Write-Host "Skipping ./server/generate tests (--skip-generate)"
+            Write-Host "Skipping ./server/generate compilation tests (--skip-generate)"
         }
     }
 } finally {

go-tests.sh

@@ -27,16 +27,70 @@ fi
 
 SKIP_GENERATE=0
 UNIT_ONLY=0
-for arg in "$@"; do
-	if [ "$arg" = "--skip-generate" ]; then
+GENERATE_ONLY=0
+GENERATE_RUN=""
+
+usage() {
+	cat <<'EOF'
+Usage: ./go-tests.sh [--skip-generate | --unit-only | --generate-only] [--run REGEX]
+
+  --skip-generate  Skip implant compilation tests.
+  --unit-only      Skip implant compilation and C2 end-to-end tests.
+  --generate-only  Run only the server/generate package tests.
+  --run REGEX      Select generation tests by Go test name regex; requires --generate-only.
+  --help          Show this help.
+EOF
+}
+
+while [ "$#" -gt 0 ]; do
+	case "$1" in
+	--skip-generate)
 		SKIP_GENERATE=1
-	fi
-	if [ "$arg" = "--unit-only" ]; then
+		;;
+	--unit-only)
 		UNIT_ONLY=1
 		SKIP_GENERATE=1
-	fi
+		;;
+	--generate-only)
+		GENERATE_ONLY=1
+		;;
+	--run)
+		if [ "$#" -lt 2 ] || [ -z "$2" ] || [[ "$2" == --* ]]; then
+			echo "--run requires a non-empty regex" >&2
+			exit 2
+		fi
+		GENERATE_RUN="$2"
+		shift
+		;;
+	--run=*)
+		GENERATE_RUN="${1#--run=}"
+		if [ -z "$GENERATE_RUN" ]; then
+			echo "--run requires a non-empty regex" >&2
+			exit 2
+		fi
+		;;
+	--help | -h)
+		usage
+		exit 0
+		;;
+	*)
+		echo "Unknown argument: $1" >&2
+		usage >&2
+		exit 2
+		;;
+	esac
+	shift
 done
 
+if [ "$GENERATE_ONLY" -eq 1 ] && [ "$SKIP_GENERATE" -eq 1 ]; then
+	echo "--generate-only cannot be combined with --skip-generate or --unit-only" >&2
+	exit 2
+fi
+if [ -n "$GENERATE_RUN" ] && [ "$GENERATE_ONLY" -eq 0 ]; then
+	echo "--run requires --generate-only" >&2
+	exit 2
+fi
+
 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"
@@ -175,8 +229,35 @@ unpack_server_assets() {
 	"$HOST_GO" run -tags=server,go_sqlite ./server unpack --force
 }
 
+run_generate_tests() {
+	# Reclaim the host build cache before generation so long CI runs do not
+	# exhaust the runner's disk while retaining the implant compiler cache.
+	run_test_cmd "clean Go build cache" "$HOST_GO" clean -cache || return 1
+	export GOPROXY=off
+	# Generate tests are memory intensive (garble/c-shared cross-builds). Keep
+	# parallelism conservative by default to avoid OOM kills on smaller systems.
+	local generate_gomaxprocs="${SLIVER_GENERATE_GOMAXPROCS:-2}"
+	local generate_go_p="${SLIVER_GENERATE_GO_P:-1}"
+	local generate_test_parallel="${SLIVER_GENERATE_TEST_PARALLEL:-1}"
+	local generate_test_cmd=(
+		"$HOST_GO" test -timeout 6h -p "$generate_go_p" -parallel "$generate_test_parallel" -tags="server,$TAGS,sliver_generate"
+	)
+	if [ -n "$GENERATE_RUN" ]; then
+		# Keep test selection out of GOFLAGS, which implant compiler children inherit.
+		generate_test_cmd+=(-run "$GENERATE_RUN" -count=1 -v)
+	fi
+	generate_test_cmd+=(./server/generate)
+	run_test_cmd "./server/generate" \
+		env GOMAXPROCS="$generate_gomaxprocs" "${generate_test_cmd[@]}"
+}
+
 run_test_cmd "unpack server assets" unpack_server_assets || exit 1
 
+if [ "$GENERATE_ONLY" -eq 1 ]; then
+	run_generate_tests || exit 1
+	exit 0
+fi
+
 TEST_DIRS=()
 while IFS= read -r test_dir; do
 	TEST_DIRS+=("$test_dir")
@@ -191,7 +272,7 @@ E2E_SUPPORT_TEST_PKGS=()
 for test_dir in "${TEST_DIRS[@]}"; do
 	pkg="./$test_dir"
 	case "$pkg" in
-	./server/c2 | ./server/generate)
+	./server/c2)
 		# handled separately below
 		;;
 	./client/*)
@@ -249,7 +330,7 @@ for pkg in "${IMPLANT_TEST_PKGS[@]}"; do
 	run_test_cmd "$pkg" "$HOST_GO" test "$pkg" || exit 1
 done
 
-## Server + Util
+## Server + Util, including generation helpers without compilation tests.
 for pkg in "${SERVER_UTIL_TEST_PKGS[@]}"; do
 	if should_skip_package "$pkg" "server,$TAGS"; then
 		continue
@@ -277,26 +358,14 @@ else
 	run_test_cmd "./server/c2 (e2e dns)" "$HOST_GO" test -tags="server,$TAGS,sliver_e2e" ./server/c2 -run 'TestDNS_' -count=1 || exit 1
 fi
 
-## Server generate
+## Implant compilation tests opt in with sliver_generate; helper tests run above.
 if [ "$SKIP_GENERATE" -eq 0 ]; then
-	# The package tests above no longer need their Go build cache. Reclaim it
-	# before the cross-platform generation matrix starts so long CI runs do not
-	# exhaust the runner's disk while retaining the implant compiler cache.
-	run_test_cmd "clean Go build cache" "$HOST_GO" clean -cache || exit 1
-	export GOPROXY=off
-	# Generate tests are memory intensive (garble/c-shared cross-builds). Keep
-	# parallelism conservative by default to avoid OOM kills on smaller systems.
-	GENERATE_GOMAXPROCS="${SLIVER_GENERATE_GOMAXPROCS:-2}"
-	GENERATE_GO_P="${SLIVER_GENERATE_GO_P:-1}"
-	GENERATE_TEST_PARALLEL="${SLIVER_GENERATE_TEST_PARALLEL:-1}"
-	run_test_cmd "./server/generate" \
-		env GOMAXPROCS="$GENERATE_GOMAXPROCS" \
-		"$HOST_GO" test -timeout 6h -p "$GENERATE_GO_P" -parallel "$GENERATE_TEST_PARALLEL" -tags="server,$TAGS" ./server/generate || exit 1
+	run_generate_tests || exit 1
 else
 	echo
 	if [ "$UNIT_ONLY" -eq 1 ]; then
-		echo "Skipping ./server/generate tests (--unit-only)"
+		echo "Skipping ./server/generate compilation tests (--unit-only)"
 	else
-		echo "Skipping ./server/generate tests (--skip-generate)"
+		echo "Skipping ./server/generate compilation tests (--skip-generate)"
 	fi
 fi

implant/sliver/constants/constants.go

@@ -25,10 +25,17 @@ var (
 	SliverName = `{{.Name}}`
 )
 
-// Message - Fake message for embedding canaries
+// Message carries canaries in struct tags, which remain plaintext under Garble.
 type Message struct {
 	Command string `c2:"[[GenerateCanary]]"`
+	// {{if .Config.IncludeMTLS}}
+	MTLS string `c2:"[[GenerateCanary]]"`
+	// {{end}}
+	// {{if .Config.IncludeWG}}
+	WireGuard string `c2:"[[GenerateCanary]]"`
+	// {{end}}
 }
 
-// never obfuscate the Message type
-var _ = reflect.TypeOf(Message{})
+// Access a field at runtime to retain the type metadata and its canary tags.
+// A discarded TypeOf call alone is optimized away, removing the canaries.
+var _ = reflect.TypeOf(Message{}).Field(0)

implant/sliver/extension/memfs_test.go

@@ -19,6 +19,8 @@ package extension
 */
 
 import (
+	"bytes"
+	"io"
 	"io/fs"
 	"testing"
 )
@@ -33,7 +35,11 @@ func TestBasicMemFSOpenFile(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if stat, _ := fi.Stat(); stat.Name() != "memfs/test.1" {
+	stat, err := fi.Stat()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stat.Name() != "memfs/test.1" {
 		t.Fatalf("expected memfs/test.1, got %s", stat.Name())
 	}
 	defer fi.Close()
@@ -58,20 +64,24 @@ func TestMemFSOpenRoot(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if len(entries) != 3 {
-		for entry := range entries {
-			t.Logf("entry: %s", entries[entry].Name())
+	wantEntries := map[string]bool{"memfs/test.1": false, "memfs/test.2": false, "memfs/a": true}
+	if len(entries) != len(wantEntries) {
+		t.Fatalf("expected %d entries, got %d", len(wantEntries), len(entries))
+	}
+	for _, entry := range entries {
+		wantDir, ok := wantEntries[entry.Name()]
+		if !ok || entry.IsDir() != wantDir {
+			t.Fatalf("unexpected entry %q (directory=%v)", entry.Name(), entry.IsDir())
 		}
-		t.Fatalf("expected 3 entries, got %d", len(entries))
+		delete(wantEntries, entry.Name())
 	}
 
 	di, err := dirNode.Stat()
 	if err != nil {
 		t.Fatal(err)
 	}
-	t.Logf("name: %s, is dir: %v", di.Name(), di.IsDir())
-	for _, entry := range entries {
-		t.Logf("entry: %s", entry.Name())
+	if di.Name() != "memfs" || !di.IsDir() {
+		t.Fatalf("expected memfs root directory, got %q (directory=%v)", di.Name(), di.IsDir())
 	}
 }
 
@@ -80,7 +90,7 @@ func TestBasicMemFSOpenDir(t *testing.T) {
 		"test/foo.1":             {0x00},
 		"test/foo":               {0x00},
 		"testing/foo":            {0x00},
-		"a/b/c/d/e/f/test/a.txt": {0x00},
+		"a/b/c/d/e/f/test/a.txt": {0x31, 0x00, 0x7f},
 	}}
 
 	// Test Open File
@@ -91,7 +101,11 @@ func TestBasicMemFSOpenDir(t *testing.T) {
 	if f == nil {
 		t.Fatal("expected file, got nil and no error")
 	}
-	if stat, _ := f.Stat(); stat.Name() != "memfs/test" || !stat.IsDir() {
+	stat, err := f.Stat()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stat.Name() != "memfs/test" || !stat.IsDir() {
 		t.Fatalf("expected 'memfs/test' dir, got '%s'", stat.Name())
 	}
 	defer f.Close()
@@ -100,7 +114,16 @@ func TestBasicMemFSOpenDir(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if stat, _ := dir.Stat(); stat.Name() != "memfs/a/b/c/d/e/f/test" || !stat.IsDir() {
+	defer func() {
+		if err := dir.Close(); err != nil {
+			t.Errorf("close directory: %v", err)
+		}
+	}()
+	stat, err = dir.Stat()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stat.Name() != "memfs/a/b/c/d/e/f/test" || !stat.IsDir() {
 		t.Fatalf("expected 'memfs/a/b/c/d/e/f/test' dir, got '%s'", stat.Name())
 	}
 
@@ -108,19 +131,17 @@ func TestBasicMemFSOpenDir(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	data := make([]byte, 1)
-	_, err = f.Read(data)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if data[0] != 0x00 {
-		t.Fatalf("expected 0x00, got %x", data[0])
-	}
+	defer func() {
+		if err := f.Close(); err != nil {
+			t.Errorf("close file: %v", err)
+		}
+	}()
+	assertMemFSFileContents(t, f, []byte{0x31, 0x00, 0x7f})
 }
 
 func TestMemFSReadFileData(t *testing.T) {
 	wasmFS := WasmMemoryFS{memFS: map[string][]byte{
-		"test/foo.1": {0x00},
+		"test/foo.1": {0x41, 0x00, 0xff, 0x42},
 	}}
 
 	// Test Open File
@@ -128,19 +149,31 @@ func TestMemFSReadFileData(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if stat, _ := fi.Stat(); stat.Name() != "memfs/test/foo.1" {
+	stat, err := fi.Stat()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if stat.Name() != "memfs/test/foo.1" {
 		t.Fatalf("expected memfs/test/foo.1, got %s", stat.Name())
 	}
 	defer fi.Close()
 
-	// Read file
-	data := make([]byte, 1)
-	_, err = fi.Read(data)
-	if err != nil {
-		t.Fatal(err)
+	assertMemFSFileContents(t, fi, []byte{0x41, 0x00, 0xff, 0x42})
+}
+
+func assertMemFSFileContents(t *testing.T, file fs.File, want []byte) {
+	t.Helper()
+	// Read in chunks to verify byte counts and advancement of the file offset.
+	for offset := 0; offset < len(want); {
+		end := min(offset+2, len(want))
+		data := make([]byte, end-offset)
+		n, err := file.Read(data)
+		if err != nil || n != len(data) || !bytes.Equal(data, want[offset:end]) {
+			t.Fatalf("Read at %d = (%x, %d, %v), want (%x, %d, nil)", offset, data, n, err, want[offset:end], len(data))
+		}
+		offset = end
 	}
-	if data[0] != 0x00 {
-		t.Fatalf("expected 0x00, got %x", data[0])
+	if n, err := file.Read(make([]byte, 1)); n != 0 || err != io.EOF {
+		t.Fatalf("Read after file contents = (%d, %v), want (0, EOF)", n, err)
 	}
-
 }

implant/sliver/transports/mtls/mtls.go

@@ -189,7 +189,7 @@ func ReadEnvelope(r io.Reader) (*pb.Envelope, error) {
 	rawSigBuf := make([]byte, cryptography.RawSigSize)
 	dataLengthBuf := make([]byte, 4) // Size of uint32
 	if len(rawSigBuf) == 0 || len(dataLengthBuf) == 0 || r == nil {
-		panic("[[GenerateCanary]]")
+		panic("invalid envelope reader")
 	}
 
 	n, err := io.ReadFull(r, rawSigBuf)

implant/sliver/transports/wireguard/wireguard.go

@@ -182,7 +182,7 @@ func ReadEnvelope(connection net.Conn) (*pb.Envelope, error) {
 	rawSigBuf := make([]byte, cryptography.RawSigSize)
 	dataLengthBuf := make([]byte, 4) // Size of uint32
 	if len(rawSigBuf) == 0 || len(dataLengthBuf) == 0 || connection == nil {
-		panic("[[GenerateCanary]]")
+		panic("invalid envelope connection")
 	}
 
 	n, err := io.ReadFull(connection, rawSigBuf)

server/assets/traffic-encoders/traffic-encoder_test.go

@@ -20,11 +20,10 @@ package traffic_encoders
 
 import (
 	"bytes"
-	"crypto/rand"
 	_ "embed"
 	"encoding/hex"
+	"fmt"
 	"testing"
-	"time"
 
 	implantEncoders "github.com/bishopfox/sliver/implant/sliver/encoders/traffic"
 	serverEncoders "github.com/bishopfox/sliver/util/encoders/traffic"
@@ -34,102 +33,124 @@ import (
 var hexWASM []byte
 
 func TestTrafficEncoderCompatibility_hex(t *testing.T) {
-
-	// Hex
-
 	implantSideHex, err := implantEncoders.CreateTrafficEncoder("hex", hexWASM, func(msg string) {
 		t.Log(msg)
 	})
 	if err != nil {
 		t.Fatal(err)
 	}
+	t.Cleanup(func() {
+		if err := implantSideHex.Close(); err != nil {
+			t.Errorf("close implant encoder: %v", err)
+		}
+	})
 	serverSideHex, err := serverEncoders.CreateTrafficEncoder("hex", hexWASM, func(msg string) {
 		t.Log(msg)
 	})
 	if err != nil {
 		t.Fatal(err)
 	}
-
-	data := make([]byte, 1024)
-	_, err = rand.Read(data)
-	if err != nil {
-		t.Fatal(err)
-	}
-	encodedData, err := implantSideHex.Encode(data)
-	if err != nil {
-		t.Fatal(err)
-	}
-	decodedData, err := serverSideHex.Decode(encodedData)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if !bytes.Equal(data, decodedData) {
-		t.Fatal("Decoded data does not match original")
-	}
-
-	data = make([]byte, 1024)
-	_, err = rand.Read(data)
-	if err != nil {
-		t.Fatal(err)
-	}
-	encodedData, err = serverSideHex.Encode(data)
-	if err != nil {
-		t.Fatal(err)
-	}
-	decodedData, err = implantSideHex.Decode(encodedData)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if !bytes.Equal(data, decodedData) {
-		t.Fatal("Decoded data does not match original")
-	}
-}
-
-// Encoder specific tests
-func TestHexPerformance(t *testing.T) {
-
-	sizes := []int{1024, 1024 * 1024, 2 * 1024 * 1024, 4 * 1024 * 1024}
-
-	// Stock encoder
-	for i := 0; i < len(sizes); i++ {
-		originalValue := make([]byte, sizes[i])
-		rand.Read(originalValue)
-		stock := time.Now()
-		encodedValue := hex.EncodeToString(originalValue)
-		decodedValue, err := hex.DecodeString(encodedValue)
-		if err != nil {
-			t.Fatal(err)
+	t.Cleanup(func() {
+		if err := serverSideHex.Close(); err != nil {
+			t.Errorf("close server encoder: %v", err)
 		}
-		t.Logf("Stock encoder took %v (%d bytes)", time.Since(stock), sizes[i])
-		if !bytes.Equal(originalValue, decodedValue) {
-			t.Fatalf("Expected %v but got %v", originalValue, decodedValue)
-		}
-	}
+	})
 
-	// Traffic encoder
-	for i := 0; i < len(sizes); i++ {
-		encoder, err := serverEncoders.CreateTrafficEncoder("hex", hexWASM, func(msg string) {
-			t.Log(msg)
+	directions := []struct {
+		name   string
+		encode func([]byte) ([]byte, error)
+		decode func([]byte) ([]byte, error)
+	}{
+		{"implant_to_server", implantSideHex.Encode, serverSideHex.Decode},
+		{"server_to_implant", serverSideHex.Encode, implantSideHex.Decode},
+	}
+	for _, direction := range directions {
+		t.Run(direction.name, func(t *testing.T) {
+			// Exercise empty and binary inputs, both sides of a 64 KiB WASM
+			// page boundary, and the largest payload from the old timing test.
+			for _, size := range []int{0, 1, 256, 65535, 65536, 65537, 4 * 1024 * 1024} {
+				t.Run(fmt.Sprintf("bytes_%d", size), func(t *testing.T) {
+					payload := hexTestPayload(size)
+					encoded, err := direction.encode(payload)
+					if err != nil {
+						t.Fatalf("encode: %v", err)
+					}
+					if !bytes.Equal(encoded, []byte(hex.EncodeToString(payload))) {
+						t.Fatal("WASM output does not match standard hex encoding")
+					}
+					decoded, err := direction.decode(encoded)
+					if err != nil {
+						t.Fatalf("decode: %v", err)
+					}
+					if !bytes.Equal(payload, decoded) {
+						t.Fatalf("round trip changed the %d-byte payload (decoded %d bytes)", len(payload), len(decoded))
+					}
+				})
+			}
 		})
-		if err != nil {
-			t.Fatal(err)
-		}
-		defer encoder.Close()
-		originalValue := make([]byte, sizes[i])
-		rand.Read(originalValue)
-		start := time.Now()
-		encodedValue, err := encoder.Encode(originalValue)
-		if err != nil {
-			t.Fatal(err)
-		}
-		// t.Logf("Got encoded value (%d bytes)", len(encodedValue))
-		decodedValue, err := encoder.Decode(encodedValue)
-		if err != nil {
-			t.Fatal(err)
-		}
-		t.Logf("WASM Hex encoder took %v (%d bytes)", time.Since(start), sizes[i])
-		if !bytes.Equal(originalValue, decodedValue) {
-			t.Fatalf("Expected %v but got %v", originalValue, decodedValue)
-		}
 	}
 }
+
+func BenchmarkHex(b *testing.B) {
+	for _, size := range []int{1024, 1024 * 1024, 2 * 1024 * 1024, 4 * 1024 * 1024} {
+		payload := hexTestPayload(size)
+		b.Run(fmt.Sprintf("bytes_%d", size), func(b *testing.B) {
+			b.Run("standard_round_trip", func(b *testing.B) {
+				b.SetBytes(int64(size))
+				b.ReportAllocs()
+				for b.Loop() {
+					encoded := hex.EncodeToString(payload)
+					decoded, err := hex.DecodeString(encoded)
+					if err != nil {
+						b.Fatal(err)
+					}
+					if !bytes.Equal(payload, decoded) {
+						b.Fatal("decoded data does not match original")
+					}
+				}
+			})
+			b.Run("wasm_fresh_runtime_round_trip", func(b *testing.B) {
+				b.SetBytes(int64(size))
+				b.ReportAllocs()
+				// Include runtime creation and close in each iteration. The WASM
+				// API retains output buffers until the runtime is closed, so an
+				// unbounded loop over one runtime would accumulate guest memory.
+				for b.Loop() {
+					benchmarkHexWASMLifecycle(b, payload)
+				}
+			})
+		})
+	}
+}
+
+func benchmarkHexWASMLifecycle(b *testing.B, payload []byte) {
+	b.Helper()
+	encoder, err := serverEncoders.CreateTrafficEncoder("hex", hexWASM, func(string) {})
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer func() {
+		if err := encoder.Close(); err != nil {
+			b.Errorf("close encoder: %v", err)
+		}
+	}()
+	encoded, err := encoder.Encode(payload)
+	if err != nil {
+		b.Fatal(err)
+	}
+	decoded, err := encoder.Decode(encoded)
+	if err != nil {
+		b.Fatal(err)
+	}
+	if !bytes.Equal(payload, decoded) {
+		b.Fatal("decoded data does not match original")
+	}
+}
+
+func hexTestPayload(size int) []byte {
+	payload := make([]byte, size)
+	for i := range payload {
+		payload[i] = byte(i)
+	}
+	return payload
+}

server/c2/dns.go

@@ -46,6 +46,7 @@ import (
 	"time"
 
 	consts "github.com/bishopfox/sliver/client/constants"
+	"github.com/bishopfox/sliver/protobuf/clientpb"
 	"github.com/bishopfox/sliver/protobuf/dnspb"
 	"github.com/bishopfox/sliver/protobuf/sliverpb"
 	"github.com/bishopfox/sliver/server/core"
@@ -576,15 +577,7 @@ func (s *SliverDNSServer) HandleDNSRequest(domains []string, canaries bool, writ
 		return
 	}
 
-	var resp *dns.Msg
-	isC2, domain := s.isC2SubDomain(domains, req.Question[0].Name)
-	if isC2 {
-		dnsLog.Debugf("'%s' is subdomain of c2 parent '%s'", req.Question[0].Name, domain)
-		resp = s.handleC2(domain, req)
-	} else if canaries {
-		dnsLog.Debugf("checking '%s' for DNS canary matches", req.Question[0].Name)
-		resp = s.handleCanary(req)
-	}
+	resp := s.handleDNSQuery(domains, canaries, req)
 	if resp != nil {
 		// These responses often have near-max-length QNAMEs. Enable compression to
 		// keep UDP responses under the minimum DNS message size limits.
@@ -595,6 +588,27 @@ func (s *SliverDNSServer) HandleDNSRequest(domains []string, canaries bool, writ
 	}
 }
 
+func (s *SliverDNSServer) handleDNSQuery(domains []string, canaries bool, req *dns.Msg) *dns.Msg {
+	if canaries {
+		canary, err := db.CanaryByDomain(req.Question[0].Name)
+		if err == nil && canary != nil {
+			return s.handleCanary(req, canary)
+		}
+		if err != nil && !errors.Is(err, db.ErrRecordNotFound) {
+			dnsLog.Errorf("Failed to look up DNS canary: %s", err)
+		}
+	}
+
+	if isC2, domain := s.isC2SubDomain(domains, req.Question[0].Name); isC2 {
+		dnsLog.Debugf("'%s' is subdomain of c2 parent '%s'", req.Question[0].Name, domain)
+		return s.handleC2(domain, req)
+	}
+	if canaries {
+		return s.nameErrorResp(req)
+	}
+	return nil
+}
+
 // Returns true if the requested domain is a c2 subdomain, and the domain it matched with
 func (s *SliverDNSServer) isC2SubDomain(domains []string, reqDomain string) (bool, string) {
 	for _, parentDomain := range domains {
@@ -1142,39 +1156,24 @@ func (s *SliverDNSServer) handleNOP(domain string, msg *dnspb.DNSMessage, checks
 // Canary Record Handler
 // ---------------------------
 // Canary -> valid? -> trigger alert event
-func (s *SliverDNSServer) handleCanary(req *dns.Msg) *dns.Msg {
-	// Don't block, return error as fast as possible
-	go func() {
-		defer recoverAndLogPanic(dnsLog.Errorf, "dns handleCanary")
+func (s *SliverDNSServer) handleCanary(req *dns.Msg, canary *clientpb.DNSCanary) *dns.Msg {
+	firstTrigger, err := generate.TriggerCanary(canary)
+	if err != nil {
+		dnsLog.Errorf("Failed to record DNS canary trigger for '%s': %s", canary.Domain, err)
+		return s.nameErrorResp(req)
+	}
 
-		reqDomain := strings.ToLower(req.Question[0].Name)
-		if !strings.HasSuffix(reqDomain, ".") {
-			reqDomain += "." // Ensure we have the FQDN
-		}
-		canary, err := db.CanaryByDomain(reqDomain)
-		if err != nil {
-			dnsLog.Errorf("Failed to find canary: %s", err)
-			return
-		}
-		if canary != nil {
-			dnsLog.Warnf("DNS canary tripped for '%s'", canary.ImplantName)
-			if !canary.Triggered {
-				// Defer publishing the event until we're sure the db is sync'd
-				defer core.EventBroker.Publish(core.Event{
-					Session: &core.Session{
-						Name: canary.ImplantName,
-					},
-					Data:      []byte(canary.Domain),
-					EventType: consts.CanaryEvent,
-				})
-				canary.Triggered = true
-				canary.FirstTriggered = time.Now().Format(time.RFC1123)
-			}
-			canary.LatestTrigger = time.Now().Format(time.RFC1123)
-			canary.Count++
-			generate.UpdateCanary(canary)
-		}
-	}()
+	dnsLog.Warnf("DNS canary tripped for '%s'", canary.ImplantName)
+	if firstTrigger {
+		// Publish only after the trigger transaction commits successfully.
+		core.EventBroker.Publish(core.Event{
+			Session: &core.Session{
+				Name: canary.ImplantName,
+			},
+			Data:      []byte(canary.Domain),
+			EventType: consts.CanaryEvent,
+		})
+	}
 	return s.nameErrorResp(req)
 }
 

server/c2/dns_canaries_test.go

@@ -0,0 +1,297 @@
+package c2
+
+import (
+	"encoding/binary"
+	"errors"
+	"hash/crc32"
+	"net"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	consts "github.com/bishopfox/sliver/client/constants"
+	"github.com/bishopfox/sliver/protobuf/dnspb"
+	"github.com/bishopfox/sliver/server/core"
+	"github.com/bishopfox/sliver/server/db"
+	"github.com/bishopfox/sliver/server/db/models"
+	"github.com/miekg/dns"
+	"google.golang.org/protobuf/proto"
+	"gorm.io/gorm"
+)
+
+func TestDNSCanarySameParentTriggersOnce(t *testing.T) {
+	const parent = "canary-parent.invalid."
+	// Legacy records can contain mixed-case parent names, and DNS queries are
+	// case-insensitive independently of how the registered name was spelled.
+	canary := createDNSCanaryRecord(t, models.NewUUID().String()+".CaNaRy-PaReNt.InVaLiD.")
+	events := subscribeDNSCanaryEvents(t)
+	lookups := countDNSCanaryLookups(t)
+	server := newTestDNSServer()
+
+	response := queryDNSCanary(t, server, parent, strings.ToUpper(canary.Domain), true)
+	assertDNSCanaryResponse(t, response)
+	if got := lookups.Load(); got != 1 {
+		t.Fatalf("first query performed %d canary lookups, want 1", got)
+	}
+	first := assertDNSCanaryRecord(t, canary, 1)
+	assertDNSCanaryEvents(t, collectDNSCanaryEvents(t, events, false), canary, 1)
+
+	lookups.Store(0)
+	response = queryDNSCanary(t, server, parent, strings.ToLower(canary.Domain), true)
+	assertDNSCanaryResponse(t, response)
+	if got := lookups.Load(); got != 1 {
+		t.Fatalf("repeat query performed %d canary lookups, want 1", got)
+	}
+	second := assertDNSCanaryRecord(t, canary, 2)
+	if !second.FirstTrigger.Equal(first.FirstTrigger) || second.LatestTrigger.Before(first.LatestTrigger) {
+		t.Error("repeat query changed the first trigger or moved the latest trigger backwards")
+	}
+	assertDNSCanaryEvents(t, collectDNSCanaryEvents(t, events, false), canary, 0)
+}
+
+func TestDNSCanaryC2Routing(t *testing.T) {
+	tests := []struct {
+		name       string
+		registered bool
+		enabled    bool
+	}{
+		{name: "unknown_name", enabled: true},
+		{name: "disabled_detection", registered: true},
+		{name: "registered_name_precedes_c2", registered: true, enabled: true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server, parent, domain, checksum := newDNSCanaryNOP(t)
+			var canary *models.DNSCanary
+			if tt.registered {
+				canary = createDNSCanaryRecord(t, domain)
+			}
+			events := subscribeDNSCanaryEvents(t)
+			lookups := countDNSCanaryLookups(t)
+			response := queryDNSCanary(t, server, parent, domain, tt.enabled)
+			wantLookups := int64(0)
+			if tt.enabled {
+				wantLookups = 1
+			}
+			if got := lookups.Load(); got != wantLookups {
+				t.Fatalf("performed %d canary lookups, want %d", got, wantLookups)
+			}
+			if tt.registered && tt.enabled {
+				assertDNSCanaryResponse(t, response)
+				assertDNSCanaryRecord(t, canary, 1)
+				assertDNSCanaryEvents(t, collectDNSCanaryEvents(t, events, false), canary, 1)
+				return
+			}
+			assertDNSCanaryNOP(t, response, checksum)
+			assertDNSCanaryEvents(t, collectDNSCanaryEvents(t, events, false), canary, 0)
+			if tt.registered {
+				assertDNSCanaryRecord(t, canary, 0)
+			}
+		})
+	}
+}
+
+func TestDNSCanaryFailedTriggerDoesNotPublish(t *testing.T) {
+	const parent = "canary-failure.invalid."
+	canary := createDNSCanaryRecord(t, models.NewUUID().String()+"."+parent)
+	events := subscribeDNSCanaryEvents(t)
+	reject := rejectDNSCanaryUpdates(t)
+	server := newTestDNSServer()
+
+	assertDNSCanaryResponse(t, queryDNSCanary(t, server, parent, canary.Domain, true))
+	assertDNSCanaryRecord(t, canary, 0)
+	assertDNSCanaryEvents(t, collectDNSCanaryEvents(t, events, false), canary, 0)
+
+	reject.Store(false)
+	assertDNSCanaryResponse(t, queryDNSCanary(t, server, parent, canary.Domain, true))
+	assertDNSCanaryRecord(t, canary, 1)
+	assertDNSCanaryEvents(t, collectDNSCanaryEvents(t, events, false), canary, 1)
+}
+
+func createDNSCanaryRecord(t *testing.T, domain string) *models.DNSCanary {
+	t.Helper()
+	canary := &models.DNSCanary{ImplantName: "dns-canary-" + models.NewUUID().String(), Domain: domain}
+	if err := db.Session().Create(canary).Error; err != nil {
+		t.Fatalf("create canary: %v", err)
+	}
+	t.Cleanup(func() {
+		if err := db.Session().Where("id = ?", canary.ID).Delete(&models.DNSCanary{}).Error; err != nil {
+			t.Errorf("delete canary: %v", err)
+		}
+	})
+	return canary
+}
+
+func assertDNSCanaryRecord(t *testing.T, canary *models.DNSCanary, count uint32) models.DNSCanary {
+	t.Helper()
+	var stored models.DNSCanary
+	if err := db.Session().Where("id = ?", canary.ID).First(&stored).Error; err != nil {
+		t.Fatalf("reload canary: %v", err)
+	}
+	if stored.Count != count || stored.Triggered != (count > 0) {
+		t.Fatalf("stored count=%d triggered=%v, want count=%d triggered=%v", stored.Count, stored.Triggered, count, count > 0)
+	}
+	if stored.FirstTrigger.IsZero() != (count == 0) || stored.LatestTrigger.IsZero() != (count == 0) {
+		t.Error("trigger timestamps do not match the persisted count")
+	}
+	return stored
+}
+
+func countDNSCanaryLookups(t *testing.T) *atomic.Int64 {
+	t.Helper()
+	var count atomic.Int64
+	name := "test:count-dns-canary-lookups-" + models.NewUUID().String()
+	if err := db.Client.Callback().Query().Before("gorm:query").Register(name, func(tx *gorm.DB) {
+		if tx.Statement.Table == "dns_canaries" {
+			count.Add(1)
+		}
+	}); err != nil {
+		t.Fatalf("register canary query counter: %v", err)
+	}
+	t.Cleanup(func() { _ = db.Client.Callback().Query().Remove(name) })
+	return &count
+}
+
+func rejectDNSCanaryUpdates(t *testing.T) *atomic.Bool {
+	t.Helper()
+	var reject atomic.Bool
+	reject.Store(true)
+	name := "test:reject-dns-canary-updates-" + models.NewUUID().String()
+	if err := db.Client.Callback().Update().Before("gorm:update").Register(name, func(tx *gorm.DB) {
+		if reject.Load() && tx.Statement.Table == "dns_canaries" {
+			_ = tx.AddError(errors.New("injected DNS canary update failure"))
+		}
+	}); err != nil {
+		t.Fatalf("register canary update failure: %v", err)
+	}
+	t.Cleanup(func() { _ = db.Client.Callback().Update().Remove(name) })
+	return &reject
+}
+
+func subscribeDNSCanaryEvents(t *testing.T) chan core.Event {
+	t.Helper()
+	events := core.EventBroker.Subscribe()
+	t.Cleanup(func() { core.EventBroker.Unsubscribe(events) })
+	// Subscription registration is asynchronous. Complete a broker round trip
+	// before sending requests so the first canary event cannot be missed.
+	collectDNSCanaryEvents(t, events, true)
+	return events
+}
+
+func collectDNSCanaryEvents(t *testing.T, events <-chan core.Event, subscribing bool) []core.Event {
+	t.Helper()
+	marker := core.Event{EventType: "dns-canary-test-barrier", Data: []byte(models.NewUUID().String())}
+	core.EventBroker.Publish(marker)
+	timer := time.NewTimer(5 * time.Second)
+	defer timer.Stop()
+	ticker := time.NewTicker(time.Millisecond)
+	defer ticker.Stop()
+	var canaries []core.Event
+	for {
+		select {
+		case event := <-events:
+			if event.EventType == marker.EventType && string(event.Data) == string(marker.Data) {
+				return canaries
+			}
+			if event.EventType == consts.CanaryEvent {
+				canaries = append(canaries, event)
+			}
+		case <-ticker.C:
+			if subscribing {
+				marker.Data = []byte(models.NewUUID().String())
+				core.EventBroker.Publish(marker)
+			}
+		case <-timer.C:
+			t.Fatal("timed out waiting for DNS canary event barrier")
+		}
+	}
+}
+
+func assertDNSCanaryEvents(t *testing.T, events []core.Event, canary *models.DNSCanary, count int) {
+	t.Helper()
+	if len(events) != count {
+		t.Fatalf("received %d canary events, want %d", len(events), count)
+	}
+	for _, event := range events {
+		if event.Session == nil || event.Session.Name != canary.ImplantName || string(event.Data) != canary.Domain {
+			t.Errorf("unexpected canary event: %+v", event)
+		}
+	}
+}
+
+func newDNSCanaryNOP(t *testing.T) (*SliverDNSServer, string, string, uint32) {
+	t.Helper()
+	server := newTestDNSServer()
+	const sessionID = uint32(0x123456)
+	server.sessions.Store(sessionID, newTestDNSSession(sessionID))
+	parent := models.NewUUID().String() + ".canary-c2.invalid."
+	msg := &dnspb.DNSMessage{Type: dnspb.DNSMessageType_NOP, ID: (1 << 24) | sessionID, Data: []byte("canary-routing")}
+	raw, err := proto.Marshal(msg)
+	if err != nil {
+		t.Fatalf("marshal DNS NOP: %v", err)
+	}
+	domain, err := encodeMessageDomain("."+parent, msg)
+	if err != nil {
+		t.Fatalf("encode DNS NOP: %v", err)
+	}
+	return server, parent, domain, crc32.ChecksumIEEE(raw)
+}
+
+func queryDNSCanary(t *testing.T, server *SliverDNSServer, parent, domain string, enabled bool) *dns.Msg {
+	t.Helper()
+	req := new(dns.Msg)
+	req.SetQuestion(domain, dns.TypeA)
+	writer := &canaryDNSResponseWriter{}
+	server.HandleDNSRequest([]string{parent}, enabled, writer, req)
+	if writer.response == nil {
+		t.Fatal("DNS handler did not write a response")
+	}
+	return writer.response
+}
+
+func assertDNSCanaryResponse(t *testing.T, response *dns.Msg) {
+	t.Helper()
+	if response.Rcode != dns.RcodeNameError || !response.Authoritative {
+		t.Fatalf("canary response rcode=%d authoritative=%v, want authoritative NXDOMAIN", response.Rcode, response.Authoritative)
+	}
+}
+
+func assertDNSCanaryNOP(t *testing.T, response *dns.Msg, checksum uint32) {
+	t.Helper()
+	if response.Rcode != dns.RcodeSuccess || len(response.Answer) != 1 {
+		t.Fatalf("NOP response rcode=%d answers=%d, want success with one answer", response.Rcode, len(response.Answer))
+	}
+	answer, ok := response.Answer[0].(*dns.A)
+	if !ok || len(answer.A) != 4 {
+		t.Fatalf("unexpected DNS NOP answer: %v", response.Answer[0])
+	}
+	if got := binary.LittleEndian.Uint32(answer.A); got != checksum {
+		t.Fatalf("NOP checksum=%d, want %d", got, checksum)
+	}
+}
+
+type canaryDNSResponseWriter struct {
+	response *dns.Msg
+}
+
+func (w *canaryDNSResponseWriter) LocalAddr() net.Addr  { return &net.UDPAddr{} }
+func (w *canaryDNSResponseWriter) RemoteAddr() net.Addr { return &net.UDPAddr{} }
+func (w *canaryDNSResponseWriter) Close() error         { return nil }
+func (w *canaryDNSResponseWriter) TsigStatus() error    { return nil }
+func (w *canaryDNSResponseWriter) TsigTimersOnly(bool)  {}
+func (w *canaryDNSResponseWriter) Hijack()              {}
+
+func (w *canaryDNSResponseWriter) WriteMsg(msg *dns.Msg) error {
+	w.response = msg.Copy()
+	return nil
+}
+
+func (w *canaryDNSResponseWriter) Write(data []byte) (int, error) {
+	msg := new(dns.Msg)
+	if err := msg.Unpack(data); err != nil {
+		return 0, err
+	}
+	w.response = msg
+	return len(data), nil
+}

server/c2/http_test.go

@@ -1,173 +0,0 @@
-package c2
-
-/*
-	Sliver Implant Framework
-	Copyright (C) 2021  Bishop Fox
-
-	This program is free software: you can redistribute it and/or modify
-	it under the terms of the GNU General Public License as published by
-	the Free Software Foundation, either version 3 of the License, or
-	(at your option) any later version.
-
-	This program is distributed in the hope that it will be useful,
-	but WITHOUT ANY WARRANTY; without even the implied warranty of
-	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-	GNU General Public License for more details.
-
-	You should have received a copy of the GNU General Public License
-	along with this program.  If not, see <https://www.gnu.org/licenses/>.
-*/
-
-// import (
-// 	"bytes"
-// 	"fmt"
-// 	insecureRand "math/rand"
-// 	"net/http"
-// 	"net/http/httptest"
-// 	"net/url"
-// 	"testing"
-
-// 	implantCrypto "github.com/bishopfox/sliver/implant/sliver/cryptography"
-// 	implantEncoders "github.com/bishopfox/sliver/implant/sliver/encoders"
-// 	implantTransports "github.com/bishopfox/sliver/implant/sliver/transports/httpclient"
-// 	"github.com/bishopfox/sliver/protobuf/sliverpb"
-// 	"github.com/bishopfox/sliver/server/configs"
-// 	"github.com/bishopfox/sliver/server/cryptography"
-// 	"google.golang.org/protobuf/proto"
-// )
-
-// func TestStartSessionHandler(t *testing.T) {
-
-// 	implantTransports.SetNonceQueryArgChars("abcdedfghijklmnopqrstuvwxyz")
-
-// 	server, err := StartHTTPListener(&HTTPServerConfig{
-// 		Addr:       "127.0.0.1:8888",
-// 		Secure:     false,
-// 		EnforceOTP: true,
-// 	})
-// 	if err != nil {
-// 		t.Fatalf("Listener failed to start %s", err)
-// 		return
-// 	}
-
-// 	client := implantTransports.SliverHTTPClient{}
-// 	baseURL := &url.URL{
-// 		Scheme: "http",
-// 		Host:   "127.0.0.1:8888",
-// 		Path:   fmt.Sprintf("/test/foo.%s", c2Config.ImplantConfig.StartSessionFileExt),
-// 	}
-// 	nonce, encoder := implantEncoders.RandomEncoder(0)
-// 	testURL := client.NonceQueryArgument(baseURL, nonce)
-// 	testURL = client.OTPQueryArgument(testURL, implantCrypto.GetOTPCode())
-
-// 	// Generate key exchange request
-// 	sKey := cryptography.RandomKey()
-// 	httpSessionInit := &sliverpb.HTTPSessionInit{Key: sKey[:]}
-// 	data, _ := proto.Marshal(httpSessionInit)
-// 	encryptedSessionInit, err := implantCrypto.ECCEncryptToServer(data)
-// 	if err != nil {
-// 		t.Fatalf("Failed to encrypt session init %s", err)
-// 	}
-// 	payload, _ := encoder.Encode(encryptedSessionInit)
-// 	body := bytes.NewReader(payload)
-
-// 	validReq := httptest.NewRequest(http.MethodPost, testURL.String(), body)
-// 	t.Logf("[http] req request uri: '%v'", validReq.RequestURI)
-// 	rr := httptest.NewRecorder()
-// 	server.HTTPServer.Handler.ServeHTTP(rr, validReq)
-// 	if status := rr.Code; status != http.StatusOK {
-// 		t.Fatalf("handler returned wrong status code: got %d want %d", status, http.StatusOK)
-// 	}
-
-// }
-
-// func TestGetOTPFromURL(t *testing.T) {
-
-// 	implantTransports.SetNonceQueryArgChars("abcdedfghijklmnopqrstuvwxyz")
-
-// 	client := implantTransports.SliverHTTPClient{}
-// 	for i := 0; i < 100; i++ {
-// 		baseURL := &url.URL{
-// 			Scheme: "http",
-// 			Host:   "127.0.0.1:8888",
-// 			Path:   "/test/foo.txt",
-// 		}
-// 		value := fmt.Sprintf("%d", insecureRand.Intn(99999999))
-// 		testURL := client.OTPQueryArgument(baseURL, value)
-// 		urlValue, err := getOTPFromURL(testURL)
-// 		if err != nil {
-// 			t.Fatal(err)
-// 		}
-// 		if urlValue != value {
-// 			t.Fatalf("Mismatched OTP values %s (%s != %s)", testURL.String(), value, urlValue)
-// 		}
-// 	}
-// }
-
-// func TestGetNonceFromURL(t *testing.T) {
-
-// 	implantTransports.SetNonceQueryArgChars("abcdedfghijklmnopqrstuvwxyz")
-
-// 	client := implantTransports.SliverHTTPClient{}
-// 	for i := 0; i < 100; i++ {
-// 		baseURL := &url.URL{
-// 			Scheme: "http",
-// 			Host:   "127.0.0.1:8888",
-// 			Path:   "/test/foo.txt",
-// 		}
-// 		nonce, encoder := implantEncoders.RandomEncoder(0)
-// 		testURL := client.NonceQueryArgument(baseURL, nonce)
-// 		t.Log(testURL.String())
-// 		urlNonce, err := getNonceFromURL(testURL)
-// 		if err != nil {
-// 			t.Errorf("Nonce '%d' triggered error from %#v", nonce, encoder)
-// 			t.Fatal(err)
-// 		}
-// 		if urlNonce != nonce {
-// 			t.Fatalf("Mismatched encoder nonces %s (%d != %d)", testURL.String(), nonce, urlNonce)
-// 		}
-// 	}
-// }
-
-// func TestGetNonceFromURLWithCustomQueryArgs(t *testing.T) {
-// 	for j := 0; j < 100; j++ {
-
-// 		queryArgs := randomArgs(1)
-// 		implantTransports.SetNonceQueryArgChars(queryArgs)
-// 		t.Logf("Using query args: %s", queryArgs)
-
-// 		client := implantTransports.SliverHTTPClient{}
-// 		for i := 0; i < 10; i++ {
-// 			baseURL := &url.URL{
-// 				Scheme: "http",
-// 				Host:   "127.0.0.1:8888",
-// 				Path:   "/test/foo.txt",
-// 			}
-// 			nonce, encoder := implantEncoders.RandomEncoder(0)
-// 			testURL := client.NonceQueryArgument(baseURL, nonce)
-// 			t.Log(testURL.String())
-// 			urlNonce, err := getNonceFromURL(testURL)
-// 			if err != nil {
-// 				t.Errorf("Nonce '%d' triggered error from %#v", nonce, encoder)
-// 				t.Fatal(err)
-// 			}
-// 			if urlNonce != nonce {
-// 				t.Fatalf("Mismatched encoder nonces %s (%d != %d)", testURL.String(), nonce, urlNonce)
-// 			}
-// 		}
-// 	}
-// }
-
-// func randomArgs(min int) string {
-// 	var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!@#$%^&*()")
-// 	size := insecureRand.Intn(25) + min
-// 	b := make(map[rune]interface{}, size)
-// 	for len(b) <= size {
-// 		b[letters[insecureRand.Intn(len(letters))]] = nil
-// 	}
-// 	var keys string
-// 	for k := range b {
-// 		keys += string(k)
-// 	}
-// 	return keys
-// }

server/core/rtunnels/authorization_test.go

@@ -465,16 +465,48 @@ func TestRegistryRevokeSessionIsScopedAndIdempotent(t *testing.T) {
 func TestRegistryGeneratesUniqueOpaqueIDs(t *testing.T) {
 	t.Parallel()
 
-	registry := NewRegistry()
-	seen := map[AuthorizationID]struct{}{}
-	for index := 0; index < 256; index++ {
-		id, err := registry.Begin("session", "127.0.0.1:8080", 0)
-		must.NoError(t, err)
-		assert.Len(t, id, 43)
-		assert.False(t, strings.ContainsAny(id.String(), "+/="))
-		_, duplicate := seen[id]
-		assert.False(t, duplicate)
-		seen[id] = struct{}{}
+	firstBytes := bytes.Repeat([]byte{0xff}, authorizationIDBytes)
+	secondBytes := bytes.Repeat([]byte{0x42}, authorizationIDBytes)
+	for _, test := range []struct {
+		name          string
+		random        []byte
+		wantError     bool
+		wantRemaining int
+	}{
+		{name: "distinct IDs", random: secondBytes},
+		{name: "retry collision", random: bytes.Join([][]byte{firstBytes, secondBytes}, nil)},
+		{
+			name:          "exhaust collisions",
+			random:        bytes.Join([][]byte{bytes.Repeat(firstBytes, maxAuthorizationIDAttempts), secondBytes}, nil),
+			wantError:     true,
+			wantRemaining: authorizationIDBytes,
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			random := bytes.NewReader(bytes.Join([][]byte{firstBytes, test.random}, nil))
+			registry := newRegistry(random)
+			firstID, err := registry.Begin("session", "127.0.0.1:8080", 0)
+			must.NoError(t, err)
+			assert.Len(t, firstID, 43)
+			assert.False(t, strings.ContainsAny(firstID.String(), "+/="))
+
+			secondID, err := registry.Begin("session", "127.0.0.1:8081", 0)
+			if test.wantError {
+				must.ErrorIs(t, err, ErrAuthorizationIDGeneration)
+				assert.Empty(t, secondID)
+				assert.Len(t, registry.List("session"), 1)
+			} else {
+				must.NoError(t, err)
+				assert.NotEqual(t, firstID, secondID)
+				assert.Len(t, secondID, 43)
+				assert.False(t, strings.ContainsAny(secondID.String(), "+/="))
+				assert.Len(t, registry.List("session"), 2)
+			}
+			assert.Equal(t, test.wantRemaining, random.Len())
+			first, ok := registry.Lookup("session", firstID)
+			must.True(t, ok)
+			assert.Equal(t, "127.0.0.1:8080", first.Address)
+		})
 	}
 }
 
@@ -484,13 +516,6 @@ func (errorReader) Read(_ []byte) (int, error) {
 	return 0, io.ErrUnexpectedEOF
 }
 
-func TestAuthorizationErrorsAreComparable(t *testing.T) {
-	t.Parallel()
-
-	_, err := canonicalizeAddress("invalid")
-	assert.True(t, errors.Is(err, ErrInvalidForwardAddress))
-}
-
 func acquirePlanForTest(registry *Registry, sessionID string, id AuthorizationID) (dialPlan, error) {
 	plan, _, reservation, err := registry.reserve(sessionID, id, "")
 	if err == nil {

server/cryptography/cryptography_test.go

@@ -134,26 +134,27 @@ func TestAgeKeyExTamper(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	// Tamper with the ciphertext
-	ciphertext := allCiphertext[32:]
-	ciphertext[insecureRand.Intn(len(ciphertext))] ^= 0xFF
-	_, err = AgeKeyExFromImplant(
-		serverAgeKeyPair.Private,
-		implantPeerAgeKeyPair.Private,
-		ciphertext,
-	)
-	if err == nil {
-		t.Fatal(err)
-	}
+	tamperedCiphertext := bytes.Clone(allCiphertext[32:])
+	tamperedCiphertext[len(tamperedCiphertext)-1] ^= 0xFF
 
-	// Leave an invalid header with valid ciphertext
-	_, err = AgeKeyExFromImplant(
-		serverAgeKeyPair.Private,
-		implantPeerAgeKeyPair.Private,
-		allCiphertext,
-	)
-	if err == nil {
-		t.Fatal(err)
+	for _, test := range []struct {
+		name       string
+		ciphertext []byte
+	}{
+		{name: "tampered ciphertext", ciphertext: tamperedCiphertext},
+		// Retaining the public key hash adds an invalid header to otherwise valid ciphertext.
+		{name: "invalid header", ciphertext: bytes.Clone(allCiphertext)},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			_, err := AgeKeyExFromImplant(
+				serverAgeKeyPair.Private,
+				implantPeerAgeKeyPair.Private,
+				test.ciphertext,
+			)
+			if err == nil {
+				t.Fatal("expected key exchange decryption to reject invalid ciphertext")
+			}
+		})
 	}
 }
 

server/db/canaries_test.go

@@ -0,0 +1,52 @@
+package db
+
+import (
+	"errors"
+	"strings"
+	"testing"
+
+	"github.com/bishopfox/sliver/server/db/models"
+)
+
+func TestCanaryByDomain(t *testing.T) {
+	canary := &models.DNSCanary{
+		ImplantName: "canary_lookup_" + models.NewUUID().String(),
+		Domain:      "lookup-" + models.NewUUID().String() + ".Canary.Example.",
+	}
+	if err := Session().Create(canary).Error; err != nil {
+		t.Fatalf("create canary: %v", err)
+	}
+	t.Cleanup(func() {
+		if err := Session().Delete(canary).Error; err != nil {
+			t.Errorf("delete canary: %v", err)
+		}
+	})
+
+	for _, domain := range []string{
+		canary.Domain,
+		strings.ToLower(canary.Domain),
+		strings.ToUpper(canary.Domain),
+		" " + strings.TrimSuffix(canary.Domain, ".") + " ",
+	} {
+		t.Run(domain, func(t *testing.T) {
+			found, err := CanaryByDomain(domain)
+			if err != nil {
+				t.Fatalf("lookup canary: %v", err)
+			}
+			if found.ID != canary.ID.String() || found.Domain != canary.Domain {
+				t.Errorf("lookup returned %v, want ID %s and domain %s", found, canary.ID, canary.Domain)
+			}
+		})
+	}
+}
+
+func TestCanaryByDomainNotFound(t *testing.T) {
+	for _, domain := range []string{"", " \t", "missing-" + models.NewUUID().String() + ".invalid."} {
+		t.Run(domain, func(t *testing.T) {
+			found, err := CanaryByDomain(domain)
+			if !errors.Is(err, ErrRecordNotFound) || found != nil {
+				t.Fatalf("lookup returned (%v, %v), want (nil, record not found)", found, err)
+			}
+		})
+	}
+}

server/db/helpers.go

@@ -1,3 +1,4 @@
+// Package db provides persistence and queries for Sliver server state.
 package db
 
 /*
@@ -742,13 +743,25 @@ func ListCanaries() ([]*clientpb.DNSCanary, error) {
 
 // CanaryByDomain - Check if a canary exists
 func CanaryByDomain(domain string) (*clientpb.DNSCanary, error) {
-	if len(domain) < 1 {
+	domain = strings.ToLower(strings.TrimSpace(domain))
+	if domain == "" {
 		return nil, ErrRecordNotFound
 	}
+	if !strings.HasSuffix(domain, ".") {
+		domain += "."
+	}
 	dbSession := Session()
 	canary := models.DNSCanary{}
-	err := dbSession.Where(&models.DNSCanary{Domain: domain}).First(&canary).Error
-	return canary.ToProtobuf(), err
+	// Older builds may have stored mixed-case parent domains. DNS names are
+	// case-insensitive, so normalize both sides without rewriting those records.
+	result := dbSession.Where("LOWER(domain) = ?", domain).Limit(1).Find(&canary)
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	if result.RowsAffected == 0 {
+		return nil, ErrRecordNotFound
+	}
+	return canary.ToProtobuf(), nil
 }
 
 // WebsiteByName - Get website by name

server/db/models/canary.go

@@ -48,6 +48,7 @@ func (c *DNSCanary) BeforeCreate(tx *gorm.DB) (err error) {
 // ToProtobuf - Converts to protobuf object
 func (c *DNSCanary) ToProtobuf() *clientpb.DNSCanary {
 	return &clientpb.DNSCanary{
+		ID:             c.ID.String(),
 		ImplantName:    c.ImplantName,
 		Domain:         c.Domain,
 		Triggered:      c.Triggered,

server/db/models/uuid_test.go

@@ -19,13 +19,10 @@ package models
 */
 
 import (
-	"database/sql"
-	"database/sql/driver"
 	"encoding/json"
 	"path/filepath"
 	"strings"
 	"testing"
-	uuid "uuid"
 
 	gosqlite "github.com/glebarez/sqlite"
 	"gorm.io/gorm"
@@ -89,9 +86,6 @@ func TestUUIDSQLAndJSONEncoding(t *testing.T) {
 	want := NewUUID()
 	standard := want.Standard()
 
-	var _ sql.Scanner = (*UUID)(nil)
-	var _ driver.Valuer = UUID{}
-
 	value, err := want.Value()
 	if err != nil {
 		t.Fatalf("Value: %v", err)
@@ -188,11 +182,3 @@ func TestUUIDGORMRoundTrip(t *testing.T) {
 		t.Fatalf("stored UUID = %q, want canonical %q", storedID, want.ID.String())
 	}
 }
-
-func TestUUIDMatchesStandardLibrary(t *testing.T) {
-	standard := uuid.NewV4()
-	model := UUIDFrom(standard)
-	if model.Standard() != standard {
-		t.Fatalf("Standard = %q, want %q", model.Standard(), standard)
-	}
-}

server/db/wasmsqlite/sqlite_version_test.go

@@ -1,23 +0,0 @@
-package wasmsqlite
-
-import (
-	"database/sql"
-	"log"
-	"testing"
-)
-
-func TestSQLiteVersion(t *testing.T) {
-	var version string
-
-	db, err := sql.Open(DriverName, ":memory:")
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	row := db.QueryRow("select sqlite_version()")
-	if row.Scan(&version) != nil {
-		log.Fatal(err)
-	}
-
-	t.Log(version)
-}

server/encoders/encoders_test.go

@@ -20,206 +20,65 @@ package encoders
 
 import (
 	"bytes"
-	"crypto/rand"
-	insecureRand "math/rand"
+	"fmt"
 	"testing"
 
 	implantEncoders "github.com/bishopfox/sliver/implant/sliver/encoders"
 	util "github.com/bishopfox/sliver/util/encoders"
 )
 
-const (
-	sampleSizeMax = 8192
-)
-
-func randomDataRandomSize(maxSize int) []byte {
-	buf := make([]byte, insecureRand.Intn(maxSize))
-	rand.Read(buf)
-	return buf
-}
-
-func TestCompatibilityBase64(t *testing.T) {
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(sampleSizeMax)
-		output, _ := Base64.Encode(sample)
-		data, err := implantEncoders.Base64.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into base64")
-			return
-		}
-		if !bytes.Equal(sample, data) {
-			t.Errorf("sample does not match returned\n%#v != %#v", sample, data)
-		}
-
-		sample2 := randomDataRandomSize(sampleSizeMax)
-		output, _ = implantEncoders.Base64.Encode(sample2)
-		data, err = Base64.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into base64")
-			return
-		}
-		if !bytes.Equal(sample2, data) {
-			t.Errorf("sample2 does not match returned\n%#v != %#v", sample2, data)
-		}
-	}
-}
-
-func TestCompatibilityBase58(t *testing.T) {
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(sampleSizeMax)
-		output, _ := Base58.Encode(sample)
-		data, err := implantEncoders.Base58.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into base58")
-			return
-		}
-		if !bytes.Equal(sample, data) {
-			t.Errorf("sample does not match returned\n%#v != %#v", sample, data)
-		}
-
-		sample2 := randomDataRandomSize(sampleSizeMax)
-		output, _ = Base58.Encode(sample2)
-		data, err = implantEncoders.Base58.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into base58")
-			return
-		}
-		if !bytes.Equal(sample2, data) {
-			t.Errorf("sample2 does not match returned\n%#v != %#v", sample2, data)
-		}
-	}
-}
-
-func TestCompatibilityBase32(t *testing.T) {
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(sampleSizeMax)
-		output, _ := Base32.Encode(sample)
-		data, err := implantEncoders.Base32.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into base32")
-			return
-		}
-		if !bytes.Equal(sample, data) {
-			t.Errorf("sample does not match returned\n%#v != %#v", sample, data)
-		}
-
-		sample2 := randomDataRandomSize(sampleSizeMax)
-		output, _ = Base32.Encode(sample2)
-		data, err = implantEncoders.Base32.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into base32")
-			return
-		}
-		if !bytes.Equal(sample2, data) {
-			t.Errorf("sample2 does not match returned\n%#v != %#v", sample2, data)
-		}
-	}
-}
-
-func TestCompatibilityEnglish(t *testing.T) {
-
+func TestEncoderCompatibility(t *testing.T) {
 	util.SetEnglishDictionary(getTestEnglishDictionary())
 
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(sampleSizeMax)
-		output, _ := English.Encode(sample)
-		data, err := implantEncoders.English.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into english")
-			return
-		}
-		if !bytes.Equal(sample, data) {
-			t.Errorf("sample does not match returned\n%#v != %#v", sample, data)
-		}
-
-		sample2 := randomDataRandomSize(sampleSizeMax)
-		output, _ = implantEncoders.English.Encode(sample2)
-		data, err = English.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into english")
-			return
-		}
-		if !bytes.Equal(sample2, data) {
-			t.Errorf("sample2 does not match returned\n%#v != %#v", sample2, data)
-		}
+	encoders := []struct {
+		name    string
+		server  util.Encoder
+		implant util.Encoder
+	}{
+		{"base64", Base64, implantEncoders.Base64},
+		{"base58", Base58, implantEncoders.Base58},
+		{"base32", Base32, implantEncoders.Base32},
+		{"english", English, implantEncoders.English},
+		{"hex", Hex, implantEncoders.Hex},
+		{"gzip", Gzip, implantEncoders.Gzip},
+		{"png", PNG, implantEncoders.PNG},
 	}
-}
-
-func TestCompatibilityHex(t *testing.T) {
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(sampleSizeMax)
-		output, _ := Hex.Encode(sample)
-		data, err := implantEncoders.Hex.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into hex")
-			return
-		}
-		if !bytes.Equal(sample, data) {
-			t.Errorf("sample does not match returned\n%#v != %#v", sample, data)
-		}
-
-		sample2 := randomDataRandomSize(sampleSizeMax)
-		output, _ = Hex.Encode(sample2)
-		data, err = implantEncoders.Hex.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into hex")
-			return
-		}
-		if !bytes.Equal(sample2, data) {
-			t.Errorf("sample2 does not match returned\n%#v != %#v", sample2, data)
-		}
-	}
-}
-
-func TestCompatibilityGzip(t *testing.T) {
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(sampleSizeMax)
-		output, _ := Gzip.Encode(sample)
-		data, err := implantEncoders.Gzip.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into gzip")
-			return
-		}
-		if !bytes.Equal(sample, data) {
-			t.Errorf("sample does not match returned\n%#v != %#v", sample, data)
-		}
-
-		sample2 := randomDataRandomSize(sampleSizeMax)
-		output, _ = Gzip.Encode(sample2)
-		data, err = implantEncoders.Gzip.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into gzip")
-			return
-		}
-		if !bytes.Equal(sample2, data) {
-			t.Errorf("sample2 does not match returned\n%#v != %#v", sample2, data)
-		}
-	}
-}
-
-func TestCompatibilityPNG(t *testing.T) {
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(sampleSizeMax)
-		output, _ := PNG.Encode(sample)
-		data, err := implantEncoders.PNG.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into png")
-			return
-		}
-		if !bytes.Equal(sample, data) {
-			t.Errorf("sample does not match returned\n%#v != %#v", sample, data)
-		}
-
-		sample2 := randomDataRandomSize(sampleSizeMax)
-		output, _ = PNG.Encode(sample2)
-		data, err = implantEncoders.PNG.Decode(output)
-		if err != nil {
-			t.Error("Failed to encode/decode sample data into png")
-			return
-		}
-		if !bytes.Equal(sample2, data) {
-			t.Errorf("sample2 does not match returned\n%#v != %#v", sample2, data)
-		}
+	for _, encoder := range encoders {
+		t.Run(encoder.name, func(t *testing.T) {
+			directions := []struct {
+				name    string
+				encoder util.Encoder
+				decoder util.Encoder
+			}{
+				{"server_to_implant", encoder.server, encoder.implant},
+				{"implant_to_server", encoder.implant, encoder.server},
+			}
+			for _, direction := range directions {
+				t.Run(direction.name, func(t *testing.T) {
+					// Cover encoding block boundaries, every byte value, leading and
+					// trailing zero bytes, and the old random corpus's upper bound.
+					for _, size := range []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 255, 256, 257, 8192} {
+						t.Run(fmt.Sprintf("bytes_%d", size), func(t *testing.T) {
+							payload := make([]byte, size)
+							for i := range payload {
+								payload[i] = byte(i)
+							}
+							encoded, err := direction.encoder.Encode(payload)
+							if err != nil {
+								t.Fatalf("encode: %v", err)
+							}
+							decoded, err := direction.decoder.Decode(encoded)
+							if err != nil {
+								t.Fatalf("decode: %v", err)
+							}
+							if !bytes.Equal(payload, decoded) {
+								t.Fatalf("round trip changed the %d-byte payload (decoded %d bytes)", len(payload), len(decoded))
+							}
+						})
+					}
+				})
+			}
+		})
 	}
 }
 

server/generate/binaries_test.go

@@ -1,3 +1,5 @@
+//go:build sliver_generate
+
 package generate
 
 /*

server/generate/canaries.go

@@ -21,11 +21,13 @@ package generate
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/bishopfox/sliver/protobuf/clientpb"
 	"github.com/bishopfox/sliver/server/db"
 	"github.com/bishopfox/sliver/server/db/models"
 	"github.com/bishopfox/sliver/util"
+	"gorm.io/gorm"
 )
 
 const (
@@ -50,11 +52,42 @@ func canarySubDomain() string {
 	return string(subdomain)
 }
 
-// UpdateCanary - Update an existing canary
-func UpdateCanary(canary *clientpb.DNSCanary) error {
-	dbSession := db.Session()
-	result := dbSession.Save(&canary)
-	return result.Error
+// TriggerCanary records a DNS query and reports whether it first triggered the
+// existing canary. The caller's status may be stale; only its ID is used.
+func TriggerCanary(canary *clientpb.DNSCanary) (bool, error) {
+	if canary == nil || canary.ID == "" {
+		return false, fmt.Errorf("missing canary ID")
+	}
+	id, err := models.ParseUUID(canary.ID)
+	if err != nil || id == models.NilUUID() {
+		return false, fmt.Errorf("invalid canary ID %q", canary.ID)
+	}
+
+	firstTrigger := false
+	err = db.Session().Transaction(func(tx *gorm.DB) error {
+		count := tx.Model(&models.DNSCanary{}).Where("id = ?", id).
+			UpdateColumn("count", gorm.Expr("count + ?", 1))
+		if count.Error != nil {
+			return count.Error
+		}
+		if count.RowsAffected != 1 {
+			return db.ErrRecordNotFound
+		}
+
+		// The increment locks the existing row before timestamps are sampled,
+		// so waiting queries cannot move LatestTrigger backwards.
+		first := tx.Model(&models.DNSCanary{}).
+			Where("id = ? AND triggered = ?", id, false).
+			Updates(models.DNSCanary{Triggered: true, FirstTrigger: time.Now()})
+		if first.Error != nil {
+			return first.Error
+		}
+		firstTrigger = first.RowsAffected == 1
+
+		return tx.Model(&models.DNSCanary{}).Where("id = ?", id).
+			Updates(models.DNSCanary{LatestTrigger: time.Now()}).Error
+	})
+	return firstTrigger && err == nil, err
 }
 
 // CanaryGenerator - Holds data related to canary generation
@@ -75,7 +108,7 @@ func (g *CanaryGenerator) GenerateCanary() string {
 	// Don't need secure random here
 	index := util.Intn(len(g.ParentDomains))
 
-	parentDomain := g.ParentDomains[index]
+	parentDomain := strings.ToLower(strings.TrimSpace(g.ParentDomains[index]))
 	parentDomain = strings.TrimPrefix(parentDomain, ".")
 	if !strings.HasSuffix(parentDomain, ".") {
 		parentDomain += "." // Ensure we have the FQDN

server/generate/canaries_test.go

@@ -0,0 +1,126 @@
+//go:build sliver_generate
+
+package generate
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/bishopfox/sliver/protobuf/clientpb"
+	"github.com/bishopfox/sliver/server/certs"
+	"github.com/bishopfox/sliver/server/configs"
+	"github.com/bishopfox/sliver/server/db"
+	"github.com/bishopfox/sliver/server/db/models"
+)
+
+// Canary domains must remain visible to a binary scanner after the production
+// compiler and obfuscator run. Checking rendered source alone misses both
+// literal obfuscation and linker removal of unused type metadata.
+func TestCanariesSurviveObfuscation(t *testing.T) {
+	const parentDomain = "canary-regression.invalid."
+	certs.SetupCAs()
+	tests := []struct {
+		name         string
+		goos         string
+		protocol     string
+		beacon       bool
+		wantCanaries int
+	}{
+		{name: "windows_http_session", goos: "windows", protocol: "http", wantCanaries: 1},
+		{name: "linux_http_beacon", goos: "linux", protocol: "http", beacon: true, wantCanaries: 1},
+		{name: "windows_mtls_session", goos: "windows", protocol: "mtls", wantCanaries: 2},
+		{name: "linux_wireguard_session", goos: "linux", protocol: "wg", wantCanaries: 2},
+		{name: "linux_http_without_canaries", goos: "linux", protocol: "http"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			name := "canary_" + strings.ReplaceAll(models.NewUUID().String(), "-", "")
+			config := &clientpb.ImplantConfig{
+				GOOS:             tt.goos,
+				GOARCH:           "amd64",
+				Debug:            false,
+				ObfuscateSymbols: true,
+				NetGoEnabled:     true,
+				IsBeacon:         tt.beacon,
+				C2: []*clientpb.ImplantC2{
+					{URL: tt.protocol + "://c2.invalid"},
+				},
+			}
+			if tt.wantCanaries > 0 {
+				config.CanaryDomains = []string{parentDomain}
+			}
+			switch tt.protocol {
+			case "http":
+				config.IncludeHTTP = true
+			case "mtls":
+				config.IncludeMTLS = true
+			case "wg":
+				config.IncludeWG = true
+				config.WGKeyExchangePort = 1234
+				config.WGTcpCommsPort = 5678
+				certs.SetupWGKeys()
+			}
+
+			// Register cleanup before generation so failed builds do not leave
+			// source trees or canary records behind.
+			projectDir := filepath.Join(GetSliversDir(), config.GOOS, config.GOARCH, name)
+			t.Cleanup(func() {
+				if err := os.RemoveAll(projectDir); err != nil {
+					t.Errorf("remove generated project: %v", err)
+				}
+				if err := db.Session().Where("implant_name = ?", name).Delete(&models.DNSCanary{}).Error; err != nil {
+					t.Errorf("remove generated canaries: %v", err)
+				}
+			})
+
+			build, err := GenerateConfig(name, config)
+			if err != nil {
+				t.Fatalf("generate implant config: %v", err)
+			}
+			httpConfig := configs.GenerateDefaultHTTPC2Config()
+			binPath, err := SliverExecutable(name, build, config, httpConfig.ImplantConfig)
+			if err != nil {
+				t.Fatalf("build obfuscated implant: %v", err)
+			}
+			binary, err := os.ReadFile(binPath)
+			if err != nil {
+				t.Fatalf("read obfuscated implant: %v", err)
+			}
+			var canaries []models.DNSCanary
+			if err := db.Session().Where("implant_name = ?", name).Find(&canaries).Error; err != nil {
+				t.Fatalf("read generated canaries: %v", err)
+			}
+			assertCanaryBinary(t, binary, canaries, parentDomain, tt.wantCanaries)
+			t.Logf("checked %d registered canaries in %s/%s obfuscated binary (%d bytes)", len(canaries), config.GOOS, config.GOARCH, len(binary))
+		})
+	}
+}
+
+func assertCanaryBinary(t *testing.T, binary []byte, canaries []models.DNSCanary, parentDomain string, wantCanaries int) {
+	t.Helper()
+
+	if bytes.Contains(binary, []byte("[[GenerateCanary]]")) {
+		t.Error("unexpanded canary placeholder remains in binary")
+	}
+	if len(canaries) != wantCanaries {
+		t.Fatalf("registered %d canaries, want %d", len(canaries), wantCanaries)
+	}
+	if wantCanaries == 0 {
+		if bytes.Contains(binary, []byte(parentDomain)) {
+			t.Error("build without canaries contains the canary parent domain")
+		}
+		return
+	}
+	for _, canary := range canaries {
+		if canary.Domain == "" || !strings.HasSuffix(canary.Domain, "."+parentDomain) {
+			t.Errorf("unexpected generated canary domain %q", canary.Domain)
+			continue
+		}
+		if !bytes.Contains(binary, []byte(canary.Domain)) {
+			t.Errorf("registered canary %q is absent from the obfuscated binary", canary.Domain)
+		}
+	}
+}

server/generate/canary_state_test.go

@@ -0,0 +1,199 @@
+package generate
+
+import (
+	"errors"
+	"strings"
+	"sync/atomic"
+	"testing"
+
+	"github.com/bishopfox/sliver/protobuf/clientpb"
+	"github.com/bishopfox/sliver/server/db"
+	"github.com/bishopfox/sliver/server/db/models"
+	"gorm.io/gorm"
+)
+
+func createCanaryStateFixture(t *testing.T) *models.DNSCanary {
+	t.Helper()
+	canary := &models.DNSCanary{
+		ImplantName: "canary_state_" + models.NewUUID().String(),
+		Domain:      models.NewUUID().String() + ".canary.invalid.",
+	}
+	if err := db.Session().Create(canary).Error; err != nil {
+		t.Fatalf("create canary: %v", err)
+	}
+	t.Cleanup(func() {
+		if err := db.Session().Delete(canary).Error; err != nil {
+			t.Errorf("delete canary: %v", err)
+		}
+	})
+	return canary
+}
+
+func readCanaryState(t *testing.T, id models.UUID) models.DNSCanary {
+	t.Helper()
+	var canary models.DNSCanary
+	if err := db.Session().First(&canary, "id = ?", id).Error; err != nil {
+		t.Fatalf("read canary: %v", err)
+	}
+	return canary
+}
+
+func assertCanaryTriggerCount(t *testing.T, canary models.DNSCanary, count uint32) {
+	t.Helper()
+	if !canary.Triggered || canary.Count != count {
+		t.Errorf("triggered/count = %v/%d, want true/%d", canary.Triggered, canary.Count, count)
+	}
+	if canary.FirstTrigger.IsZero() || canary.LatestTrigger.Before(canary.FirstTrigger) {
+		t.Errorf("invalid trigger times: first=%v latest=%v", canary.FirstTrigger, canary.LatestTrigger)
+	}
+}
+
+func TestTriggerCanaryPersistsRepeatedQueries(t *testing.T) {
+	canary := createCanaryStateFixture(t)
+	// Reuse a stale caller snapshot; persistence must use the database state.
+	snapshot := &clientpb.DNSCanary{ID: canary.ID.String(), Count: 99}
+	first, err := TriggerCanary(snapshot)
+	if err != nil || !first {
+		t.Fatalf("first trigger = (%v, %v), want (true, nil)", first, err)
+	}
+	initial := readCanaryState(t, canary.ID)
+	assertCanaryTriggerCount(t, initial, 1)
+	first, err = TriggerCanary(snapshot)
+	if err != nil || first {
+		t.Fatalf("repeat trigger = (%v, %v), want (false, nil)", first, err)
+	}
+	latest := readCanaryState(t, canary.ID)
+	assertCanaryTriggerCount(t, latest, 2)
+	if !latest.FirstTrigger.Equal(initial.FirstTrigger) || latest.LatestTrigger.Before(initial.LatestTrigger) {
+		t.Error("repeat query changed first trigger time or moved latest trigger backwards")
+	}
+	if latest.Domain != canary.Domain || latest.ImplantName != canary.ImplantName {
+		t.Error("trigger overwrote canary identity from the caller snapshot")
+	}
+}
+
+func TestTriggerCanaryConcurrentQueries(t *testing.T) {
+	const queries = 12
+	canary := createCanaryStateFixture(t)
+	snapshot := &clientpb.DNSCanary{ID: canary.ID.String()}
+	type result struct {
+		first bool
+		err   error
+	}
+	results := make(chan result, queries)
+	start := make(chan struct{})
+	for range queries {
+		go func() {
+			<-start
+			first, err := TriggerCanary(snapshot)
+			results <- result{first: first, err: err}
+		}()
+	}
+	close(start)
+	firstTriggers := 0
+	for range queries {
+		got := <-results
+		if got.err != nil {
+			t.Errorf("concurrent trigger: %v", got.err)
+		}
+		if got.first {
+			firstTriggers++
+		}
+	}
+	if firstTriggers != 1 {
+		t.Errorf("first triggers = %d, want 1", firstTriggers)
+	}
+	assertCanaryTriggerCount(t, readCanaryState(t, canary.ID), queries)
+}
+
+func TestTriggerCanaryRejectsMissingIdentity(t *testing.T) {
+	canary := createCanaryStateFixture(t)
+	var before int64
+	if err := db.Session().Model(&models.DNSCanary{}).Count(&before).Error; err != nil {
+		t.Fatalf("count canaries: %v", err)
+	}
+	for _, snapshot := range []*clientpb.DNSCanary{
+		nil,
+		{},
+		{ID: "not-a-uuid"},
+		{ID: models.NilUUID().String()},
+		{ID: models.NewUUID().String()},
+	} {
+		first, err := TriggerCanary(snapshot)
+		if err == nil || first {
+			t.Errorf("trigger for %v = (%v, %v), want (false, error)", snapshot, first, err)
+		}
+	}
+	var after int64
+	if err := db.Session().Model(&models.DNSCanary{}).Count(&after).Error; err != nil {
+		t.Fatalf("count canaries: %v", err)
+	}
+	if after != before {
+		t.Errorf("invalid identity inserted canaries: count changed from %d to %d", before, after)
+	}
+	unchanged := readCanaryState(t, canary.ID)
+	if unchanged.Triggered || unchanged.Count != 0 {
+		t.Error("invalid identity changed an existing canary")
+	}
+}
+
+func TestTriggerCanaryRollsBackWriteFailure(t *testing.T) {
+	canary := createCanaryStateFixture(t)
+	injectedErr := errors.New("injected canary timestamp write failure")
+	callbackName := "test:canary-trigger-write-" + canary.ID.String()
+	var rejectWrite atomic.Bool
+	rejectWrite.Store(true)
+	if err := db.Client.Callback().Update().Before("gorm:update").Register(callbackName, func(tx *gorm.DB) {
+		update, ok := tx.Statement.Dest.(models.DNSCanary)
+		if rejectWrite.Load() && ok && !update.LatestTrigger.IsZero() {
+			_ = tx.AddError(injectedErr)
+		}
+	}); err != nil {
+		t.Fatalf("register write failure: %v", err)
+	}
+	t.Cleanup(func() {
+		rejectWrite.Store(false)
+		if err := db.Client.Callback().Update().Remove(callbackName); err != nil {
+			t.Errorf("remove write failure callback: %v", err)
+		}
+	})
+	first, err := TriggerCanary(&clientpb.DNSCanary{ID: canary.ID.String()})
+	if first || !errors.Is(err, injectedErr) {
+		t.Fatalf("failed trigger = (%v, %v), want (false, injected error)", first, err)
+	}
+	unchanged := readCanaryState(t, canary.ID)
+	if unchanged.Triggered || unchanged.Count != 0 || !unchanged.FirstTrigger.IsZero() || !unchanged.LatestTrigger.IsZero() {
+		t.Error("failed trigger did not roll back count and trigger timestamps")
+	}
+	rejectWrite.Store(false)
+	first, err = TriggerCanary(&clientpb.DNSCanary{ID: canary.ID.String()})
+	if err != nil || !first {
+		t.Fatalf("trigger after rollback = (%v, %v), want (true, nil)", first, err)
+	}
+	assertCanaryTriggerCount(t, readCanaryState(t, canary.ID), 1)
+}
+
+func TestGenerateCanaryNormalizesParent(t *testing.T) {
+	for _, parent := range []string{"MiXeD.Example", "MIXED.EXAMPLE.", " \t.MiXeD.Example \t"} {
+		t.Run(parent, func(t *testing.T) {
+			name := "canary_normalization_" + models.NewUUID().String()
+			t.Cleanup(func() {
+				if err := db.Session().Where("implant_name = ?", name).Delete(&models.DNSCanary{}).Error; err != nil {
+					t.Errorf("delete generated canary: %v", err)
+				}
+			})
+			generator := CanaryGenerator{ImplantName: name, ParentDomains: []string{parent}}
+			domain := generator.GenerateCanary()
+			if !strings.HasSuffix(domain, ".mixed.example.") || domain != strings.ToLower(domain) {
+				t.Fatalf("generated domain %q was not normalized", domain)
+			}
+			canary, err := db.CanaryByDomain(domain)
+			if err != nil {
+				t.Fatalf("look up normalized canary: %v", err)
+			}
+			if canary.Domain != domain || canary.ImplantName != name {
+				t.Errorf("stored canary = %v, want domain %s and implant %s", canary, domain, name)
+			}
+		})
+	}
+}

server/generate/profiles_test.go

@@ -1,19 +0,0 @@
-package generate
-
-/*
-	Sliver Implant Framework
-	Copyright (C) 2019  Bishop Fox
-
-	This program is free software: you can redistribute it and/or modify
-	it under the terms of the GNU General Public License as published by
-	the Free Software Foundation, either version 3 of the License, or
-	(at your option) any later version.
-
-	This program is distributed in the hope that it will be useful,
-	but WITHOUT ANY WARRANTY; without even the implied warranty of
-	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-	GNU General Public License for more details.
-
-	You should have received a copy of the GNU General Public License
-	along with this program.  If not, see <https://www.gnu.org/licenses/>.
-*/

server/generate/shellcode_test.go

@@ -1,3 +1,5 @@
+//go:build sliver_generate
+
 package generate
 
 /*

server/loot/loot_test.go

@@ -123,40 +123,6 @@ func TestAllLoot(t *testing.T) {
 	}
 }
 
-func TestAllLootOf(t *testing.T) {
-	lootStore := GetLootStore()
-	_, err := lootStore.Add(&clientpb.Loot{
-		Name:     name1,
-		FileType: clientpb.FileType_BINARY,
-		File: &commonpb.File{
-			Name: name1,
-			Data: data1,
-		},
-	})
-	if err != nil {
-		t.Fatal(err)
-	}
-	_, err = lootStore.Add(&clientpb.Loot{
-		Name:     name2,
-		FileType: clientpb.FileType_TEXT,
-		File: &commonpb.File{
-			Name: name1,
-			Data: []byte("hello world"),
-		},
-	})
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	// Cleanup
-	for _, loot := range lootStore.All().Loot {
-		err = lootStore.Rm(loot.ID)
-		if err != nil {
-			t.Fatal(err)
-		}
-	}
-}
-
 func TestLootErrors(t *testing.T) {
 	lootStore := GetLootStore()
 	loot, err := lootStore.Add(&clientpb.Loot{

server/website/website_test.go

@@ -48,29 +48,6 @@ func randomData() []byte {
 	return buf
 }
 
-func TestAddContent(t *testing.T) {
-	webContent := clientpb.WebContent{
-		Path:        "/data1",
-		ContentType: contentType1,
-		Size:        uint64(len(data1)),
-		Content:     data1,
-	}
-	err := AddContent(website1, &webContent)
-	if err != nil {
-		t.Fatal(err)
-	}
-	webContent2 := clientpb.WebContent{
-		Path:        "/data2",
-		ContentType: contentType2,
-		Size:        uint64(len(data2)),
-		Content:     data1,
-	}
-	err = AddContent(website2, &webContent2)
-	if err != nil {
-		t.Fatal(err)
-	}
-}
-
 func TestGetContent(t *testing.T) {
 
 	webContent := clientpb.WebContent{

util/encoders/images_test.go

@@ -26,43 +26,45 @@ import (
 	"testing"
 )
 
-var (
-	imageTests = []struct {
-		Input []byte
-	}{
-		{[]byte("abc")},   // byte count on image pixel alignment
-		{[]byte("abcde")}, // byte count offset of image pixel alignment
-		{[]byte{0x0, 0x01, 0x02, 0x03, 0x04}},
-		{[]byte{0x01, 0x02, 0x03, 0x04, 0x0}},
-	}
-)
-
 func TestPNG(t *testing.T) {
-	pngEncoder := new(PNGEncoder)
-	for _, test := range imageTests {
-		buf, _ := pngEncoder.Encode(test.Input)
-		decodeOutput, err := pngEncoder.Decode(buf)
-		if err != nil {
-			t.Errorf("png decode returned error: %q", err)
-		}
-		if !bytes.Equal(test.Input, decodeOutput) {
-			t.Errorf("png Decode(img) => %q, expected %q", decodeOutput, test.Input)
-		}
+	allBytes := make([]byte, 256)
+	for i := range allBytes {
+		allBytes[i] = byte(i)
+	}
+	tests := []struct {
+		name    string
+		payload []byte
+	}{
+		{"empty", nil},
+		{"one_zero", []byte{0}},
+		{"one_escape", []byte{1}},
+		{"partial_pixel", []byte{2, 3}},
+		{"pixel_aligned", []byte("abc")},
+		{"pixel_overflow", []byte("abcde")},
+		{"before_square_boundary", bytes.Repeat([]byte{2}, 11)},
+		{"at_square_boundary", bytes.Repeat([]byte{2}, 12)},
+		{"after_square_boundary", bytes.Repeat([]byte{2}, 13)},
+		{"leading_zero", []byte{0, 1, 2, 3, 4}},
+		{"trailing_zero", []byte{1, 2, 3, 4, 0}},
+		{"adjacent_escapes", []byte{0, 1, 0, 1, 1, 0}},
+		{"all_bytes", allBytes},
+		{"one_megabyte", bytes.Repeat(allBytes, 4096)},
 	}
-}
-
-func TestPNGRandomDataRandomSize(t *testing.T) {
 	pngEncoder := new(PNGEncoder)
-	for i := 0; i < 100; i++ {
-		sample := randomDataRandomSize(1024 * 1024)
-		buf, _ := pngEncoder.Encode(sample)
-		decodeOutput, err := pngEncoder.Decode(buf)
-		if err != nil {
-			t.Errorf("png decode returned error: %q", err)
-		}
-		if !bytes.Equal(sample, decodeOutput) {
-			t.Errorf("png Decode(img) => %q, expected %q", decodeOutput, sample)
-		}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			encoded, err := pngEncoder.Encode(test.payload)
+			if err != nil {
+				t.Fatalf("encode: %v", err)
+			}
+			decoded, err := pngEncoder.Decode(encoded)
+			if err != nil {
+				t.Fatalf("decode: %v", err)
+			}
+			if !bytes.Equal(test.payload, decoded) {
+				t.Fatalf("round trip changed the %d-byte payload (decoded %d bytes)", len(test.payload), len(decoded))
+			}
+		})
 	}
 }