tools/vsymlink: use the win32 api to update the system environment

cmd/tools/vsymlink.v

@@ -1,6 +1,11 @@
 import os
 import v.pref
 
+const (
+	hkey_local_machine = voidptr(0x80000002)
+	hwnd_broadcast = voidptr(0xffff)
+)
+
 fn main(){
 	$if windows {
 		setup_symlink_on_windows()
@@ -34,56 +39,141 @@ fn setup_symlink_on_unix(){
 }
 
 fn setup_symlink_on_windows(){
-	vexe := pref.vexe_path()
-	// NB: Putting $vdir directly into PATH will also result in
-	// make.bat being global, which is NOT what we want.
-	//
-	// Instead, we create a small launcher v.bat, in a new local
-	// folder .symlink/ . That .symlink/ folder can then be put
-	// in PATH without poluting it with anything else - just a
-	// `v` command will be available, simillar to unix.
-	//
-	// Creating a real NTFS symlink to the real executable was also
-	// tried, but then os.real_path( os.executable() ) returns the
-	// path to the symlink, unfortunately, unlike on posix systems
-	// ¯\_(ツ)_/¯
-	vdir := os.real_path(os.dir(vexe))
-	vsymlinkdir := os.join_path(vdir, '.symlink')
-	vsymlinkbat := os.join_path(vsymlinkdir, 'v.bat')
-	os.rmdir_all(vsymlinkdir)
-	os.mkdir_all(vsymlinkdir)
-	os.write_file(vsymlinkbat, '@echo off\n${vexe} %*')
-	if !os.exists( vsymlinkbat ) {
-		eprintln('Could not create $vsymlinkbat')
-		exit(1)
-	}
-	println('Created $vsymlinkbat .')
-	current_paths := os.getenv('PATH').split(';').map(it.trim('/\\'))
-	if vsymlinkdir in current_paths {
-		println('$vsymlinkdir is already on your PATH')
-		println('Try running `v version`')
-		exit(0)
-	}
-	// put vsymlinkdir first, prevent duplicates:
-	mut new_paths := [ vsymlinkdir ]
-	for p in current_paths {
-		if p !in new_paths {
-			new_paths << p
+	$if windows {
+		vexe := pref.vexe_path()
+		// NB: Putting $vdir directly into PATH will also result in
+		// make.bat being global, which is NOT what we want.
+		//
+		// Instead, we create a small launcher v.bat, in a new local
+		// folder .symlink/ . That .symlink/ folder can then be put
+		// in PATH without poluting it with anything else - just a
+		// `v` command will be available, similar to unix.
+		//
+		// Creating a real NTFS symlink to the real executable was also
+		// tried, but then os.real_path( os.executable() ) returns the
+		// path to the symlink, unfortunately, unlike on posix systems
+		// ¯\_(ツ)_/¯
+		vdir := os.real_path(os.dir(vexe))
+		vsymlinkdir := os.join_path(vdir, '.symlink')
+		vsymlinkbat := os.join_path(vsymlinkdir, 'v.bat')
+		if os.exists(vsymlinkbat) {
+			print('Batch script $vsymlinkbat already exists, checking system %PATH%...')
 		}
-	}
-	//
-	change_path_cmd := 'setx /M PATH "' + new_paths.join(';') +'"'
-	println('Changing global PATH with:')
-	println(change_path_cmd)
-	res := os.system(change_path_cmd)
-	if res == 0 {
-		println('')
-		println('$vsymlinkdir has been prepended to PATH.')
-		println('Try running `v version`.')
-		exit(0)
-	} else {
-		println('Could not run `setx`, probably you are not an administrator.')
-		println('`v symlink` should be launched with admin privileges.')
-		exit(1)
+		else {
+			os.rmdir_all(vsymlinkdir)
+			os.mkdir_all(vsymlinkdir)
+			os.write_file(vsymlinkbat, '@echo off\n${vexe} %*')
+			if !os.exists(vsymlinkbat) {
+				eprintln('Could not create $vsymlinkbat')
+				exit(1)
+			}
+			else {
+				print('Created $vsymlinkbat, checking system %PATH%...')
+			}
+		}
+
+		reg_sys_env_handle  := get_reg_sys_env_handle() or {
+			warn_and_exit(err)
+			return
+		}
+		defer {
+			C.RegCloseKey(reg_sys_env_handle)
+		}
+
+		sys_env_path := get_reg_value(reg_sys_env_handle, 'Path') or {
+			warn_and_exit(err)
+			return
+		}
+
+		current_sys_paths := sys_env_path.split(os.path_delimiter).map(it.trim('/$os.path_separator'))
+		mut new_paths := [ vsymlinkdir ]
+		for p in current_sys_paths {
+			if p !in new_paths {
+				new_paths << p
+			}
+		}
+
+		new_sys_env_path := new_paths.join(';')
+
+		if new_sys_env_path == sys_env_path {
+			println('configured.')
+		}
+		else {
+			print('not configured.\nSetting system %PATH%...')
+			set_reg_value(reg_sys_env_handle, 'Path', new_sys_env_path) or {
+				warn_and_exit(err)
+				return
+			}
+			println('done.')
+		}
+
+		print('Letting running process know to update their Environment...')
+		send_setting_change_msg('Environment') or {
+			eprintln('\n' + err)
+			warn_and_exit('You might need to run this again to have `v` in your %PATH%')
+			return
+		}
+
+		println('finished.\n\nNote: restart your shell/IDE to load the new %PATH%.')
+		println('\nAfter restarting your shell/IDE, give `v version` a try in another dir!')
 	}
 }
+
+fn warn_and_exit(err string) {
+	eprintln(err)
+	exit(1)
+}
+
+// get the system environment registry handle
+fn get_reg_sys_env_handle() ?voidptr {
+	$if windows {
+		// open the registry key
+		reg_key_path   := 'SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment'
+		reg_env_key    := voidptr(0) // or HKEY (HANDLE)
+		if C.RegOpenKeyEx(hkey_local_machine, reg_key_path.to_wide(), 0, 1 | 2, &reg_env_key) != 0 {
+			return error('Could not open "$reg_key_path" in the registry')
+		}
+
+		return reg_env_key
+	}
+	return error('not on windows')
+}
+
+// get a value from a given $key
+fn get_reg_value(reg_env_key voidptr, key string) ?string {
+	$if windows {
+		// query the value (shortcut the sizing step)
+		reg_value_size := 4095 // this is the max length (not for the registry, but for the system %PATH%)
+		mut reg_value  := &u16(malloc(reg_value_size))
+		if C.RegQueryValueEx(reg_env_key, key.to_wide(), 0, 0, reg_value, &reg_value_size) != 0 {
+			return error('Unable to get registry value for "$key", are you running as an Administrator?')
+		}
+
+		return string_from_wide(reg_value)
+	}
+	return error('not on windows')
+}
+
+// sets the value for the given $key to the given  $value
+fn set_reg_value(reg_key voidptr, key string, value string) ?bool {
+	$if windows {
+		if C.RegSetValueEx(reg_key, key.to_wide(), 0, 1, value.to_wide(), 4095) != 0 {
+			return error('Unable to set registry value for "$key", are you running as an Administrator?')
+		}
+
+		return true
+	}
+	return error('not on windows')
+}
+
+// broadcasts a message to all listening windows (explorer.exe in particular)
+// letting them know that the system environment has changed and should be reloaded
+fn send_setting_change_msg(message_data string) ?bool {
+	$if windows {
+		if C.SendMessageTimeout(hwnd_broadcast, 0x001A, 0, message_data.to_wide(), 2, 5000, 0) == 0 {
+			return error('Could not broadcast WM_SETTINGCHANGE')
+		}
+		return true
+	}
+	return error('not on windows')
+}