From 20f7e889c6242a6796988925ab12a344cad8bf2c Mon Sep 17 00:00:00 2001
From: Jesse Peterson <jesse.c.peterson@gmail.com>
Date: Sat, 20 Apr 2024 15:52:50 -0700
Subject: [PATCH] add token muxer and static token helper

---
 service/nanomdm/token.go      | 69 +++++++++++++++++++++++++++++++++++
 service/nanomdm/token_test.go | 28 ++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 service/nanomdm/token.go
 create mode 100644 service/nanomdm/token_test.go

diff --git a/service/nanomdm/token.go b/service/nanomdm/token.go
new file mode 100644
index 0000000..ca1a718
--- /dev/null
+++ b/service/nanomdm/token.go
@@ -0,0 +1,69 @@
+package nanomdm
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/micromdm/nanomdm/mdm"
+	"github.com/micromdm/nanomdm/service"
+)
+
+type StaticToken struct {
+	token []byte
+}
+
+func NewStaticToken(token []byte) *StaticToken {
+	return &StaticToken{token: token}
+}
+
+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.
+// A TokenServiceType string is associated with a GetToken handler and
+// then dispatched appropriately with a matching TokenServiceType.
+type TokenServiceTypeMux struct {
+	typesMu sync.RWMutex
+	types   map[string]service.GetToken
+}
+
+// NewTokenServiceTypeMux creates a new TokenServiceTypeMux.
+func NewTokenServiceTypeMux() *TokenServiceTypeMux { return &TokenServiceTypeMux{} }
+
+// 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) {
+	if serviceType == "" {
+		panic("tokenmux: invalid service type")
+	}
+	if handler == nil {
+		panic("tokenmux: invalid handler")
+	}
+	mux.typesMu.Lock()
+	defer mux.typesMu.Unlock()
+	if mux.types == nil {
+		mux.types = make(map[string]service.GetToken)
+	} else if _, exists := mux.types[serviceType]; exists {
+		panic("tokenmux: multiple registrations for " + serviceType)
+	}
+	mux.types[serviceType] = handler
+}
+
+// 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
+	}
+	mux.typesMu.RLock()
+	if mux.types != nil {
+		next = mux.types[serviceType]
+	}
+	mux.typesMu.RUnlock()
+	if next == nil {
+		return nil, fmt.Errorf("no handler for TokenServiceType: %v", serviceType)
+	}
+	return next.GetToken(r, t)
+}
diff --git a/service/nanomdm/token_test.go b/service/nanomdm/token_test.go
new file mode 100644
index 0000000..74ffd8b
--- /dev/null
+++ b/service/nanomdm/token_test.go
@@ -0,0 +1,28 @@
+package nanomdm
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/micromdm/nanomdm/mdm"
+)
+
+func TestToken(t *testing.T) {
+	m := NewTokenServiceTypeMux()
+	inTok := []byte("hello")
+	m.Handle("com.apple.maid", NewStaticToken(inTok))
+	inMDMGetToken := &mdm.GetToken{TokenServiceType: "com.apple.maid"}
+	outTok, err := m.GetToken(nil, inMDMGetToken)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(inTok, outTok.TokenData) {
+		t.Error("input and output not equal")
+	}
+	// invalid type
+	inMDMGetToken = &mdm.GetToken{TokenServiceType: "com.apple.does-not-exist"}
+	_, err = m.GetToken(nil, inMDMGetToken)
+	if err == nil {
+		t.Fatal("should be an error")
+	}
+}