Bug 1627805 - Obtain path for the download XML r=agashlin

Differential Revision: https://phabricator.services.mozilla.com/D69875

toolkit/components/updateagent/src/update_paths.rs

@@ -2,7 +2,12 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
-use std::{ffi::OsString, path::PathBuf};
+use std::{env::current_exe, ffi::OsString, path::PathBuf};
+use winapi::shared::{
+    minwindef::MAX_PATH,
+    winerror::{FAILED, HRESULT},
+};
+use wio::wide::{FromWide, ToWide};
 
 pub fn get_download_xml() -> Result<OsString, String> {
     let mut update_dir = PathBuf::from(get_update_directory()?);
@@ -18,7 +23,39 @@ pub fn get_install_hash() -> Result<OsString, String> {
         .ok_or("Couldn't get update directory name".to_string())
 }
 
-pub fn get_update_directory() -> Result<OsString, String> {
-    // TODO: Get the real update directory
-    Ok(OsString::from("C:\\test\\"))
+#[link(name = "updatecommon", kind = "static")]
+extern "C" {
+    fn get_common_update_directory(install_path: *const u16, result: *mut u16) -> HRESULT;
+}
+pub fn get_update_directory() -> Result<OsString, String> {
+    let mut binary_path = current_exe()
+        .map_err(|e| format!("Couldn't get executing binary's path: {}", e))?
+        .parent()
+        .ok_or("Couldn't get executing binary's parent directory")?
+        .as_os_str()
+        .to_wide_null();
+    // It looks like Path.as_os_str returns a path without a trailing slash, so this may be
+    // unnecessary. But the documentation does not appear to guarantee the lack of trailing slash.
+    // get_common_update_directory, however, must be given the installation path without a trailing
+    // slash. So remove any trailing path slashes, just to be safe.
+    strip_trailing_path_separators(&mut binary_path);
+
+    let mut buffer: Vec<u16> = Vec::with_capacity(MAX_PATH + 1);
+    let hr = unsafe { get_common_update_directory(binary_path.as_ptr(), buffer.as_mut_ptr()) };
+    if FAILED(hr) {
+        return Err(format!("Failed to get update directory: HRESULT={}", hr));
+    }
+    Ok(unsafe { OsString::from_wide_ptr_null(buffer.as_ptr()) })
+}
+
+fn strip_trailing_path_separators(path: &mut Vec<u16>) {
+    let forward_slash = '/' as u16;
+    let backslash = '\\' as u16;
+    for character in path.iter_mut().rev() {
+        if *character == forward_slash || *character == backslash {
+            *character = 0;
+        } else if *character != 0 {
+            break;
+        }
+    }
 }

toolkit/mozapps/update/common/commonupdatedir.cpp

@@ -897,6 +897,23 @@ GetUserUpdateDirectory(const wchar_t* installPath, const char* vendor,
       result);
 }
 
+/**
+ * This is a much more limited version of the GetCommonUpdateDirectory that can
+ * be called from Rust.
+ * The result parameter must be a valid pointer to a buffer of length
+ * MAX_PATH + 1
+ */
+extern "C" HRESULT get_common_update_directory(const wchar_t* installPath,
+                                               wchar_t* result) {
+  mozilla::UniquePtr<wchar_t[]> uniqueResult;
+  HRESULT hr = GetCommonUpdateDirectory(
+      installPath, SetPermissionsOf::BaseDirIfNotExists, uniqueResult);
+  if (FAILED(hr)) {
+    return hr;
+  }
+  return StringCchCopyW(result, MAX_PATH + 1, uniqueResult.get());
+}
+
 /**
  * This is a helper function that does all of the work for
  * GetCommonUpdateDirectory and GetUserUpdateDirectory. It partially exists to