Add support for TryBlock and CatchBlock in UiaOperationAbstraction wr…

…apper (#73)

This PR add the Try/Catch helpers in UiaOperationAbstraction wrappers. It also add a helper for GetCurrentFailureCode.

src/UIAutomation/FunctionalTests/UiaOperationAbstractionTests.cpp

@@ -2148,5 +2148,97 @@ namespace UiaOperationAbstractionTests
         {
             ElementAsVariantTest(true);
         }
+
+        // Tests the try block and catch block implementation both locally and remote. When an exception is encountered, the execution
+        // continues till the end of instructions without returning from the exception location, also records the exception
+        // in catch block.
+        void TryAndCatchBlockTest(const bool useRemoteOperations)
+        {
+            ModernApp app(L"Microsoft.WindowsCalculator_8wekyb3d8bbwe!App");
+            app.Activate();
+            auto calc = WaitForElementFocus(L"Display is 0");
+            UiaElement element = calc;
+
+            auto guard = InitializeUiaOperationAbstraction(useRemoteOperations);
+
+            auto scope = UiaOperationScope::StartNew();
+
+            scope.BindInput(element);
+
+            UiaInt failureCode = 0;
+            UiaArray<UiaInt> numbers;
+            numbers.Append(1);
+            numbers.Append(2);
+            numbers.Append(3);
+            scope.TryCatch([&]()
+                {
+                    auto number = numbers.GetAt(3);
+                },
+                [&](UiaFailure failure)
+                {
+                    failureCode = failure.GetCurrentFailureCode();
+                });
+
+            UiaInt numberAfterTryCatch = numbers.GetAt(2);
+            scope.BindResult(failureCode);
+            scope.BindResult(numberAfterTryCatch);
+            scope.Resolve();
+
+            auto hresult = static_cast<HRESULT>(failureCode);
+            Assert::AreEqual(E_BOUNDS, hresult);
+            Assert::AreEqual(3, static_cast<int>(numberAfterTryCatch));
+        }
+
+        TEST_METHOD(TestTryAndCatchBlockAsRemoteTest)
+        {
+            TryAndCatchBlockTest(true);
+        }
+
+        TEST_METHOD(TestTryAndCatchBlockAsLocalTest)
+        {
+            TryAndCatchBlockTest(false);
+        }
+
+        // Tests the try block implementation both locally and remote. When an exception is encountered, the execution
+       // continues till the end of instructions without returning from the exception location.
+        void TryBlockTest(const bool useRemoteOperations)
+        {
+            ModernApp app(L"Microsoft.WindowsCalculator_8wekyb3d8bbwe!App");
+            app.Activate();
+            auto calc = WaitForElementFocus(L"Display is 0");
+            UiaElement element = calc;
+
+            auto guard = InitializeUiaOperationAbstraction(useRemoteOperations);
+
+            auto scope = UiaOperationScope::StartNew();
+
+            scope.BindInput(element);
+
+            UiaInt failureCode = 0;
+            UiaArray<UiaInt> numbers;
+            numbers.Append(1);
+            numbers.Append(2);
+            numbers.Append(3);
+            scope.Try([&]()
+                {
+                    auto number = numbers.GetAt(3);
+                });
+
+            UiaString testString{ L"Text after the try block" };
+            scope.BindResult(testString);
+            scope.Resolve();
+
+            Assert::AreEqual(std::wstring(L"Text after the try block"), testString.GetLocalWstring());
+        }
+
+        TEST_METHOD(TestTryBlockAsRemoteTest)
+        {
+            TryBlockTest(true);
+        }
+
+        TEST_METHOD(TestTryBlockAsLocalTest)
+        {
+            TryBlockTest(false);
+        }
     };
-}
+}

src/UIAutomation/FunctionalTests/pch.h

@@ -18,6 +18,7 @@
 #include <string>
 #include <variant>
 
+#include <wil/cppwinrt.h>
 #include <wil/resource.h>
 #include <wil/result.h>
 #include <wil/com.h>

src/UIAutomation/UiaOperationAbstraction/UiaOperationAbstraction.cpp

@@ -760,6 +760,21 @@ namespace UiaOperationAbstraction
         g_automation.get().reset();
     }
 
+    // UiaFailure
+    UiaInt UiaFailure::GetCurrentFailureCode()
+    {
+        if (m_useRemoteApi)
+        {
+            return m_remoteOperation.GetCurrentFailureCode();
+        }
+
+        // wil::ResultFromCaughtException needs to be used only inside the catch block.
+        // Usually it rethrows the error in the catch block and using it outside the catch
+        // block could crash the process.
+        return wil::ResultFromCaughtException();
+    }
+
+    // UiaOperationDelegator
     UiaOperationDelegator::UiaOperationDelegator() :
         UiaOperationDelegator(g_useRemoteOperations)
     {

src/UIAutomation/UiaOperationAbstraction/UiaOperationAbstraction.h

@@ -71,6 +71,38 @@ namespace UiaOperationAbstraction
         const char* what() const override { return "RemoteOperationsWrapper local loop break"; }
     };
 
+    class UiaFailure
+    {
+    public:
+        // Make UiaOperationDelegator a friend class for UiaFailure and make the UiaFailure
+        // constructor private so that only UiaOperationDelegator can create instances for
+        // UiaFailure in try/catch. This is necessary as UiaFailure contains GetCurrentFailureCode() which
+        // returns the last failure code in both remote/local cases. The local case uses 
+        // wil::ResultFromCaughtException which should be invoked only inside catch block as it
+        // rethows the exception and when called outside the catch block can crash the application.
+        friend class UiaOperationDelegator;
+
+        // Delete copy/move constructors.
+        UiaFailure(const UiaFailure&) = delete;
+        UiaFailure& operator=(const UiaFailure&) = delete;
+
+        UiaFailure(UiaFailure&&) = delete;
+        UiaFailure& operator=(UiaFailure&&) = delete;
+
+        ~UiaFailure() = default;
+
+        UiaInt GetCurrentFailureCode();
+
+    private:
+        // Private constructor used only by the TryCatch/Try of UiaOperationDelegator.
+        UiaFailure(const winrt::Microsoft::UI::UIAutomation::AutomationRemoteOperation remoteOperation, const bool useRemoteApi) :
+            m_remoteOperation(remoteOperation),
+            m_useRemoteApi(useRemoteApi) {}
+
+        const bool m_useRemoteApi;
+        const winrt::Microsoft::UI::UIAutomation::AutomationRemoteOperation m_remoteOperation;
+    };
+
     // The UiaOperationDelegator class in is charge of delegating calls on UIA wrapper types in this abstraction
     // to either the UIA remote operations API or the non-remote UIA API. This class also exposes basic constructs
     // such as if-statements and loops which operate on UIA abstraction wrapper types which should be used by
@@ -123,6 +155,44 @@ namespace UiaOperationAbstraction
             }
         }
 
+        template<class TryBody, class CatchBody>
+        void TryCatch(TryBody&& tryBody, CatchBody&& catchBody) const
+        {
+            if (m_useRemoteApi)
+            {
+                // A new catch body is defined to hold the caller sent catch body as the caller
+                // sent catch body takes UiaFailure as the input param which can be used to get the
+                // failure code on any failure. For local execution, GetCurrentFailureCode calls
+                // wil::ResultFromCaughtException and this needs to be called only in catch block
+                // as it rethrows the exception. Calling it outside the catch block can crash the process.
+                // Having UiaFailure glued with CatchBody can prevent the GetCurrentFailureCode to be accessed
+                // outside catch block
+                auto newCatchBody = [this, catchBody(std::forward<CatchBody>(catchBody))]()
+                {
+                    catchBody(UiaFailure(m_remoteOperation, m_useRemoteApi));
+                };
+
+                m_remoteOperation.TryBlock(std::forward<TryBody>(tryBody), std::move(newCatchBody));
+            }
+            else
+            {
+                try
+                {
+                    tryBody();
+                }
+                catch(...)
+                {
+                    catchBody(UiaFailure(m_remoteOperation, m_useRemoteApi));
+                }
+            }
+        }
+
+        template<class TryBody>
+        void Try(TryBody&& tryBody) const
+        {
+            TryCatch(std::forward<TryBody>(tryBody), [](UiaFailure /*failure*/) {});
+        }
+
         // This method handles a pure lvalue conditional.
         // This method MUST NOT be modified to take a const lvalue ref or a non-ref value, because those
         // would allow an expression as the first argument.
@@ -2372,6 +2442,18 @@ namespace UiaOperationAbstraction
             GetCurrentDelegator()->If<OnTrue>(conditionBool, std::forward<OnTrue>(onTrue));
         }
 
+        template<class TryBody, class CatchBody>
+        inline void TryCatch(TryBody&& tryBody, CatchBody&& catchBody) const
+        {
+            GetCurrentDelegator()->TryCatch<TryBody, CatchBody>(std::forward<TryBody>(tryBody), std::forward<CatchBody>(catchBody));
+        }
+
+        template<class TryBody>
+        inline void Try(TryBody&& tryBody) const
+        {
+            GetCurrentDelegator()->Try<TryBody>(std::forward<TryBody>(tryBody));
+        }
+
         // This method handles a pure lvalue conditional.
         // This method MUST NOT be modified to take a const lvalue ref or a non-ref value, because those
         // would allow an expression as the first argument.

src/UIAutomation/UiaOperationAbstraction/pch.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <wil/cppwinrt.h>
 #include <unknwn.h>
 #include <Windows.h>
 #include <UIAutomation.h>