From 29e819330e6a96b7e43a8bc316d71ac3c8945df2 Mon Sep 17 00:00:00 2001
From: Fresnel Imaging Team <noreply@google.com>
Date: Sun, 14 Jun 2026 21:50:33 -0700
Subject: [PATCH] Small refactoring.

PiperOrigin-RevId: 932219546
---
 cli/commands/server/server_windows.go | 77 +++++++++++++++++++--------
 cli/config/config_windows.go          |  3 +-
 2 files changed, 55 insertions(+), 25 deletions(-)

diff --git a/cli/commands/server/server_windows.go b/cli/commands/server/server_windows.go
index 0ed4ff6..780e11c 100644
--- a/cli/commands/server/server_windows.go
+++ b/cli/commands/server/server_windows.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync"
 	"net/http"
 	"net/rpc"
 	"net/rpc/jsonrpc"
@@ -28,8 +29,19 @@ import (
 // PipeName is the name of the named pipe used for RPC communication.
 const PipeName = `\\.\pipe\fresnel_service`
 
+var (
+	// PreWriteDiskHook allows executing custom setup logic prior to writing a disk.
+	PreWriteDiskHook = func() error { return nil }
+	// PostWriteDiskHook allows executing custom teardown logic after writing a disk.
+	PostWriteDiskHook = func() {}
+	// StartupHook allows executing custom logic when the server service starts.
+	StartupHook = func() {}
+)
+
 // FresnelService is the RPC service exposed over the named pipe.
-type FresnelService struct{}
+type FresnelService struct {
+	writeMu sync.Mutex
+}
 
 // WriteRequest represents a request to write an image to disk.
 type WriteRequest struct {
@@ -55,6 +67,15 @@ func (s *FresnelService) WriteDisk(req *WriteRequest, resp *WriteResponse) error
 		resp.Error = "Access denied: Could not securely identify the calling user."
 		return nil
 	}
+	s.writeMu.Lock()
+	defer s.writeMu.Unlock()
+
+	if err := PreWriteDiskHook(); err != nil {
+		resp.Error = fmt.Sprintf("pre-write hook error: %v", err)
+		return nil
+	}
+	defer PostWriteDiskHook()
+
 	conf, err := config.New(req.Cleanup, req.Warning, req.Eject, req.FFU, req.Update, req.Devices, req.Distro, req.Track, req.ConfTrack, req.SeedServer)
 	if err != nil {
 		resp.Error = fmt.Sprintf("config error: %v", err)
@@ -71,14 +92,15 @@ func (s *FresnelService) WriteDisk(req *WriteRequest, resp *WriteResponse) error
 
 	if req.SSOCookie != "" {
 		tlsClient, err := sso.TLSClient(nil, nil)
+		if err != nil {
+			deck.Errorf("TLSClient error: %v", err)
+		}
 		if err == nil {
 			client := &ssoHTTPClient{
 				cookie: req.SSOCookie,
 				client: tlsClient,
 			}
 			i.SetHTTPClient(client)
-		} else {
-			deck.Errorf("TLSClient error: %v", err)
 		}
 	}
 
@@ -144,6 +166,8 @@ func (m *tokenCodec) Close() error {
 func (m *fresnelSvc) Execute(args []string, r <-chan svc.ChangeRequest, changes chan<- svc.Status) (svcSpecificEC bool, exitCode uint32) {
 	changes <- svc.Status{State: svc.StartPending}
 
+	StartupHook()
+
 	service := new(FresnelService)
 	rpc.Register(service)
 
@@ -166,26 +190,7 @@ func (m *fresnelSvc) Execute(args []string, r <-chan svc.ChangeRequest, changes
 			if err != nil {
 				return
 			}
-			var clientToken windows.Token
-
-			// Extract the raw Windows pipe handle using reflection.
-			fd := getPipeHandle(conn)
-			if fd != 0 {
-				runtime.LockOSThread()
-				// Impersonate the pipe client.
-				if err := installer.ImpersonateNamedPipeClient(fd); err == nil {
-					// Grab a copy of their token.
-					if err := windows.OpenThreadToken(windows.CurrentThread(), windows.TOKEN_ALL_ACCESS, true, &clientToken); err != nil {
-						deck.Errorf("Failed to open thread token: %v", err)
-					}
-					windows.RevertToSelf()
-				} else {
-					deck.Errorf("Failed to impersonate pipe client: %v", err)
-				}
-				runtime.UnlockOSThread()
-			} else {
-				deck.Errorf("Could not extract raw pipe handle from connection")
-			}
+			clientToken := getClientToken(conn)
 			// Serve the RPC connection using our injecting codec.
 			codec := &tokenCodec{
 				ServerCodec: jsonrpc.NewServerCodec(conn),
@@ -235,6 +240,32 @@ func getPipeHandle(conn net.Conn) windows.Handle {
 	return 0
 }
 
+// getClientToken impersonates the connected pipe client to retrieve their Windows token.
+func getClientToken(conn net.Conn) windows.Token {
+	fd := getPipeHandle(conn)
+	if fd == 0 {
+		deck.Errorf("Could not extract raw pipe handle from connection")
+		return 0
+	}
+
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+
+	if err := installer.ImpersonateNamedPipeClient(fd); err != nil {
+		deck.Errorf("Failed to impersonate pipe client: %v", err)
+		return 0
+	}
+	defer windows.RevertToSelf()
+
+	var clientToken windows.Token
+	if err := windows.OpenThreadToken(windows.CurrentThread(), windows.TOKEN_ALL_ACCESS, true, &clientToken); err != nil {
+		deck.Errorf("Failed to open thread token: %v", err)
+		return 0
+	}
+
+	return clientToken
+}
+
 type ssoHTTPClient struct {
 	cookie string
 	client *http.Client
diff --git a/cli/config/config_windows.go b/cli/config/config_windows.go
index 4e100fb..2443f30 100644
--- a/cli/config/config_windows.go
+++ b/cli/config/config_windows.go
@@ -29,7 +29,6 @@ var (
 	// IsElevatedCmd injects the command to determine the elevation state of the
 	// user context.
 	IsElevatedCmd      = isAdmin
-	funcUSBPermissions = HasWritePermissions
 
 	denyWriteRegKey = `SOFTWARE\Policies\Microsoft\Windows\RemovableStorageDevices\{53f5630d-b6bf-11d0-94f2-00a0c91efb8b}`
 )
@@ -69,7 +68,7 @@ func isAdmin() (bool, error) {
 }
 
 // HasWritePermissions determines if the local machine is blocked from writing to removable media via policy.
-func HasWritePermissions() error {
+var HasWritePermissions = func() error {
 	v, err := registry.GetInteger(denyWriteRegKey, "Deny_Write")
 	if err != nil && err != registry.ErrNotExist {
 		return err