MDM proxy for seamless migrations (#19779)

Implementation for the proxy described in #19387.

---------

Co-authored-by: Robert Fairburn <8029478+rfairburn@users.noreply.github.com>

tools/mdm/migration/mdmproxy/Dockerfile

@@ -0,0 +1,17 @@
+FROM golang:1.22.4-alpine3.20@sha256:ace6cc3fe58d0c7b12303c57afe6d6724851152df55e08057b43990b927ad5e8
+ARG TAG
+RUN apk update && apk add --no-cache git
+RUN git clone -b $TAG --depth=1 --no-tags --progress --no-recurse-submodules https://github.com/fleetdm/fleet.git && cd /go/fleet/tools/mdm/migration/mdmproxy && go build .
+
+FROM alpine:3.20.1@sha256:b89d9c93e9ed3597455c90a0b88a8bbb5cb7188438f70953fede212a0c4394e0
+LABEL maintainer="Fleet Developers"
+
+RUN apk update && apk add --no-cache tini
+COPY --from=0 /go/fleet/tools/mdm/migration/mdmproxy/mdmproxy /usr/bin/mdmproxy
+ADD --chmod=0755 ./entrypoint.sh /usr/bin/entrypoint.sh
+
+# Create mdmproxy group and user
+RUN addgroup -S mdmproxy && adduser -S mdmproxy -G mdmproxy
+USER mdmproxy
+
+ENTRYPOINT ["/sbin/tini", "/usr/bin/entrypoint.sh"]

tools/mdm/migration/mdmproxy/README.md

@@ -0,0 +1,28 @@
+Proxy for MDM requests used in seamless migrations, as described in
+https://github.com/fleetdm/fleet/issues/19387.
+
+
+### Usage
+
+```
+Usage of ./mdmproxy:
+  -auth-token string
+        Auth token for remote flag updates (remote updates disabled if not provided)
+  -existing-hostname string
+        Hostname for existing MDM server (eg. 'mdm.example.com') (required)
+  -existing-url string
+        Existing MDM server URL (full path) (required)
+  -fleet-url string
+        Fleet MDM server URL (full path) (required)
+  -migrate-percentage int
+        Percentage of clients to migrate from existing MDM to Fleet
+  -migrate-udids string
+        Space/newline-delimited list of UDIDs to migrate always
+  -server-address string
+        Address for server to listen on (default ":8080")
+```
+
+### Example invocation
+```
+mdmproxy --migrate-udids '' --auth-token foo --existing-url https://3.14.233.249 --existing-hostname micromdm.example.com --fleet-url https://example.cloud.fleetdm.com --migrate-percentage 0
+```

tools/mdm/migration/mdmproxy/entrypoint.sh

@@ -0,0 +1,31 @@
+#!/bin/sh
+
+set -e 
+AUTH_TOKEN_ARG=""
+MIGRATE_PERCENTAGE_ARG=""
+MIGRATE_UDIDS_ARG=""
+
+if [ -z "${MDMPROXY_SERVER_ADDRESS}" ]; then
+	MDMPROXY_SERVER_ADDRESS=":8080"
+fi
+
+if [ -n "${MDMPROXY_AUTH_TOKEN}" ]; then
+	AUTH_TOKEN_ARG="-auth-token \"${MDMPROXY_AUTH_TOKEN:?}\""
+fi
+
+if [ -n "${MDMPROXY_MIGRATE_PERCENTAGE}" ]; then
+	MIGRATE_PERCENTAGE_ARG="-migrate-percentage \"${MDMPROXY_MIGRATE_PERCENTAGE:?}\""
+fi
+
+if [ -n "${MDMPROXY_MIGRATE_UDIDS}" ]; then
+	MIGRATE_UDIDS_ARG="-migrate-udids \"${MDMPROXY_MIGRATE_UDIDS:?}\""
+fi
+
+eval exec /usr/bin/mdmproxy \
+	${AUTH_TOKEN_ARG} \
+	-existing-hostname "${MDMPROXY_EXISTING_HOSTNAME:?}" \
+	-existing-url "${MDMPROXY_EXISTING_URL:?}" \
+	-fleet-url "${MDMPROXY_FLEET_URL:?}"  \
+	${MIGRATE_PERCENTAGE_ARG} \
+	${MIGRATE_UDIDS_ARG} \
+	-server-address "${MDMPROXY_SERVER_ADDRESS:?}"

tools/mdm/migration/mdmproxy/mdmproxy.go

@@ -0,0 +1,322 @@
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"errors"
+	"flag"
+	"fmt"
+	"hash/fnv"
+	"io"
+	"log"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"howett.net/plist"
+)
+
+type mdmProxy struct {
+	migrateUDIDs      map[string]struct{}
+	migratePercentage int
+	existingServerURL string
+	existingHostname  string
+	fleetServerURL    string
+	existingProxy     *httputil.ReverseProxy
+	fleetProxy        *httputil.ReverseProxy
+	// mutex is used to sync reads/updates to the migrateUDIDs and migratePercentage
+	mutex sync.RWMutex
+	// token is used to authenticate updates to the migrateUDIDs and migratePercentage
+	token string
+}
+
+func (m *mdmProxy) handleProxy(w http.ResponseWriter, r *http.Request) {
+	if r.URL.Host != "" {
+		log.Printf("%s %s Forbidden", r.Method, r.URL.String())
+		http.Error(w, "Forbidden", http.StatusForbidden)
+		return
+	}
+
+	// Send all SCEP requests to the existing server
+	if strings.Contains(r.URL.Path, "scep") {
+		log.Printf("%s %s -> Existing (SCEP)", r.Method, r.URL.String())
+		m.existingProxy.ServeHTTP(w, r)
+		return
+	}
+
+	// Send all micromdm API requests to the existing server
+	if strings.HasPrefix(r.URL.Path, "/v1") || strings.HasPrefix(r.URL.Path, "/push") {
+		log.Printf("%s %s -> Existing (API)", r.Method, r.URL.String())
+		m.existingProxy.ServeHTTP(w, r)
+		return
+
+	}
+
+	// Read the body of the request
+	body, err := io.ReadAll(r.Body)
+	_ = r.Body.Close()
+	if err != nil {
+		log.Println("Failed to read request body: ", err.Error())
+		http.Error(w, "Unable to read request body", http.StatusUnprocessableEntity)
+		return
+	}
+	// Reset body so that the reverse proxy request includes it
+	r.Body = io.NopCloser(bytes.NewReader(body))
+
+	// Get the UDID from request
+	udid, err := udidFromRequestBody(body)
+	if err != nil {
+		log.Printf("%s %s Failed to get UDID: %v", r.Method, r.URL.String(), err)
+		http.Error(w, err.Error(), http.StatusUnprocessableEntity)
+		return
+	}
+
+	// Migrated UDIDs go to the Fleet server, otherwise requests go to the existing server.
+	if udid != "" && m.isUDIDMigrated(udid) {
+		log.Printf("%s %s (%s) -> Fleet", r.Method, r.URL.String(), udid)
+		m.fleetProxy.ServeHTTP(w, r)
+	} else {
+		log.Printf("%s %s (%s) -> Existing", r.Method, r.URL.String(), udid)
+		m.existingProxy.ServeHTTP(w, r)
+	}
+}
+
+func (m *mdmProxy) handleUpdatePercentage(w http.ResponseWriter, r *http.Request) {
+	if m.token == "" {
+		http.Error(w, "Set auth token to enable remote updates", http.StatusUnauthorized)
+		return
+	}
+	authHeader := r.Header.Get("Authorization")
+	if authHeader == "" {
+		http.Error(w, "Authorization header must be provided", http.StatusUnauthorized)
+		return
+
+	}
+	if !strings.HasPrefix(authHeader, "Bearer ") {
+		http.Error(w, "Authorization header must start with \"Bearer \"", http.StatusUnauthorized)
+		return
+	}
+	if authHeader != "Bearer "+m.token {
+		http.Error(w, "Authorization header does not match", http.StatusUnauthorized)
+		return
+	}
+
+	body, err := io.ReadAll(r.Body)
+	defer r.Body.Close()
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Failed to read body: %v", err), http.StatusInternalServerError)
+		return
+	}
+	percentage, err := strconv.Atoi(string(body))
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Cannot read body as integer: %v", err), http.StatusUnprocessableEntity)
+		return
+	}
+	if percentage < 0 || percentage > 100 {
+		http.Error(w, "Percentage should be in range (0, 100)", http.StatusBadRequest)
+		return
+	}
+
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.migratePercentage = percentage
+
+	msg := fmt.Sprintf("Migrate percentage updated: %v\n", percentage)
+	log.Printf(msg)
+	fmt.Fprintf(w, msg)
+}
+
+func (m *mdmProxy) handleUpdateMigrateUDIDs(w http.ResponseWriter, r *http.Request) {
+	if m.token == "" {
+		http.Error(w, "Set auth token to enable remote updates", http.StatusUnauthorized)
+		return
+	}
+	authHeader := r.Header.Get("Authorization")
+	if authHeader == "" {
+		http.Error(w, "Authorization header must be provided", http.StatusUnauthorized)
+		return
+
+	}
+	if !strings.HasPrefix(authHeader, "Bearer ") {
+		http.Error(w, "Authorization header must start with \"Bearer \"", http.StatusUnauthorized)
+		return
+	}
+	if authHeader != "Bearer "+m.token {
+		http.Error(w, "Authorization header does not match", http.StatusUnauthorized)
+		return
+	}
+
+	defer r.Body.Close()
+	udids, err := processUDIDs(r.Body)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.migrateUDIDs = udids
+
+	msg := fmt.Sprintf("Migrate UDIDs updated: %v\n", udids)
+	log.Printf(msg)
+	fmt.Fprintf(w, msg)
+}
+
+func processUDIDs(in io.Reader) (map[string]struct{}, error) {
+	scanner := bufio.NewScanner(in)
+	scanner.Split(bufio.ScanWords)
+	udids := make(map[string]struct{})
+	for scanner.Scan() {
+		udids[strings.TrimSpace(scanner.Text())] = struct{}{}
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("Failed to scan UDIDs: %w", err)
+	}
+	return udids, nil
+}
+
+func (m *mdmProxy) isUDIDMigrated(udid string) bool {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+	// If the UDID is manually included, it's always migrated
+	if _, ok := m.migrateUDIDs[udid]; ok {
+		return true
+	}
+
+	// Otherwise migrate by percentage
+	return udidIncludedByPercentage(udid, m.migratePercentage)
+}
+
+func udidFromRequestBody(body []byte) (string, error) {
+	// Not all requests (eg. SCEP) contain a UDID. Return empty without an error in this case.
+	if body == nil || len(body) == 0 {
+		return "", nil
+	}
+
+	type mdmRequest struct {
+		UDID string `plist:""`
+	}
+	var req mdmRequest
+	_, err := plist.Unmarshal(body, &req)
+	if err != nil {
+		return "", fmt.Errorf("unmarshal request: %w body: %s", err, string(body))
+	}
+	if req.UDID == "" {
+		return "", errors.New("request body does not contain UDID")
+	}
+
+	return req.UDID, nil
+}
+
+func hashUDID(udid string) uint {
+	hash := fnv.New32a()
+	hash.Write([]byte(udid))
+	return uint(hash.Sum32())
+}
+
+func udidIncludedByPercentage(udid string, percentage int) bool {
+	index := hashUDID(udid) % 100
+	return int(index) < percentage
+}
+
+func makeExistingProxy(existingURL, existingDNSName string) *httputil.ReverseProxy {
+	targetURL, err := url.Parse(existingURL)
+	if err != nil {
+		panic("failed to parse fleet-url: " + err.Error())
+	}
+	proxy := httputil.NewSingleHostReverseProxy(targetURL)
+
+	// Allow TLS validation to use the "old" server name
+	transport := http.DefaultTransport.(*http.Transport).Clone()
+	transport.TLSClientConfig.ServerName = existingDNSName
+	proxy.Transport = transport
+
+	return proxy
+}
+
+func makeFleetProxy(fleetURL string) *httputil.ReverseProxy {
+	targetURL, err := url.Parse(fleetURL)
+	if err != nil {
+		panic("failed to parse fleet-url: " + err.Error())
+	}
+	proxy := httputil.NewSingleHostReverseProxy(targetURL)
+
+	return proxy
+}
+
+func main() {
+	authToken := flag.String("auth-token", "", "Auth token for remote flag updates (remote updates disabled if not provided)")
+	existingURL := flag.String("existing-url", "", "Existing MDM server URL (full path) (required)")
+	existingHostname := flag.String("existing-hostname", "", "Hostname for existing MDM server (eg. 'mdm.example.com') (required)")
+	fleetURL := flag.String("fleet-url", "", "Fleet MDM server URL (full path) (required)")
+	migratePercentage := flag.Int("migrate-percentage", 0, "Percentage of clients to migrate from existing MDM to Fleet")
+	migrateUDIDs := flag.String("migrate-udids", "", "Space/newline-delimited list of UDIDs to migrate always")
+	serverAddr := flag.String("server-address", ":8080", "Address for server to listen on")
+	flag.Parse()
+
+	// Check required flags
+	if *existingURL == "" {
+		log.Fatal("--existing-url must be set")
+	}
+	if *existingHostname == "" {
+		log.Fatal("--existing-hostname must be set")
+	}
+	if *fleetURL == "" {
+		log.Fatal("--fleet-url must be set")
+	}
+
+	udids, err := processUDIDs(bytes.NewBufferString(*migrateUDIDs))
+	if err != nil {
+		panic(err)
+	}
+	log.Printf("--migrate-udids set: %v", udids)
+	log.Printf("--migrate-percentage set: %d", *migratePercentage)
+	log.Printf("--existing-url set: %s", *existingURL)
+	log.Printf("--existing-hostname set: %s", *existingHostname)
+	log.Printf("--fleet-url set: %s", *fleetURL)
+	if *authToken != "" {
+		log.Printf("--auth-token set. Remote configuration enabled.")
+	} else {
+		log.Printf("--auth-token is empty. Remote configuration disabled.")
+	}
+
+	proxy := mdmProxy{
+		token:             *authToken,
+		existingServerURL: *existingURL,
+		fleetServerURL:    *fleetURL,
+		existingHostname:  *existingHostname,
+		migratePercentage: *migratePercentage,
+		migrateUDIDs:      udids,
+		existingProxy:     makeExistingProxy(*existingURL, *existingHostname),
+		fleetProxy:        makeFleetProxy(*fleetURL),
+	}
+
+	mux := http.NewServeMux()
+	// Health check endpoint used for load balancers
+	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
+		_, err := w.Write([]byte("OK"))
+		if err != nil {
+			log.Printf("/healthz error: %v", err)
+		}
+	})
+	// Remote management of migration (enabled if auth token set)
+	mux.HandleFunc("/admin/udids", proxy.handleUpdateMigrateUDIDs)
+	mux.HandleFunc("/admin/percentage", proxy.handleUpdatePercentage)
+	// Handler for the actual proxying
+	mux.HandleFunc("/", proxy.handleProxy)
+
+	log.Printf("Starting server on %s", *serverAddr)
+	server := &http.Server{
+		Addr:              *serverAddr,
+		ReadHeaderTimeout: 10 * time.Second,
+		Handler:           mux,
+	}
+	err = server.ListenAndServe()
+	if err != nil {
+		fmt.Printf("Error starting server: %s\n", err)
+	}
+}

tools/mdm/migration/mdmproxy/mdmproxy_test.go

@@ -0,0 +1,136 @@
+package main
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_udidFromRequestBody(t *testing.T) {
+	type args struct {
+		body []byte
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    string
+		wantErr bool
+	}{
+		{
+			name: "apple example",
+			// Modified from https://developer.apple.com/documentation/devicemanagement/implementing_device_management/sending_mdm_commands_to_a_device
+			args: args{[]byte(`
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0"> 
+<dict>
+    <key>UDID</key>
+    <string>EFCAF06F-127C-42EA-BF01-E1923A836991</string>
+    <key>CommandUUID</key>
+    <string>9F09D114-BCFD-42AD-A974-371AA7D6256E</string>
+    <key>Status</key>
+    <string>Acknowledged</string> 
+</dict>
+</plist>
+`)},
+			want: "EFCAF06F-127C-42EA-BF01-E1923A836991",
+		},
+		{
+			name: "fleet example",
+			args: args{[]byte(`
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>Status</key>
+	<string>Idle</string>
+	<key>UDID</key>
+	<string>419D33EC-06E6-558D-AD52-601BA1867730</string>
+</dict>
+</plist>
+`)},
+			want: "419D33EC-06E6-558D-AD52-601BA1867730",
+		},
+		{
+			name: "empty request",
+			args: args{[]byte("")},
+			want: "",
+		},
+		{
+			name:    "invalid plist",
+			args:    args{[]byte("<")},
+			wantErr: true,
+		},
+		{
+			name: "plist missing udid",
+			args: args{[]byte(`
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>Status</key>
+	<string>Idle</string>
+</dict>
+</plist>
+`)},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := udidFromRequestBody(tt.args.body)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("udidFromRequestBody() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("udidFromRequestBody() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+// Run the function with the given percentage and return the number of included UDIDs
+func countUdidIncludedByPercentage(percentage, runs int) int {
+	included := 0
+	for i := 0; i < runs; i++ {
+		if udidIncludedByPercentage(uuid.NewString(), percentage) {
+			included++
+		}
+	}
+	return included
+}
+
+func Test_udidIncludedByPercentageNone(t *testing.T) {
+	const percentage int = 0
+	const runs int = 100000
+	included := countUdidIncludedByPercentage(percentage, runs)
+	require.Equal(t, 0, included, "expected no UDIDs to be included")
+}
+
+func Test_udidIncludedByPercentageAll(t *testing.T) {
+	const percentage int = 100
+	const runs int = 100000
+	included := countUdidIncludedByPercentage(percentage, runs)
+	require.Equal(t, runs, included, "expected all UDIDs to be included")
+}
+
+func Test_udidIncludedByPercentage(t *testing.T) {
+	tests := []struct {
+		percentage int
+	}{
+		{1}, {5}, {10}, {25}, {42}, {50}, {75}, {95},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprint(tt.percentage), func(t *testing.T) {
+			const runs int = 100000
+			included := countUdidIncludedByPercentage(tt.percentage, runs)
+			percentage := float64(included) / float64(runs)
+			// Test is nondeterministic, so assert that the actual value is within 1% of the
+			// expected value (to avoid flakiness). In 1000 runs this did not flake once.
+			require.InDelta(t, float64(tt.percentage)/100, percentage, .01)
+		})
+	}
+}