Repository: kolide/launcher

Issue:
Update comshim usage to prevent panics
We use comshim to manage our calls to `ole.CoInitializeEx` in our [WMI queries](https://github.com/kolide/launcher/blob/b06b3a2de3c280a9164bb22b015790dda4d15cc5/ee/wmi/wmi.go#L134) and [windows update table](https://github.com/kolide/launcher/blob/b06b3a2de3c280a9164bb22b015790dda4d15cc5/ee/tables/windowsupdatetable/windowsupdate.go#L99).

Unfortunately there are a few potential avenues where this can introduce a panic, at least when called multiple times without being uninitialized, or when changing thread modes (watchdog calls `ole.CoInitialize(0)`, not multithreaded, which can result in race condition here).

We should fork and update, or somehow avoid the panic [here](https://github.com/scjalliance/comshim/blob/b2ef9d8d7e0fe205a6f89f706c48da27411b9cd9/shim.go#L64-L65).

We should also update watchdog to use the same pattern (or at least use `ole.CoInitializeEx` with `ole.COINIT_MULTITHREADED`) [here](https://github.com/kolide/launcher/blob/main/ee/watchdog/controller_windows.go)

Discussion:

Solution patch:
diff --git a/cmd/launcher/query_windowsupdates_windows.go b/cmd/launcher/query_windowsupdates_windows.go
index dd1cf2be..5565a9d8 100644
--- a/cmd/launcher/query_windowsupdates_windows.go
+++ b/cmd/launcher/query_windowsupdates_windows.go
@@ -10,11 +10,11 @@ import (
 	"os"
 	"strconv"
 
+	comshim "github.com/NozomiNetworks/go-comshim"
 	"github.com/kolide/launcher/ee/tables/windowsupdatetable"
 	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/windows/windowsupdate"
 	"github.com/peterbourgon/ff/v3"
-	"github.com/scjalliance/comshim"
 )
 
 // runQueryWindowsUpdates is a subcommand allowing us to call the Windows Update Agent
@@ -54,7 +54,10 @@ func runQueryWindowsUpdates(systemMultiSlogger *multislogger.MultiSlogger, args
 }
 
 func searchLocale(locale string, tableMode int) ([]byte, string, int, error) {
-	comshim.Add(1)
+	if err := comshim.TryAdd(1); err != nil {
+		comshim.Done() // ensure we decrement the global shim counter that TryAdd increments immediately
+		return nil, "", 0, fmt.Errorf("unable to init comshim: %w", err)
+	}
 	defer comshim.Done()
 
 	searcher, setLocale, isDefaultLocale, err := getSearcher(locale)
diff --git a/ee/localserver/krypto-ec-middleware.go b/ee/localserver/krypto-ec-middleware.go
index d5385a75..11a77d18 100644
--- a/ee/localserver/krypto-ec-middleware.go
+++ b/ee/localserver/krypto-ec-middleware.go
@@ -31,6 +31,8 @@ import (
 
 const (
 	timestampValidityRange                            = 150
+	maxDesiredCallbackQueueSize                       = 10 // we may want to increase this in the future
+	callbackQueueCapacity                             = 100
 	kolideKryptoEccHeader20230130Value                = "2023-01-30"
 	kolideKryptoHeaderKey                             = "X-Kolide-Krypto"
 	kolideSessionIdHeaderKey                          = "X-Kolide-Session"
@@ -81,6 +83,7 @@ type kryptoEcMiddleware struct {
 	presenceDetector      presenceDetector
 	presenceDetectionLock sync.Mutex
 	tenantMunemo          *atomic.String
+	callbackQueue         chan *http.Request
 
 	// presenceDetectionStatusUpdateInterval is the interval at which the presence detection
 	// callback is sent while waiting on user to complete presence detection
@@ -90,6 +93,15 @@ type kryptoEcMiddleware struct {
 
 func newKryptoEcMiddleware(slogger *slog.Logger, localDbSigner crypto.Signer, counterParty ecdsa.PublicKey, presenceDetector presenceDetector, tenantMunemo string) *kryptoEcMiddleware {
 	atomicMunemo := atomic.NewString(tenantMunemo)
+
+	// Set up our callback queue with a worker to send callbacks. The callback queue has a buffer
+	// because we never want to block on sending to the queue.
+	callbackQueue := make(chan *http.Request, callbackQueueCapacity)
+	workerSlogger := slogger.With("subcomponent", "middleware_callback_worker")
+	gowrapper.Go(context.TODO(), workerSlogger, func() {
+		callbackWorker(callbackQueue, workerSlogger)
+	})
+
 	return &kryptoEcMiddleware{
 		localDbSigner:                         localDbSigner,
 		counterParty:                          counterParty,
@@ -98,6 +110,7 @@ func newKryptoEcMiddleware(slogger *slog.Logger, localDbSigner crypto.Signer, co
 		timestampValidityRange:                timestampValidityRange,
 		presenceDetectionStatusUpdateInterval: 30 * time.Second,
 		tenantMunemo:                          atomicMunemo,
+		callbackQueue:                         callbackQueue,
 	}
 }
 
@@ -118,21 +131,70 @@ type callbackDataStruct struct {
 	UserAgent string
 }
 
+// callbackWorker processes requests from the given channel one at a time;
+// this ensures that we don't have multiple callback requests firing simultaneously,
+// to avoid data races during secretless registration.
+func callbackWorker(requests <-chan *http.Request, slogger *slog.Logger) {
+	client := http.Client{
+		Timeout: 8 * time.Second,
+	}
+
+	// Worker loop
+	for req := range requests {
+		// Anonymous function to avoid piling up defers
+		if err := func() error {
+			ctx, cancel := context.WithTimeout(context.Background(), client.Timeout)
+			defer cancel()
+			ctx, span := observability.StartSpan(ctx)
+			defer span.End()
+
+			callbackStart := time.Now()
+			resp, err := client.Do(req.WithContext(ctx))
+			if err != nil {
+				err = fmt.Errorf("sending request: %w", err)
+				observability.SetError(span, err)
+				return err
+			}
+
+			if resp != nil && resp.Body != nil {
+				resp.Body.Close()
+			}
+
+			slogger.Log(req.Context(), slog.LevelDebug,
+				"finished callback",
+				"response_status", resp.Status,
+				"took", time.Since(callbackStart),
+			)
+			return nil
+		}(); err != nil {
+			slogger.Log(context.TODO(), slog.LevelError,
+				"got error in callback",
+				"err", err,
+			)
+		}
+	}
+
+	slogger.Log(context.TODO(), slog.LevelInfo,
+		"callback worker shut down",
+	)
+}
+
 // sendCallback is a command to allow launcher to callback to the SaaS side with krypto responses. As the URL it inside
 // the signed data, and the response is encrypted, this is reasonably secure.
 //
 // Also, because the URL is the box, we cannot cleanly do this through middleware. It reqires a lot of passing data
 // around through context. Doing it here, as part of kryptoEcMiddleware, allows for a fairly succint defer.
 //
-// Note that because this is a network call, it should be called in a goroutine.
-func sendCallback(slogger *slog.Logger, req *http.Request, data *callbackDataStruct) {
+// Because we only want to allow one callback request through at a time, we pass the request off to the callback queue,
+// where the callback worker will pick it up.
+func (e *kryptoEcMiddleware) sendCallback(req *http.Request, data *callbackDataStruct) {
 	if req == nil {
 		return
 	}
 
 	b, err := json.Marshal(data)
 	if err != nil {
-		slogger.Log(req.Context(), slog.LevelError,
+		e.slogger.Log(req.Context(), slog.LevelError,
 			"unable to marshal callback data",
 			"err", err,
 		)
@@ -140,30 +202,15 @@ func sendCallback(slogger *slog.Logger, req *http.Request, data *callbackDataStr
 
 	req.Body = io.NopCloser(bytes.NewReader(b))
 
-	// TODO: This feels like it would be cleaner if we passed in an http client at initialzation time
-	client := http.Client{
-		Timeout: 5 * time.Second,
-	}
-	ctx, cancel := context.WithTimeout(context.Background(), client.Timeout)
-	defer cancel()
-
-	resp, err := client.Do(req.WithContext(ctx))
-	if err != nil {
-		slogger.Log(req.Context(), slog.LevelError,
-			"got error in callback",
-			"err", err,
+	// Check to make sure our queue isn't filling up too rapidly -- drop oldest callback
+	if len(e.callbackQueue) >= maxDesiredCallbackQueueSize {
+		e.slogger.Log(req.Context(), slog.LevelWarn,
+			"callback queue exceeds desired max callback queue size, dropping oldest callback from queue",
+			"queue_len", len(e.callbackQueue),
 		)
-		return
-	}
-
-	if resp != nil && resp.Body != nil {
-		resp.Body.Close()
+		<-e.callbackQueue
 	}
-
-	slogger.Log(req.Context(), slog.LevelDebug,
-		"finished callback",
-		"response_status", resp.Status,
-	)
+	e.callbackQueue <- req
 }
 
 func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
@@ -254,9 +301,8 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 					)
 				}
 
-				gowrapper.Go(r.Context(), e.slogger, func() {
-					sendCallback(e.slogger, callbackReq, callbackData)
-				})
+				// Sending the callback passes the request off to our callback queue, so this is non-blocking.
+				e.sendCallback(callbackReq, callbackData)
 
 				gowrapper.Go(r.Context(), e.slogger, func() {
 					e.detectPresence(challengeBox)
@@ -318,6 +364,11 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 	})
 }
 
+// Close ensures we shut down our callback worker
+func (e *kryptoEcMiddleware) Close() {
+	close(e.callbackQueue)
+}
+
 // extractChallenge finds the challenge in an http request. It prefers the GET parameter, but will fall back to POST data.
 func extractChallenge(r *http.Request) (*challenge.OuterChallenge, error) {
 	// first check query parmeters
@@ -548,7 +599,7 @@ func (e *kryptoEcMiddleware) detectPresence(challengeBox *challenge.OuterChallen
 			return
 		}
 
-		sendCallback(e.slogger, req, callBackData)
+		e.sendCallback(req, callBackData)
 
 		if finalPresenceDetectionResult != nil || hasPresenceDetectionTimedout {
 			return
diff --git a/ee/localserver/server.go b/ee/localserver/server.go
index 36706c59..d30aeb8b 100644
--- a/ee/localserver/server.go
+++ b/ee/localserver/server.go
@@ -12,6 +12,7 @@ import (
 	"net"
 	"net/http"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/kolide/krypto/pkg/echelper"
@@ -48,6 +49,7 @@ type localServer struct {
 	querier                Querier
 	kolideServer           string
 	cancel                 context.CancelFunc
+	interrupted            *atomic.Bool
 
 	myLocalDbSigner crypto.Signer
 	serverEcKey     *ecdsa.PublicKey
@@ -73,6 +75,7 @@ func New(ctx context.Context, k types.Knapsack, presenceDetector presenceDetecto
 		dt4aLimiter:     rate.NewLimiter(defaultRateLimit, defaultRateBurst),
 		kolideServer:    k.KolideServerURL(),
 		myLocalDbSigner: agent.LocalDbKeys(),
+		interrupted:     &atomic.Bool{},
 	}
 
 	// TODO: As there may be things that adjust the keys during runtime, we need to persist that across
@@ -356,6 +359,10 @@ func (ls *localServer) Stop() error {
 }
 
 func (ls *localServer) Interrupt(_ error) {
+	if ls.interrupted.Swap(true) {
+		return
+	}
+
 	ctx := context.TODO()
 
 	ls.slogger.Log(ctx, slog.LevelDebug,
@@ -370,6 +377,8 @@ func (ls *localServer) Interrupt(_ error) {
 	}
 
 	ls.cancel()
+
+	ls.kryptoMiddleware.Close()
 }
 
 func (ls *localServer) startListener() (net.Listener, error) {
diff --git a/ee/tables/windowsupdatetable/windowsupdate.go b/ee/tables/windowsupdatetable/windowsupdate.go
index 71a06f3f..f3f512e9 100644
--- a/ee/tables/windowsupdatetable/windowsupdate.go
+++ b/ee/tables/windowsupdatetable/windowsupdate.go
@@ -14,6 +14,7 @@ import (
 	"strconv"
 	"strings"
 
+	comshim "github.com/NozomiNetworks/go-comshim"
 	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/dataflatten"
 	"github.com/kolide/launcher/ee/observability"
@@ -22,7 +23,6 @@ import (
 	"github.com/kolide/launcher/ee/tables/tablewrapper"
 	"github.com/kolide/launcher/pkg/windows/windowsupdate"
 	"github.com/osquery/osquery-go/plugin/table"
-	"github.com/scjalliance/comshim"
 )
 
 // QueryResults is the data returned by execing `launcher.exe query-windowsupdates`.
@@ -212,7 +212,10 @@ func (t *Table) searchLocale(ctx context.Context, locale string, queryContext ta
 	ctx, span := observability.StartSpan(ctx)
 	defer span.End()
 
-	comshim.Add(1)
+	if err := comshim.TryAdd(1); err != nil {
+		comshim.Done() // ensure we decrement the global shim counter that TryAdd increments immediately
+		return nil, fmt.Errorf("unable to init comshim: %w", err)
+	}
 	defer comshim.Done()
 
 	var results []map[string]string
diff --git a/ee/watchdog/controller_windows.go b/ee/watchdog/controller_windows.go
index 562527d7..22ee62f2 100644
--- a/ee/watchdog/controller_windows.go
+++ b/ee/watchdog/controller_windows.go
@@ -14,6 +14,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/NozomiNetworks/go-comshim"
 	"github.com/go-ole/go-ole"
 	"github.com/go-ole/go-ole/oleutil"
 	"github.com/kolide/launcher/ee/agent/flags/keys"
@@ -225,10 +226,11 @@ func installWatchdogTask(identifier, configFilePath string) error {
 	}
 
 	taskName := launcher.TaskName(identifier, watchdogTaskType)
-	// init COM - we discard the error returned by CoInitialize because it
-	// harmlessly returns S_FALSE if we call it more than once
-	ole.CoInitialize(0)
-	defer ole.CoUninitialize()
+	if err := comshim.TryAdd(1); err != nil {
+		comshim.Done() // ensure we decrement the global shim counter that TryAdd increments immediately
+		return fmt.Errorf("unable to init comshim: %w", err)
+	}
+	defer comshim.Done()
 
 	// create our scheduler object
 	schedService, err := oleutil.CreateObject("Schedule.Service")
@@ -555,10 +557,11 @@ func RemoveWatchdogTask(identifier string) error {
 	}
 
 	taskName := launcher.TaskName(identifier, watchdogTaskType)
-	// init COM - we discard the error returned by CoInitialize because it
-	// harmlessly returns S_FALSE if we call it more than once
-	ole.CoInitialize(0)
-	defer ole.CoUninitialize()
+	if err := comshim.TryAdd(1); err != nil {
+		comshim.Done() // ensure we decrement the global shim counter that TryAdd increments immediately
+		return fmt.Errorf("unable to init comshim: %w", err)
+	}
+	defer comshim.Done()
 
 	// create our scheduler object
 	schedService, err := oleutil.CreateObject("Schedule.Service")
@@ -606,10 +609,11 @@ func watchdogTaskExists(identifier string) (bool, error) { // nolint:unused
 	}
 
 	taskName := launcher.TaskName(identifier, watchdogTaskType)
-	// init COM - we discard the error returned by CoInitialize because it
-	// harmlessly returns S_FALSE if we call it more than once
-	ole.CoInitialize(0)
-	defer ole.CoUninitialize()
+	if err := comshim.TryAdd(1); err != nil {
+		comshim.Done() // ensure we decrement the global shim counter that TryAdd increments immediately
+		return false, fmt.Errorf("unable to init comshim: %w", err)
+	}
+	defer comshim.Done()
 
 	// create our scheduler object
 	schedService, err := oleutil.CreateObject("Schedule.Service")
diff --git a/ee/wmi/wmi.go b/ee/wmi/wmi.go
index c399ea4d..1bbf6bc5 100644
--- a/ee/wmi/wmi.go
+++ b/ee/wmi/wmi.go
@@ -37,9 +37,9 @@ import (
 	"fmt"
 	"log/slog"
 
+	comshim "github.com/NozomiNetworks/go-comshim"
 	"github.com/go-ole/go-ole"
 	"github.com/go-ole/go-ole/oleutil"
-	"github.com/scjalliance/comshim"
 )
 
 const (
@@ -134,7 +134,10 @@ func Query(ctx context.Context, slogger *slog.Logger, className string, properti
 	queryString := fmt.Sprintf("SELECT * FROM %s%s", className, whereClause)
 
 	// Initialize the COM system.
-	comshim.Add(1)
+	if err := comshim.TryAdd(1); err != nil {
+		comshim.Done() // ensure we decrement the global shim counter that TryAdd increments immediately
+		return nil, fmt.Errorf("unable to init comshim: %w", err)
+	}
 	defer comshim.Done()
 
 	unknown, err := oleutil.CreateObject("WbemScripting.SWbemLocator")
diff --git a/go.mod b/go.mod
index 8f2e4ba7..1c0fd9da 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module github.com/kolide/launcher
 require (
 	github.com/Masterminds/semver v1.4.2
 	github.com/Microsoft/go-winio v0.6.2
+	github.com/NozomiNetworks/go-comshim v0.0.0-20241023091934-f8db5c9d85e0
 	github.com/clbanning/mxj v1.8.4
 	github.com/go-ini/ini v1.61.0
 	github.com/go-kit/kit v0.9.0
@@ -21,7 +22,6 @@ require (
 	github.com/osquery/osquery-go v0.0.0-20250131154556-629f995b6947
 	github.com/peterbourgon/ff/v3 v3.1.2
 	github.com/pkg/errors v0.9.1
-	github.com/scjalliance/comshim v0.0.0-20190308082608-cf06d2532c4e
 	github.com/serenize/snaker v0.0.0-20171204205717-a683aaf2d516
 	github.com/stretchr/testify v1.10.0
 	github.com/theupdateframework/go-tuf v0.5.2
diff --git a/go.sum b/go.sum
index 5d43d52f..3a0a869b 100644
--- a/go.sum
+++ b/go.sum
@@ -7,6 +7,8 @@ github.com/Masterminds/semver v1.4.2/go.mod h1:MB6lktGJrhw8PrUyiEoblNEGEQ+RzHPF0
 github.com/Microsoft/go-winio v0.4.9/go.mod h1:VhR8bwka0BXejwEJY73c50VrPtXAaKcyvVC4A4RozmA=
 github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
+github.com/NozomiNetworks/go-comshim v0.0.0-20241023091934-f8db5c9d85e0 h1:WOKXri7votvtRMpp0llnIGdb39eiMsPdT2EJUrT//YY=
+github.com/NozomiNetworks/go-comshim v0.0.0-20241023091934-f8db5c9d85e0/go.mod h1:eGli7SHrXxlN5RGaX3oYaCdB7q3/6+1xLb3EOiPym5c=
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
@@ -227,8 +229,6 @@ github.com/samber/lo v1.38.1 h1:j2XEAqXKb09Am4ebOg31SpvzUTTs6EN3VfgeLUhPdXM=
 github.com/samber/lo v1.38.1/go.mod h1:+m/ZKRl6ClXCE2Lgf3MsQlWfh4bn1bz6CXEOxnEXnEA=
 github.com/samber/slog-multi v1.0.2 h1:6BVH9uHGAsiGkbbtQgAOQJMpKgV8unMrHhhJaw+X1EQ=
 github.com/samber/slog-multi v1.0.2/go.mod h1:uLAvHpGqbYgX4FSL0p1ZwoLuveIAJvBECtE07XmYvFo=
-github.com/scjalliance/comshim v0.0.0-20190308082608-cf06d2532c4e h1:+/AzLkOdIXEPrAQtwAeWOBnPQ0BnYlBW0aCZmSb47u4=
-github.com/scjalliance/comshim v0.0.0-20190308082608-cf06d2532c4e/go.mod h1:9Tc1SKnfACJb9N7cw2eyuI6xzy845G7uZONBsi5uPEA=
 github.com/secure-systems-lab/go-securesystemslib v0.5.0 h1:oTiNu0QnulMQgN/hLK124wJD/r2f9ZhIUuKIeBsCBT8=
 github.com/secure-systems-lab/go-securesystemslib v0.5.0/go.mod h1:uoCqUC0Ap7jrBSEanxT+SdACYJTVplRXWLkGMuDjXqk=
 github.com/serenize/snaker v0.0.0-20171204205717-a683aaf2d516 h1:ofR1ZdrNSkiWcMsRrubK9tb2/SlZVWttAfqUjJi6QYc=



Test patch:
diff --git a/ee/localserver/krypto-ec-middleware_test.go b/ee/localserver/krypto-ec-middleware_test.go
index 0b7b1f31..91de6731 100644
--- a/ee/localserver/krypto-ec-middleware_test.go
+++ b/ee/localserver/krypto-ec-middleware_test.go
@@ -2,6 +2,7 @@ package localserver
 
 import (
 	"bytes"
+	"context"
 	"crypto/ecdsa"
 	"crypto/rand"
 	"encoding/base64"
@@ -16,6 +17,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -686,3 +688,45 @@ func TestMunemoCheck(t *testing.T) {
 		})
 	}
 }
+
+func Test_sendCallback(t *testing.T) {
+	t.Parallel()
+
+	// Set up a test server to receive callback requests
+	requestsReceived := &atomic.Int64{}
+	testCallbackServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestsReceived.Add(1)
+		w.Write([]byte("{}"))
+	}))
+
+	// Make sure we close the server at the end of our test
+	t.Cleanup(func() {
+		testCallbackServer.Close()
+	})
+
+	var logBytes threadsafebuffer.ThreadSafeBuffer
+	slogger := slog.New(slog.NewTextHandler(&logBytes, &slog.HandlerOptions{
+		AddSource: true,
+		Level:     slog.LevelDebug,
+	}))
+
+	requestsQueued := &atomic.Int64{}
+	mw := newKryptoEcMiddleware(slogger, nil, mustGenEcdsaKey(t).PublicKey, nil, "test-munemo")
+	for range callbackQueueCapacity {
+		go func() {
+			req, err := http.NewRequestWithContext(context.TODO(), http.MethodPost, testCallbackServer.URL, nil)
+			require.NoError(t, err)
+			mw.sendCallback(req, &callbackDataStruct{})
+			requestsQueued.Add(1)
+		}()
+	}
+
+	// Wait a little bit to give the requests a chance to enqueue
+	time.Sleep(5 * time.Second)
+
+	// We should have been able to add all requests to the queue
+	require.Equal(t, callbackQueueCapacity, int(requestsQueued.Load()), "could not add all requests to queue; logs: ", logBytes.String())
+
+	// We should have sent at least some of them
+	require.GreaterOrEqual(t, int(requestsReceived.Load()), maxDesiredCallbackQueueSize, "queue worker did not process expected number of requests; logs: ", logBytes.String())
+}
diff --git a/ee/localserver/server_test.go b/ee/localserver/server_test.go
index b2ffcbdd..781135a5 100644
--- a/ee/localserver/server_test.go
+++ b/ee/localserver/server_test.go
@@ -25,12 +25,7 @@ func TestInterrupt_Multiple(t *testing.T) {
 		Level: slog.LevelDebug,
 	}))
 	k.On("Slogger").Return(slogger)
-	k.On("Registrations").Return([]types.Registration{
-		{
-			RegistrationID: types.DefaultRegistrationID,
-			Munemo:         "test-munemo",
-		},
-	}, nil)
+	k.On("Registrations").Return([]types.Registration{}, nil) // return empty set of registrations so we will get a munemo worker
 
 	// Override the poll and recalculate interval for the test so we can be sure that the async workers
 	// do run, but then stop running on shutdown
@@ -81,6 +76,14 @@ func TestInterrupt_Multiple(t *testing.T) {
 
 	require.Equal(t, expectedInterrupts, receivedInterrupts)
 
+	// Confirm all workers shut down. Checking against logs can be a little flaky in CI, so we sleep a couple seconds first
+	// just to be safe.
+	time.Sleep(3 * time.Second)
+	logs := logBytes.String()
+	require.Contains(t, logs, "runAsyncdWorkers received shutdown signal", "id fields worker did not shut down")
+	require.Contains(t, logs, "getMunemoFromKnapsack received shutdown signal", "munemo worker did not shut down")
+	require.Contains(t, logs, "callback worker shut down", "middleware callback worker did not shut down")
+
 	k.AssertExpectations(t)
 	querier.AssertExpectations(t)
 }