If unenrolled and k2 provides enrollment secret in localserver request, save it

ee/localserver/krypto-ec-middleware.go

@@ -13,6 +13,7 @@ import (
 	"log/slog"
 	"net/http"
 	"net/url"
+	"os"
 	"runtime"
 	"strings"
 	"sync"
@@ -21,6 +22,7 @@ import (
 	"github.com/kolide/krypto"
 	"github.com/kolide/krypto/pkg/challenge"
 	"github.com/kolide/launcher/ee/agent"
+	"github.com/kolide/launcher/ee/agent/types"
 	"github.com/kolide/launcher/ee/gowrapper"
 	"github.com/kolide/launcher/ee/observability"
 	"github.com/kolide/launcher/ee/presencedetection"
@@ -40,6 +42,7 @@ const (
 	kolideOsHeaderKey                                 = "X-Kolide-Os"
 	kolideArchHeaderKey                               = "X-Kolide-Arch"
 	kolideMunemoHeaderKey                             = "X-Kolide-Munemo"
+	kolideEnrollmentSecretHeaderKey                   = "X-Kolide-Enrollment-Token"
 )
 
 type v2CmdRequestType struct {
@@ -76,26 +79,30 @@ func (cmdReq v2CmdRequestType) CallbackReq() (*http.Request, error) {
 type kryptoEcMiddleware struct {
 	localDbSigner         crypto.Signer
 	counterParty          ecdsa.PublicKey
+	knapsack              types.Knapsack
 	slogger               *slog.Logger
 	presenceDetector      presenceDetector
 	presenceDetectionLock sync.Mutex
 	tenantMunemo          string
+	tenantMunemoMutex     *sync.Mutex
 
 	// presenceDetectionStatusUpdateInterval is the interval at which the presence detection
 	// callback is sent while waiting on user to complete presence detection
 	presenceDetectionStatusUpdateInterval time.Duration
 	timestampValidityRange                int64
 }
 
-func newKryptoEcMiddleware(slogger *slog.Logger, localDbSigner crypto.Signer, counterParty ecdsa.PublicKey, presenceDetector presenceDetector, tenantMunemo string) *kryptoEcMiddleware {
+func newKryptoEcMiddleware(slogger *slog.Logger, localDbSigner crypto.Signer, counterParty ecdsa.PublicKey, knapsack types.Knapsack, presenceDetector presenceDetector, tenantMunemo string) *kryptoEcMiddleware {
 	return &kryptoEcMiddleware{
 		localDbSigner:                         localDbSigner,
 		counterParty:                          counterParty,
+		knapsack:                              knapsack,
 		slogger:                               slogger.With("keytype", "ec"),
 		presenceDetector:                      presenceDetector,
 		timestampValidityRange:                timestampValidityRange,
 		presenceDetectionStatusUpdateInterval: 30 * time.Second,
 		tenantMunemo:                          tenantMunemo,
+		tenantMunemoMutex:                     &sync.Mutex{},
 	}
 }
 
@@ -224,6 +231,10 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 			return
 		}
 
+		// The request is valid -- if we don't have an enrollment secret and k2 has provided one,
+		// then set it!
+		e.setEnrollmentSecretFromHeaders(cmdReq.Headers)
+
 		// set the kolide session id if it exists, this also the saml session id
 		kolideSessionId, ok := cmdReq.CallbackHeaders[kolideSessionIdHeaderKey]
 		if ok && len(kolideSessionId) > 0 {
@@ -750,6 +761,9 @@ func cmdReqToHttpReq(originalRequest *http.Request, cmdReq v2CmdRequestType) *ht
 }
 
 func (e *kryptoEcMiddleware) checkMunemo(headers map[string][]string) error {
+	e.tenantMunemoMutex.Lock()
+	defer e.tenantMunemoMutex.Unlock()
+
 	if e.tenantMunemo == "" {
 		e.slogger.Log(context.TODO(), slog.LevelError,
 			"no munemo set in krypto middleware, continuing",
@@ -774,3 +788,37 @@ func (e *kryptoEcMiddleware) checkMunemo(headers map[string][]string) error {
 
 	return errors.New("munemo in request does not match munemo in middleware")
 }
+
+func (e *kryptoEcMiddleware) setEnrollmentSecretFromHeaders(headers map[string][]string) {
+	e.tenantMunemoMutex.Lock()
+	defer e.tenantMunemoMutex.Unlock()
+
+	if e.tenantMunemo != "" {
+		// Already enrolled, no need to set secret
+		return
+	}
+
+	enrollmentSecretHeaders, ok := headers[kolideEnrollmentSecretHeaderKey]
+	if !ok || len(enrollmentSecretHeaders) == 0 || enrollmentSecretHeaders[0] == "" {
+		e.slogger.Log(context.TODO(), slog.LevelDebug,
+			"no enrollment secret header in request headers, cannot set enrollment secret",
+		)
+		return
+	}
+
+	enrollmentSecretPath := e.knapsack.EnrollSecretPath()
+	if enrollmentSecretPath == "" {
+		e.slogger.Log(context.TODO(), slog.LevelError,
+			"enrollment secret path not set, cannot write enrollment secret to file",
+		)
+		return
+	}
+
+	if err := os.WriteFile(enrollmentSecretPath, []byte(enrollmentSecretHeaders[0]), 0600); err != nil {
+		e.slogger.Log(context.TODO(), slog.LevelError,
+			"could not write enrollment secret to file",
+			"secret_path", enrollmentSecretPath,
+			"err", err,
+		)
+	}
+}

ee/localserver/krypto-ec-middleware_test.go

@@ -23,9 +23,8 @@ import (
 	"github.com/kolide/kit/ulid"
 	"github.com/kolide/krypto/pkg/challenge"
 	"github.com/kolide/krypto/pkg/echelper"
-	"github.com/kolide/launcher/ee/localserver/mocks"
-
 	typesmocks "github.com/kolide/launcher/ee/agent/types/mocks"
+	"github.com/kolide/launcher/ee/localserver/mocks"
 	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/threadsafebuffer"
 	"github.com/stretchr/testify/assert"
@@ -197,6 +196,8 @@ func TestKryptoEcMiddleware(t *testing.T) {
 						Level: slog.LevelDebug,
 					})).Logger
 
+					mockKnapsack := typesmocks.NewKnapsack(t)
+					mockKnapsack.On("EnrollSecretPath").Return("").Maybe()
 					mockPresenceDetector := mocks.NewPresenceDetector(t)
 
 					if runtime.GOOS != "linux" { // only doing persence detection on windows and macos for now
@@ -207,7 +208,7 @@ func TestKryptoEcMiddleware(t *testing.T) {
 					}
 
 					// set up middlewares
-					kryptoEcMiddleware := newKryptoEcMiddleware(slogger, localServerPrivateKey, remoteServerPrivateKey.PublicKey, mockPresenceDetector, "test-munemo")
+					kryptoEcMiddleware := newKryptoEcMiddleware(slogger, localServerPrivateKey, remoteServerPrivateKey.PublicKey, mockKnapsack, mockPresenceDetector, "test-munemo")
 					kryptoEcMiddleware.presenceDetectionStatusUpdateInterval = presenceDetectionCallbackInterval
 
 					rr := httptest.NewRecorder()
@@ -347,8 +348,11 @@ func TestKryptoEcMiddlewareErrors(t *testing.T) {
 					mockPresenceDetector := mocks.NewPresenceDetector(t)
 					mockPresenceDetector.On("DetectPresence", mock.AnythingOfType("string"), mock.AnythingOfType("Duration")).Return(0*time.Second, nil).Maybe()
 
+					mockKnapsack := typesmocks.NewKnapsack(t)
+					mockKnapsack.On("EnrollSecretPath").Return("").Maybe()
+
 					// set up middlewares
-					kryptoEcMiddleware := newKryptoEcMiddleware(slogger, localServerPrivateKey, remoteServerPrivateKey.PublicKey, mockPresenceDetector, "test-munemo")
+					kryptoEcMiddleware := newKryptoEcMiddleware(slogger, localServerPrivateKey, remoteServerPrivateKey.PublicKey, mockKnapsack, mockPresenceDetector, "test-munemo")
 					if tt.middlewareOpt != nil {
 						tt.middlewareOpt(kryptoEcMiddleware)
 					}
@@ -479,8 +483,11 @@ func Test_AllowedOrigin(t *testing.T) {
 			mockPresenceDetector := mocks.NewPresenceDetector(t)
 			mockPresenceDetector.On("DetectPresence", mock.AnythingOfType("string"), mock.AnythingOfType("Duration")).Return(0*time.Second, nil).Maybe()
 
+			mockKnapsack := typesmocks.NewKnapsack(t)
+			mockKnapsack.On("EnrollSecretPath").Return("").Maybe()
+
 			// set up middlewares
-			kryptoEcMiddleware := newKryptoEcMiddleware(slogger, mustGenEcdsaKey(t), counterpartyKey.PublicKey, mockPresenceDetector, "")
+			kryptoEcMiddleware := newKryptoEcMiddleware(slogger, mustGenEcdsaKey(t), counterpartyKey.PublicKey, mockKnapsack, mockPresenceDetector, "")
 
 			h := kryptoEcMiddleware.Wrap(testHandler)
 
@@ -654,6 +661,7 @@ func TestMunemoCheck(t *testing.T) {
 
 			k := typesmocks.NewKnapsack(t)
 			k.On("ReadEnrollSecret").Return(token, nil)
+			k.On("EnrollSecretPath").Return("").Maybe()
 
 			munemo, err := getMunemoFromEnrollSecret(k)
 			if tt.expectMunemoExtractionErr {
@@ -662,7 +670,7 @@ func TestMunemoCheck(t *testing.T) {
 			}
 			require.NoError(t, err)
 
-			e := newKryptoEcMiddleware(multislogger.NewNopLogger(), nil, mustGenEcdsaKey(t).PublicKey, nil, munemo)
+			e := newKryptoEcMiddleware(multislogger.NewNopLogger(), nil, mustGenEcdsaKey(t).PublicKey, k, nil, munemo)
 			err = e.checkMunemo(tt.headers)
 			if tt.expectMiddleWareCheckErr {
 				require.Error(t, err)

ee/localserver/server.go

@@ -89,7 +89,7 @@ func New(ctx context.Context, k types.Knapsack, presenceDetector presenceDetecto
 		)
 	}
 
-	ecKryptoMiddleware := newKryptoEcMiddleware(k.Slogger(), ls.myLocalDbSigner, *ls.serverEcKey, presenceDetector, munemo)
+	ecKryptoMiddleware := newKryptoEcMiddleware(k.Slogger(), ls.myLocalDbSigner, *ls.serverEcKey, k, presenceDetector, munemo)
 	ecAuthedMux := http.NewServeMux()
 	ecAuthedMux.HandleFunc("/", http.NotFound)
 	ecAuthedMux.Handle("/acceleratecontrol", ls.requestAccelerateControlHandler())