From 6af273b7d6d921b118d6c58a2fb0ed4b0a0b4e92 Mon Sep 17 00:00:00 2001
From: Jesse Peterson <jesse.c.peterson@gmail.com>
Date: Tue, 23 Apr 2024 14:04:38 -0700
Subject: [PATCH] add to main, rename muxer, add comments

---
 cmd/nanomdm/main.go           |  7 ++++++-
 service/nanomdm/token.go      | 28 +++++++++++++++-------------
 service/nanomdm/token_test.go |  2 +-
 3 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/cmd/nanomdm/main.go b/cmd/nanomdm/main.go
index 01fecdc..fccd3d0 100644
--- a/cmd/nanomdm/main.go
+++ b/cmd/nanomdm/main.go
@@ -109,8 +109,13 @@ func main() {
 		stdlog.Fatal(err)
 	}
 
+	tokenMux := nanomdm.NewTokenMux()
+
 	// create 'core' MDM service
-	nanoOpts := []nanomdm.Option{nanomdm.WithLogger(logger.With("service", "nanomdm"))}
+	nanoOpts := []nanomdm.Option{
+		nanomdm.WithGetToken(tokenMux),
+		nanomdm.WithLogger(logger.With("service", "nanomdm")),
+	}
 	if *flDMURLPfx != "" {
 		logger.Debug("msg", "declarative management setup", "url", *flDMURLPfx)
 		dm, err := nanomdm.NewDeclarativeManagementHTTPCaller(*flDMURLPfx, http.DefaultClient)
diff --git a/service/nanomdm/token.go b/service/nanomdm/token.go
index ca1a718..64c0522 100644
--- a/service/nanomdm/token.go
+++ b/service/nanomdm/token.go
@@ -8,32 +8,35 @@ import (
 	"github.com/micromdm/nanomdm/service"
 )
 
+// StaticToken holds static token bytes.
 type StaticToken struct {
 	token []byte
 }
 
+// NewStaticToken creates a new static token handler.
 func NewStaticToken(token []byte) *StaticToken {
 	return &StaticToken{token: token}
 }
 
+// GetToken always responds with the static token bytes.
 func (t *StaticToken) GetToken(_ *mdm.Request, _ *mdm.GetToken) (*mdm.GetTokenResponse, error) {
 	return &mdm.GetTokenResponse{TokenData: t.token}, nil
 }
 
-// TokenServiceTypeMux is a middleware multiplexer for GetToken check-in messages.
+// TokenMux is a middleware multiplexer for GetToken check-in messages.
 // A TokenServiceType string is associated with a GetToken handler and
-// then dispatched appropriately with a matching TokenServiceType.
-type TokenServiceTypeMux struct {
+// then dispatched appropriately.
+type TokenMux struct {
 	typesMu sync.RWMutex
 	types   map[string]service.GetToken
 }
 
-// NewTokenServiceTypeMux creates a new TokenServiceTypeMux.
-func NewTokenServiceTypeMux() *TokenServiceTypeMux { return &TokenServiceTypeMux{} }
+// NewTokenMux creates a new TokenMux.
+func NewTokenMux() *TokenMux { return &TokenMux{} }
 
 // Handle registers a GetToken handler for the given service type.
 // See https://developer.apple.com/documentation/devicemanagement/gettokenrequest
-func (mux *TokenServiceTypeMux) Handle(serviceType string, handler service.GetToken) {
+func (mux *TokenMux) Handle(serviceType string, handler service.GetToken) {
 	if serviceType == "" {
 		panic("tokenmux: invalid service type")
 	}
@@ -51,19 +54,18 @@ func (mux *TokenServiceTypeMux) Handle(serviceType string, handler service.GetTo
 }
 
 // GetToken is the middleware that dispatches a GetToken handler based on service type.
-func (mux *TokenServiceTypeMux) GetToken(r *mdm.Request, t *mdm.GetToken) (*mdm.GetTokenResponse, error) {
-	var next service.GetToken
-	var serviceType string
-	if t != nil {
-		serviceType = t.TokenServiceType
+func (mux *TokenMux) GetToken(r *mdm.Request, t *mdm.GetToken) (*mdm.GetTokenResponse, error) {
+	if t == nil {
+		return nil, fmt.Errorf("nil MDM GetToken")
 	}
+	var next service.GetToken
 	mux.typesMu.RLock()
 	if mux.types != nil {
-		next = mux.types[serviceType]
+		next = mux.types[t.TokenServiceType]
 	}
 	mux.typesMu.RUnlock()
 	if next == nil {
-		return nil, fmt.Errorf("no handler for TokenServiceType: %v", serviceType)
+		return nil, fmt.Errorf("no handler for TokenServiceType: %v", t.TokenServiceType)
 	}
 	return next.GetToken(r, t)
 }
diff --git a/service/nanomdm/token_test.go b/service/nanomdm/token_test.go
index 74ffd8b..94963ec 100644
--- a/service/nanomdm/token_test.go
+++ b/service/nanomdm/token_test.go
@@ -8,7 +8,7 @@ import (
 )
 
 func TestToken(t *testing.T) {
-	m := NewTokenServiceTypeMux()
+	m := NewTokenMux()
 	inTok := []byte("hello")
 	m.Handle("com.apple.maid", NewStaticToken(inTok))
 	inMDMGetToken := &mdm.GetToken{TokenServiceType: "com.apple.maid"}