use slogger on local server, add span_id id to logs, add kolide sessi…

…on id to logs (#1446)

ee/localserver/krypto-ec-middleware.go

@@ -2,21 +2,22 @@ package localserver
 
 import (
 	"bytes"
+	"context"
 	"crypto"
 	"crypto/ecdsa"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io"
+	"log/slog"
 	"net/http"
 	"net/url"
 	"strings"
 	"time"
 
-	"github.com/go-kit/kit/log"
-	"github.com/go-kit/kit/log/level"
 	"github.com/kolide/krypto"
 	"github.com/kolide/krypto/pkg/challenge"
+	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/traces"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
@@ -26,6 +27,7 @@ const (
 	timestampValidityRange             = 150
 	kolideKryptoEccHeader20230130Value = "2023-01-30"
 	kolideKryptoHeaderKey              = "X-Kolide-Krypto"
+	kolideSessionIdHeaderKey           = "X-Kolide-Session"
 )
 
 type v2CmdRequestType struct {
@@ -60,15 +62,15 @@ func (cmdReq v2CmdRequestType) CallbackReq() (*http.Request, error) {
 type kryptoEcMiddleware struct {
 	localDbSigner, hardwareSigner crypto.Signer
 	counterParty                  ecdsa.PublicKey
-	logger                        log.Logger
+	slogger                       *slog.Logger
 }
 
-func newKryptoEcMiddleware(logger log.Logger, localDbSigner, hardwareSigner crypto.Signer, counterParty ecdsa.PublicKey) *kryptoEcMiddleware {
+func newKryptoEcMiddleware(slogger *slog.Logger, localDbSigner, hardwareSigner crypto.Signer, counterParty ecdsa.PublicKey) *kryptoEcMiddleware {
 	return &kryptoEcMiddleware{
 		localDbSigner:  localDbSigner,
 		hardwareSigner: hardwareSigner,
 		counterParty:   counterParty,
-		logger:         log.With(logger, "keytype", "ec"),
+		slogger:        slogger.With("keytype", "ec"),
 	}
 }
 
@@ -99,9 +101,13 @@ func (e *kryptoEcMiddleware) sendCallback(req *http.Request, data *callbackDataS
 	if req == nil {
 		return
 	}
+
 	b, err := json.Marshal(data)
 	if err != nil {
-		level.Debug(e.logger).Log("msg", "unable to marshal callback data", "err", err)
+		e.slogger.Log(req.Context(), slog.LevelError,
+			"unable to marshal callback data",
+			"err", err,
+		)
 	}
 
 	req.Body = io.NopCloser(bytes.NewReader(b))
@@ -113,20 +119,28 @@ func (e *kryptoEcMiddleware) sendCallback(req *http.Request, data *callbackDataS
 
 	resp, err := client.Do(req)
 	if err != nil {
-		level.Debug(e.logger).Log("msg", "got error in callback", "err", err)
+		e.slogger.Log(req.Context(), slog.LevelError,
+			"got error in callback",
+			"err", err,
+		)
 		return
 	}
 
 	if resp != nil && resp.Body != nil {
 		resp.Body.Close()
 	}
 
-	level.Debug(e.logger).Log("msg", "Finished callback", "response-status", resp.Status)
+	e.slogger.Log(req.Context(), slog.LevelDebug,
+		"finished callback",
+		"response_status", resp.Status,
+	)
 }
 
 func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx, span := traces.StartSpan(r.Context())
+		spanCtx, span := traces.StartSpan(r.Context())
+		r = r.WithContext(spanCtx)
+
 		defer span.End()
 
 		if r.Body != nil {
@@ -136,14 +150,24 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 		challengeBox, err := extractChallenge(r)
 		if err != nil {
 			traces.SetError(span, err)
-			level.Debug(e.logger).Log("msg", "failed to extract box from request", "err", err, "path", r.URL.Path, "query_params", r.URL.RawQuery)
+			e.slogger.Log(r.Context(), slog.LevelDebug,
+				"failed to extract box from request",
+				"err", err,
+				"path", r.URL.Path,
+				"query_params", r.URL.RawQuery,
+			)
+
 			w.WriteHeader(http.StatusUnauthorized)
 			return
 		}
 
 		if err := challengeBox.Verify(e.counterParty); err != nil {
 			traces.SetError(span, err)
-			level.Debug(e.logger).Log("msg", "unable to verify signature", "err", err)
+			e.slogger.Log(r.Context(), slog.LevelDebug,
+				"unable to verify signature",
+				"err", err,
+			)
+
 			w.WriteHeader(http.StatusUnauthorized)
 			return
 		}
@@ -153,20 +177,44 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 		var cmdReq v2CmdRequestType
 		if err := json.Unmarshal(challengeBox.RequestData(), &cmdReq); err != nil {
 			traces.SetError(span, err)
-			level.Debug(e.logger).Log("msg", "unable to unmarshal cmd request", "err", err)
+			e.slogger.Log(r.Context(), slog.LevelError,
+				"unable to unmarshal cmd request",
+				"err", err,
+			)
+
 			w.WriteHeader(http.StatusUnauthorized)
 			return
 		}
 
+		// set the kolide session id if it exists, this also the saml session id
+		kolideSessionId, ok := cmdReq.CallbackHeaders[kolideSessionIdHeaderKey]
+		if ok && len(kolideSessionId) > 0 {
+			span.SetAttributes(attribute.String(kolideSessionIdHeaderKey, kolideSessionId[0]))
+			r = r.WithContext(context.WithValue(r.Context(), multislogger.KolideSessionIdKey, kolideSessionId[0]))
+		}
+
 		// Setup callback URLs and data. This is a pointer, so it can be adjusted before the defer triggers
 		callbackData := &callbackDataStruct{
 			Time:      time.Now().Unix(),
 			UserAgent: r.Header.Get("User-Agent"),
 		}
+
 		if callbackReq, err := cmdReq.CallbackReq(); err != nil {
-			level.Debug(e.logger).Log("msg", "unable to create callback req", "err", err)
+			e.slogger.Log(r.Context(), slog.LevelError,
+				"unable to create callback req",
+				"err", err,
+			)
 		} else if callbackReq != nil {
-			defer func() { go e.sendCallback(callbackReq, callbackData) }()
+			defer func() {
+				if kolideSessionId != nil && len(kolideSessionId) > 0 {
+					callbackReq = callbackReq.WithContext(
+						// adding the current request context will cause this to be cancelled before it sends
+						// so just add the session id manually
+						context.WithValue(callbackReq.Context(), multislogger.KolideSessionIdKey, kolideSessionId[0]),
+					)
+				}
+				go e.sendCallback(callbackReq, callbackData)
+			}()
 		}
 
 		// Check the timestamp, this prevents people from saving a challenge and then
@@ -175,8 +223,11 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 		if timestampDelta > timestampValidityRange || timestampDelta < -timestampValidityRange {
 			span.SetAttributes(attribute.Int64("timestamp_delta", timestampDelta))
 			span.SetStatus(codes.Error, "timestamp is out of range")
+			e.slogger.Log(r.Context(), slog.LevelError,
+				"timestamp is out of range",
+				"delta", timestampDelta,
+			)
 
-			level.Debug(e.logger).Log("msg", "timestamp is out of range", "delta", timestampDelta)
 			w.WriteHeader(http.StatusUnauthorized)
 			callbackData.Error = timeOutOfRangeErr
 			return
@@ -190,14 +241,19 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 				Path:   cmdReq.Path,
 			},
 		}
-		newReq = newReq.WithContext(ctx) // allows the trace to continue to the inner request
+
+		// setting the newReq context to the current request context
+		// allows the trace to continue to the inner request,
+		// maintains the same lifetime as the original request,
+		// allows same ctx values such as session id to be passed to the inner request
+		newReq = newReq.WithContext(r.Context())
 
 		// the body of the cmdReq become the body of the next http request
 		if cmdReq.Body != nil && len(cmdReq.Body) > 0 {
 			newReq.Body = io.NopCloser(bytes.NewBuffer(cmdReq.Body))
 		}
 
-		level.Debug(e.logger).Log("msg", "Successful challenge. Proxying")
+		e.slogger.Log(r.Context(), slog.LevelDebug, "successful challenge, proxying")
 		span.AddEvent("challenge_success")
 
 		// bhr contains the data returned by the request defined above
@@ -216,7 +272,10 @@ func (e *kryptoEcMiddleware) Wrap(next http.Handler) http.Handler {
 
 		if err != nil {
 			traces.SetError(span, err)
-			level.Debug(e.logger).Log("msg", "failed to respond", "err", err)
+			e.slogger.Log(r.Context(), slog.LevelError,
+				"failed to respond",
+				"err", err,
+			)
 			w.WriteHeader(http.StatusUnauthorized)
 			callbackData.Error = responseFailureErr
 			return

ee/localserver/krypto-ec-middleware_test.go

@@ -8,6 +8,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"io"
+	"log/slog"
 	"math/big"
 	"net/http"
 	"net/http/httptest"
@@ -16,11 +17,11 @@ import (
 	"testing"
 	"time"
 
-	"github.com/go-kit/kit/log"
 	"github.com/kolide/kit/ulid"
 	"github.com/kolide/krypto/pkg/challenge"
 	"github.com/kolide/krypto/pkg/echelper"
 	"github.com/kolide/launcher/pkg/agent/keys"
+	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -34,11 +35,16 @@ func TestKryptoEcMiddleware(t *testing.T) {
 	challengeId := []byte(ulid.New())
 	challengeData := []byte(ulid.New())
 
+	koldieSessionId := ulid.New()
+	cmdReqCallBackHeaders := map[string][]string{
+		kolideSessionIdHeaderKey: {koldieSessionId},
+	}
 	cmdReqBody := []byte(randomStringWithSqlCharacters(t, 100000))
 
 	cmdReq := mustMarshal(t, v2CmdRequestType{
-		Path: "whatevs",
-		Body: cmdReqBody,
+		Path:            "whatevs",
+		Body:            cmdReqBody,
+		CallbackHeaders: cmdReqCallBackHeaders,
 	})
 
 	var tests = []struct {
@@ -155,9 +161,12 @@ func TestKryptoEcMiddleware(t *testing.T) {
 					t.Parallel()
 
 					var logBytes bytes.Buffer
+					slogger := multislogger.New(slog.NewTextHandler(&logBytes, &slog.HandlerOptions{
+						Level: slog.LevelDebug,
+					})).Logger
 
 					// set up middlewares
-					kryptoEcMiddleware := newKryptoEcMiddleware(log.NewLogfmtLogger(&logBytes), tt.localDbKey, tt.hardwareKey, counterpartyKey.PublicKey)
+					kryptoEcMiddleware := newKryptoEcMiddleware(slogger, tt.localDbKey, tt.hardwareKey, counterpartyKey.PublicKey)
 					require.NoError(t, err)
 
 					// give our middleware with the test handler to the determiner
@@ -172,6 +181,10 @@ func TestKryptoEcMiddleware(t *testing.T) {
 						return
 					}
 
+					require.Contains(t, logBytes.String(), multislogger.KolideSessionIdKey.String())
+					require.Contains(t, logBytes.String(), koldieSessionId)
+					require.Contains(t, logBytes.String(), multislogger.SpanIdKey.String())
+
 					require.Equal(t, http.StatusOK, rr.Code)
 					require.NotEmpty(t, rr.Body.String())
 

ee/localserver/request-controlservice_test.go

@@ -12,6 +12,7 @@ import (
 	storageci "github.com/kolide/launcher/pkg/agent/storage/ci"
 	"github.com/kolide/launcher/pkg/agent/types"
 	"github.com/kolide/launcher/pkg/agent/types/mocks"
+	"github.com/kolide/launcher/pkg/log/multislogger"
 
 	"github.com/stretchr/testify/require"
 )
@@ -23,6 +24,7 @@ func Test_localServer_requestAccelerateControlFunc(t *testing.T) {
 		m := mocks.NewKnapsack(t)
 		m.On("ConfigStore").Return(storageci.NewStore(t, log.NewNopLogger(), storage.ConfigStore.String()))
 		m.On("KolideServerURL").Return("localhost")
+		m.On("Slogger").Return(multislogger.New().Logger)
 		return m
 	}
 
@@ -47,6 +49,7 @@ func Test_localServer_requestAccelerateControlFunc(t *testing.T) {
 				m.On("ConfigStore").Return(storageci.NewStore(t, log.NewNopLogger(), storage.ConfigStore.String()))
 				m.On("KolideServerURL").Return("localhost")
 				m.On("SetControlRequestIntervalOverride", 250*time.Millisecond, 1*time.Second)
+				m.On("Slogger").Return(multislogger.New().Logger)
 				return m
 			},
 		},

ee/localserver/request-id_test.go

@@ -15,6 +15,7 @@ import (
 	storageci "github.com/kolide/launcher/pkg/agent/storage/ci"
 	"github.com/kolide/launcher/pkg/agent/types"
 	typesMocks "github.com/kolide/launcher/pkg/agent/types/mocks"
+	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/osquery"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -26,6 +27,7 @@ func Test_localServer_requestIdHandler(t *testing.T) {
 	mockKnapsack := typesMocks.NewKnapsack(t)
 	mockKnapsack.On("ConfigStore").Return(storageci.NewStore(t, log.NewNopLogger(), storage.ConfigStore.String()))
 	mockKnapsack.On("KolideServerURL").Return("localhost")
+	mockKnapsack.On("Slogger").Return(multislogger.New().Logger)
 
 	var logBytes bytes.Buffer
 	server := testServer(t, mockKnapsack, &logBytes)

ee/localserver/request-query_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/kolide/launcher/pkg/agent/storage"
 	storageci "github.com/kolide/launcher/pkg/agent/storage/ci"
 	typesMocks "github.com/kolide/launcher/pkg/agent/types/mocks"
+	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/osquery/osquery-go/plugin/distributed"
 	"github.com/stretchr/testify/require"
 )
@@ -57,6 +58,7 @@ func Test_localServer_requestQueryHandler(t *testing.T) {
 			mockKnapsack := typesMocks.NewKnapsack(t)
 			mockKnapsack.On("ConfigStore").Return(storageci.NewStore(t, log.NewNopLogger(), storage.ConfigStore.String()))
 			mockKnapsack.On("KolideServerURL").Return("localhost")
+			mockKnapsack.On("Slogger").Return(multislogger.New().Logger)
 
 			//go:generate mockery --name Querier
 			// https://github.com/vektra/mockery <-- cli tool to generate mocks for usage with testify
@@ -222,6 +224,7 @@ func Test_localServer_requestRunScheduledQueryHandler(t *testing.T) {
 			mockKnapsack := typesMocks.NewKnapsack(t)
 			mockKnapsack.On("ConfigStore").Return(storageci.NewStore(t, log.NewNopLogger(), storage.ConfigStore.String()))
 			mockKnapsack.On("KolideServerURL").Return("localhost")
+			mockKnapsack.On("Slogger").Return(multislogger.New().Logger)
 
 			// set up mock querier
 			mockQuerier := mocks.NewQuerier(t)

ee/localserver/server.go

@@ -97,7 +97,7 @@ func New(k types.Knapsack, opts ...LocalServerOption) (*localServer, error) {
 	}
 	ls.myKey = privateKey
 
-	ecKryptoMiddleware := newKryptoEcMiddleware(ls.logger, ls.myLocalDbSigner, ls.myLocalHardwareSigner, *ls.serverEcKey)
+	ecKryptoMiddleware := newKryptoEcMiddleware(k.Slogger(), ls.myLocalDbSigner, ls.myLocalHardwareSigner, *ls.serverEcKey)
 	ecAuthedMux := http.NewServeMux()
 	ecAuthedMux.HandleFunc("/", http.NotFound)
 	ecAuthedMux.Handle("/acceleratecontrol", ls.requestAccelerateControlHandler())

ee/localserver/server_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/kolide/launcher/pkg/agent/storage"
 	storageci "github.com/kolide/launcher/pkg/agent/storage/ci"
 	"github.com/kolide/launcher/pkg/agent/types/mocks"
+	"github.com/kolide/launcher/pkg/log/multislogger"
 	"github.com/kolide/launcher/pkg/osquery"
 	"github.com/stretchr/testify/require"
 )
@@ -23,6 +24,7 @@ func TestInterrupt_Multiple(t *testing.T) {
 	k := mocks.NewKnapsack(t)
 	k.On("KolideServerURL").Return("localserver")
 	k.On("ConfigStore").Return(c)
+	k.On("Slogger").Return(multislogger.New().Logger)
 
 	ls, err := New(k)
 	require.NoError(t, err)