1 files changed, 74 insertions, 8 deletions

diff --git a/accessible/windows/msaa/LazyInstantiator.cpp b/accessible/windows/msaa/LazyInstantiator.cpp
index 61814d059a..3ace26d9c0 100644
--- a/accessible/windows/msaa/LazyInstantiator.cpp
+++ b/accessible/windows/msaa/LazyInstantiator.cpp
@@ -12,6 +12,7 @@
 #include "mozilla/a11y/Platform.h"
 #include "mozilla/Assertions.h"
 #include "mozilla/mscom/ProcessRuntime.h"
+#include "mozilla/StaticPrefs_accessibility.h"
 #include "mozilla/UniquePtr.h"
 #include "mozilla/WinHeaderOnlyUtils.h"
 #include "MsaaRootAccessible.h"
@@ -41,9 +42,9 @@ static const wchar_t kLazyInstantiatorProp[] =
 
 Maybe<bool> LazyInstantiator::sShouldBlockUia;
 
-/* static */
-already_AddRefed<IAccessible> LazyInstantiator::GetRootAccessible(HWND aHwnd) {
-  RefPtr<IAccessible> result;
+template <class T>
+already_AddRefed<T> LazyInstantiator::GetRoot(HWND aHwnd) {
+  RefPtr<T> result;
   // At this time we only want to check whether the acc service is running. We
   // don't actually want to create the acc service yet.
   if (!GetAccService()) {
@@ -80,7 +81,7 @@ already_AddRefed<IAccessible> LazyInstantiator::GetRootAccessible(HWND aHwnd) {
   if (!rootAcc->IsRoot()) {
     // rootAcc might represent a popup as opposed to a true root accessible.
     // In that case we just use the regular LocalAccessible::GetNativeInterface.
-    rootAcc->GetNativeInterface(getter_AddRefs(result));
+    result = MsaaAccessible::GetFrom(rootAcc);
     return result.forget();
   }
 
@@ -90,7 +91,7 @@ already_AddRefed<IAccessible> LazyInstantiator::GetRootAccessible(HWND aHwnd) {
   // don't need LazyInstantiator's capabilities anymore (since a11y is already
   // running). We can bypass LazyInstantiator by retrieving the internal
   // unknown (which is not wrapped by the LazyInstantiator) and then querying
-  // that for IID_IAccessible.
+  // that for the interface we want.
   RefPtr<IUnknown> punk(msaaRoot->GetInternalUnknown());
 
   MOZ_ASSERT(punk);
@@ -98,10 +99,24 @@ already_AddRefed<IAccessible> LazyInstantiator::GetRootAccessible(HWND aHwnd) {
     return nullptr;
   }
 
-  punk->QueryInterface(IID_IAccessible, getter_AddRefs(result));
+  punk->QueryInterface(__uuidof(T), getter_AddRefs(result));
   return result.forget();
 }
 
+/* static */
+already_AddRefed<IAccessible> LazyInstantiator::GetRootAccessible(HWND aHwnd) {
+  return GetRoot<IAccessible>(aHwnd);
+}
+
+/* static */
+already_AddRefed<IRawElementProviderSimple> LazyInstantiator::GetRootUia(
+    HWND aHwnd) {
+  if (!StaticPrefs::accessibility_uia_enable()) {
+    return nullptr;
+  }
+  return GetRoot<IRawElementProviderSimple>(aHwnd);
+}
+
 /**
  * When marshaling an interface, COM makes a whole bunch of QueryInterface
  * calls to determine what kind of marshaling the interface supports. We need
@@ -135,7 +150,8 @@ LazyInstantiator::LazyInstantiator(HWND aHwnd)
       mAllowBlindAggregation(false),
       mWeakMsaaRoot(nullptr),
       mWeakAccessible(nullptr),
-      mWeakDispatch(nullptr) {
+      mWeakDispatch(nullptr),
+      mWeakUia(nullptr) {
   MOZ_ASSERT(aHwnd);
   // Assign ourselves as the designated LazyInstantiator for aHwnd
   DebugOnly<BOOL> setPropOk =
@@ -374,9 +390,16 @@ LazyInstantiator::MaybeResolveRoot() {
     if (FAILED(hr)) {
       return hr;
     }
-
     // mWeakAccessible is weak, so don't hold a strong ref
     mWeakAccessible->Release();
+    if (StaticPrefs::accessibility_uia_enable()) {
+      hr = mRealRootUnk->QueryInterface(IID_IRawElementProviderSimple,
+                                        (void**)&mWeakUia);
+      if (FAILED(hr)) {
+        return hr;
+      }
+      mWeakUia->Release();
+    }
 
     // Now that a11y is running, we don't need to remain registered with our
     // HWND anymore.
@@ -401,6 +424,9 @@ IMPL_IUNKNOWN_QUERY_IFACE_AMBIGIOUS(IUnknown, IAccessible)
 IMPL_IUNKNOWN_QUERY_IFACE(IAccessible)
 IMPL_IUNKNOWN_QUERY_IFACE(IDispatch)
 IMPL_IUNKNOWN_QUERY_IFACE(IServiceProvider)
+if (StaticPrefs::accessibility_uia_enable()) {
+  IMPL_IUNKNOWN_QUERY_IFACE(IRawElementProviderSimple)
+}
 // See EnableBlindAggregation for comments.
 if (!mAllowBlindAggregation) {
   return E_NOINTERFACE;
@@ -771,5 +797,45 @@ LazyInstantiator::QueryService(REFGUID aServiceId, REFIID aServiceIid,
   return servProv->QueryService(aServiceId, aServiceIid, aOutInterface);
 }
 
+STDMETHODIMP
+LazyInstantiator::get_ProviderOptions(
+    __RPC__out enum ProviderOptions* aOptions) {
+  // This method is called before a UIA connection is fully established and thus
+  // before we can detect the client. We must not call RESOLVE_ROOT here because
+  // this might turn out to be a client we want to block.
+  if (!aOptions) {
+    return E_INVALIDARG;
+  }
+  *aOptions = uiaRawElmProvider::kProviderOptions;
+  return S_OK;
+}
+
+STDMETHODIMP
+LazyInstantiator::GetPatternProvider(
+    PATTERNID aPatternId, __RPC__deref_out_opt IUnknown** aPatternProvider) {
+  RESOLVE_ROOT;
+  return mWeakUia->GetPatternProvider(aPatternId, aPatternProvider);
+}
+
+STDMETHODIMP
+LazyInstantiator::GetPropertyValue(PROPERTYID aPropertyId,
+                                   __RPC__out VARIANT* aPropertyValue) {
+  RESOLVE_ROOT;
+  return mWeakUia->GetPropertyValue(aPropertyId, aPropertyValue);
+}
+
+STDMETHODIMP
+LazyInstantiator::get_HostRawElementProvider(
+    __RPC__deref_out_opt IRawElementProviderSimple** aRawElmProvider) {
+  // This method is called before a UIA connection is fully established and thus
+  // before we can detect the client. We must not call RESOLVE_ROOT here because
+  // this might turn out to be a client we want to block.
+  if (!aRawElmProvider) {
+    return E_INVALIDARG;
+  }
+  *aRawElmProvider = nullptr;
+  return UiaHostProviderFromHwnd(mHwnd, aRawElmProvider);
+}
+
 }  // namespace a11y
 }  // namespace mozilla