CLI: Only ask for admin permissions once when changing domain

cli/commands/site/set-domain.ts

@@ -10,7 +10,7 @@ import {
 	unlockAppdata,
 	updateSiteLatestCliPid,
 } from 'cli/lib/appdata';
-import { removeDomainFromHosts } from 'cli/lib/hosts-file';
+import { updateDomainInHosts } from 'cli/lib/hosts-file';
 import { connect, disconnect } from 'cli/lib/pm2-manager';
 import { setupCustomDomain } from 'cli/lib/site-utils';
 import {
@@ -59,14 +59,9 @@ export async function runCommand( sitePath: string, domainName: string ): Promis
 			await unlockAppdata();
 		}
 
-		if ( oldDomainName ) {
-			logger.reportStart(
-				LoggerAction.REMOVE_DOMAIN_FROM_HOSTS,
-				__( 'Removing domain from hosts file…' )
-			);
-			await removeDomainFromHosts( oldDomainName );
-			logger.reportSuccess( __( 'Domain removed from hosts file' ) );
-		}
+		logger.reportStart( LoggerAction.ADD_DOMAIN_TO_HOSTS, __( 'Updating hosts file…' ) );
+		await updateDomainInHosts( oldDomainName, domainName, site.port );
+		logger.reportSuccess( __( 'Hosts file updated' ) );
 
 		logger.reportStart( LoggerAction.START_DAEMON, __( 'Starting process daemon…' ) );
 		await connect();
@@ -76,7 +71,7 @@ export async function runCommand( sitePath: string, domainName: string ): Promis
 
 		if ( runningProcess ) {
 			await stopWordPressServer( site.id );
-			await setupCustomDomain( site, logger );
+			await setupCustomDomain( site, logger, { skipHostsUpdate: true } );
 			logger.reportStart( LoggerAction.START_SITE, __( 'Restarting site…' ) );
 			const processDesc = await startWordPressServer( site, logger );
 			if ( processDesc.pid ) {

cli/commands/site/tests/set-domain.test.ts

@@ -8,7 +8,7 @@ import {
 	unlockAppdata,
 	updateSiteLatestCliPid,
 } from 'cli/lib/appdata';
-import { removeDomainFromHosts } from 'cli/lib/hosts-file';
+import { updateDomainInHosts } from 'cli/lib/hosts-file';
 import { connect, disconnect } from 'cli/lib/pm2-manager';
 import { setupCustomDomain } from 'cli/lib/site-utils';
 import {
@@ -74,7 +74,7 @@ describe( 'CLI: studio site set-domain', () => {
 		( startWordPressServer as jest.Mock ).mockResolvedValue( testProcessDescription );
 		( stopWordPressServer as jest.Mock ).mockResolvedValue( undefined );
 		( arePathsEqual as jest.Mock ).mockImplementation( ( a: string, b: string ) => a === b );
-		( removeDomainFromHosts as jest.Mock ).mockResolvedValue( undefined );
+		( updateDomainInHosts as jest.Mock ).mockResolvedValue( undefined );
 		( setupCustomDomain as jest.Mock ).mockResolvedValue( undefined );
 		( getDomainNameValidationError as jest.Mock ).mockReturnValue( '' );
 		( updateSiteLatestCliPid as jest.Mock ).mockResolvedValue( undefined );
@@ -166,7 +166,11 @@ describe( 'CLI: studio site set-domain', () => {
 			expect( stopWordPressServer ).not.toHaveBeenCalled();
 			expect( startWordPressServer ).not.toHaveBeenCalled();
 			expect( updateSiteLatestCliPid ).not.toHaveBeenCalled();
-			expect( removeDomainFromHosts ).not.toHaveBeenCalled();
+			expect( updateDomainInHosts ).toHaveBeenCalledWith(
+				undefined,
+				testDomainName,
+				testSite.port
+			);
 			expect( disconnect ).toHaveBeenCalled();
 		} );
 
@@ -183,7 +187,9 @@ describe( 'CLI: studio site set-domain', () => {
 
 			expect( isServerRunning ).toHaveBeenCalledWith( testSite.id );
 			expect( stopWordPressServer ).toHaveBeenCalledWith( testSite.id );
-			expect( setupCustomDomain ).toHaveBeenCalledWith( testSite, expect.any( Logger ) );
+			expect( setupCustomDomain ).toHaveBeenCalledWith( testSite, expect.any( Logger ), {
+				skipHostsUpdate: true,
+			} );
 			expect( startWordPressServer ).toHaveBeenCalledWith( testSite, expect.any( Logger ) );
 			expect( updateSiteLatestCliPid ).toHaveBeenCalledWith(
 				testSite.id,
@@ -192,7 +198,7 @@ describe( 'CLI: studio site set-domain', () => {
 			expect( disconnect ).toHaveBeenCalled();
 		} );
 
-		it( 'should remove old domain from hosts file when replacing domain', async () => {
+		it( 'should update hosts file with old and new domain when replacing domain', async () => {
 			const oldDomain = 'old.local';
 			( readAppdata as jest.Mock ).mockResolvedValue( {
 				sites: [ { ...testSite, customDomain: oldDomain } ],
@@ -203,7 +209,11 @@ describe( 'CLI: studio site set-domain', () => {
 
 			await runCommand( testSitePath, testDomainName );
 
-			expect( removeDomainFromHosts ).toHaveBeenCalledWith( oldDomain );
+			expect( updateDomainInHosts ).toHaveBeenCalledWith(
+				oldDomain,
+				testDomainName,
+				testSite.port
+			);
 			expect( disconnect ).toHaveBeenCalled();
 		} );
 	} );

cli/lib/hosts-file.ts

@@ -127,6 +127,51 @@ export const removeDomainFromHosts = async ( domain: string ): Promise< void > =
 	}
 };
 
+/**
+ * Updates a domain in the hosts file by removing the old domain and adding the new one
+ * in a single operation (single admin privilege request).
+ */
+export const updateDomainInHosts = async (
+	oldDomain: string | undefined,
+	newDomain: string | undefined,
+	port: number
+): Promise< void > => {
+	if ( oldDomain === newDomain ) {
+		return;
+	}
+
+	if ( ! oldDomain && newDomain ) {
+		await addDomainToHosts( newDomain, port );
+		return;
+	}
+
+	if ( oldDomain && ! newDomain ) {
+		await removeDomainFromHosts( oldDomain );
+		return;
+	}
+
+	try {
+		const hostsContent = await readHostsFile();
+		const encodedOldDomain = domainToASCII( oldDomain as string );
+		const encodedNewDomain = domainToASCII( newDomain as string );
+		const oldPattern = createHostsEntryPattern( encodedOldDomain );
+		const newContent = updateStudioBlock( hostsContent, ( entries ) => {
+			const filtered = entries.filter( ( entry ) => ! entry.match( oldPattern ) );
+			return [ ...filtered, `127.0.0.1 ${ encodedNewDomain } # Port ${ port }` ];
+		} );
+
+		if ( newContent !== hostsContent ) {
+			await writeHostsFile( newContent );
+		}
+	} catch ( error ) {
+		console.error(
+			`Error replacing domain ${ oldDomain } with ${ newDomain } in hosts file:`,
+			error
+		);
+		throw error;
+	}
+};
+
 /**
  * Helper function for manipulating the "block" of entries in the hosts file
  * pertaining to WordPress Studio.

cli/lib/site-utils.ts

@@ -52,10 +52,13 @@ export function logSiteDetails( site: SiteData ): void {
 /**
  * Sets up custom domain for a site before starting.
  * Handles proxy server startup, SSL certificate generation, and hosts file configuration.
+ *
+ * @param options.skipHostsUpdate - Skip adding domain to hosts file (useful when caller already handled it)
  */
 export async function setupCustomDomain(
 	site: SiteData,
-	logger: Logger< LoggerAction >
+	logger: Logger< LoggerAction >,
+	options?: { skipHostsUpdate?: boolean }
 ): Promise< void > {
 	if ( ! site.customDomain ) {
 		return;
@@ -69,12 +72,14 @@ export async function setupCustomDomain(
 		logger.reportSuccess( __( 'SSL certificates generated' ) );
 	}
 
-	logger.reportStart( LoggerAction.ADD_DOMAIN_TO_HOSTS, __( 'Adding domain to hosts file…' ) );
-	try {
-		await addDomainToHosts( site.customDomain, site.port );
-		logger.reportSuccess( __( 'Domain added to hosts file' ) );
-	} catch ( error ) {
-		throw new LoggerError( __( 'Failed to add domain to hosts file' ), error );
+	if ( ! options?.skipHostsUpdate ) {
+		logger.reportStart( LoggerAction.ADD_DOMAIN_TO_HOSTS, __( 'Adding domain to hosts file…' ) );
+		try {
+			await addDomainToHosts( site.customDomain, site.port );
+			logger.reportSuccess( __( 'Domain added to hosts file' ) );
+		} catch ( error ) {
+			throw new LoggerError( __( 'Failed to add domain to hosts file' ), error );
+		}
 	}
 }
 

src/modules/cli/lib/execute-site-watch-command.ts

@@ -1,6 +1,5 @@
 import { z } from 'zod';
 import { sendIpcEventToRenderer } from 'src/ipc-utils';
-import { CliServerProcess } from 'src/modules/cli/lib/cli-server-process';
 import { executeCliCommand } from 'src/modules/cli/lib/execute-command';
 import { SiteServer } from 'src/site-server';
 import { loadUserData } from 'src/storage/user-data';