// Copyright 2024 The Chromium Authors // Use of this source code is governed by a BSD-style license that can be // found in the LICENSE file. #import "ios/chrome/browser/intelligence/proto_wrappers/page_context_wrapper.h" #import #import #import #import #import #import #import "base/apple/foundation_util.h" #import "base/barrier_closure.h" #import "base/base64.h" #import "base/check.h" #import "base/check_op.h" #import "base/feature_list.h" #import "base/functional/bind.h" #import "base/logging.h" #import "base/memory/weak_ptr.h" #import "base/not_fatal_until.h" #import "base/strings/string_util.h" #import "base/strings/sys_string_conversions.h" #import "base/strings/utf_string_conversions.h" #import "base/task/bind_post_task.h" #import "base/task/sequenced_task_runner.h" #import "base/task/task_traits.h" #import "base/task/thread_pool.h" #import "base/time/time.h" #import "base/timer/timer.h" #import "base/token.h" #import "base/values.h" #import "components/autofill/core/common/unique_ids.h" #import "components/autofill/ios/form_util/child_frame_registrar.h" #import "components/optimization_guide/core/page_content_proto_serializer.h" #import "components/optimization_guide/proto/features/common_quality_data.pb.h" #import "components/security_state/core/security_state.h" #import "components/security_state/ios/security_state_utils.h" #import "ios/chrome/browser/intelligence/features/features.h" #import "ios/chrome/browser/intelligence/proto_wrappers/annotated_page_content_extraction_utils.h" #import "ios/chrome/browser/intelligence/proto_wrappers/frame_grafter.h" #import "ios/chrome/browser/intelligence/proto_wrappers/page_context_extractor_java_script_feature.h" #import "ios/chrome/browser/intelligence/proto_wrappers/page_context_utils.h" #import "ios/chrome/browser/intelligence/proto_wrappers/page_context_wrapper_config.h" #import "ios/chrome/browser/intelligence/proto_wrappers/page_context_wrapper_metrics.h" #import "ios/chrome/browser/snapshots/model/snapshot_tab_helper.h" #import "ios/public/provider/chrome/browser/bwg/gemini_api.h" #import "ios/web/find_in_page/find_in_page_java_script_feature.h" #import "ios/web/public/js_messaging/content_world.h" #import "ios/web/public/js_messaging/web_frame.h" #import "ios/web/public/js_messaging/web_frames_manager.h" #import "ios/web/public/ui/crw_web_view_proxy.h" #import "ios/web/public/ui/crw_web_view_scroll_view_proxy.h" #import "ios/web/public/web_state.h" #import "net/base/schemeful_site.h" #import "url/gurl.h" #import "url/origin.h" #import "url/url_constants.h" // TODO(crbug.com/458081684): Move away from all autofill dependencies once // the migration in ios/web is done for frame registration. namespace { // The default Page Context execution timeout. base::TimeDelta kDefaultPageContextTimeout = base::Seconds(1); // Url used for for data urls. constexpr const char kDataUrl[] = "data:"; // The timeout for waiting for remote frame registration. base::TimeDelta kRegistrationTimeout = base::Milliseconds(200); // The key for whether the PageContext should be detached. The value is a // bool. constexpr const char kShouldDetachPageContext[] = "shouldDetachPageContext"; // The key for the current node's innerText in the JavaScript object. The // value is a string. constexpr const char kCurrentNodeInnerTextDictKey[] = "currentNodeInnerText"; // The key for the children frames in the JavaScript object. The value is an // array of objects. constexpr const char kChildrenFramesDictKey[] = "children"; // The key for the PageInteractionInfo of the main frame. constexpr const char kPageInteractionInfoDictKey[] = "pageInteractionInfo"; // The key for the ViewportGeometry of the main frame. constexpr const char kViewportGeometryDictKey[] = "viewportGeometry"; // The key for the links of the frame in the JavaScript object. The value is // an array of objects. constexpr const char kFrameLinksDictKey[] = "links"; // The key for a link's HREF/URL field in the JavaScript object. The value is // a string. constexpr const char kLinkHREFDictKey[] = "href"; // The key for a link's innerText in the JavaScript object. The value is a // string. constexpr const char kLinkTextDictKey[] = "linkText"; // The key for the remote frame token. constexpr const char kRemoteFrameTokenKey[] = "remoteToken"; // The JavaScript to be executed on each WebState's WebFrames, which retrieves // the innerText of the document body, and recursively traverses through // same-origin nested iframes to retrieve their innerTexts as well, // constructing a tree structure. The innerText of a node also includes text // inside open Shadow DOM roots, which is otherwise excluded from // Element.innerText. iframes are marked as processed with a nonce to avoid // duplicate text from frames, but only for the current run. Early returns if // the PageContext should be detached, or the frame is not the top-most // same-origin frame. constexpr const char16_t* kInnerTextTreeJavaScript = uR"DELIM( (() => { // Checks whether the PageContext should be detached. const shouldDetachPageContext = () => { // PageContext detachment logic injected below. $1 }; // If the PageContext should be detached, early return. if (shouldDetachPageContext()) { return { shouldDetachPageContext: true }; } // The script should only run if it has no same-origin parent. (The script // should only start execution on top-most nodes of a given origin). if (window.self !== window.top && location.ancestorOrigins?.[0] === location.origin) { // Not the top-most same-origin frame, early exit. return null; } // Direct references to `Element` members, used to read shadow roots and // query descendants without going through instance properties that a // hostile page could clobber (e.g. a form control named `querySelectorAll` // or `shadowRoot`). const shadowRootGetter = Object.getOwnPropertyDescriptor(Element.prototype, 'shadowRoot')?.get; const querySelectorAllMethod = Element.prototype.querySelectorAll; const safeShadowRoot = (element) => element ? (shadowRootGetter ? shadowRootGetter.call(element) : element.shadowRoot) : null; const safeQuerySelectorAllAsArray = (element, selector) => Array.from(querySelectorAllMethod.call(element, selector)); // Keep this implementation in sync with `getInnerTextIncludingShadowDom()` // in page_context_extractor.ts. // Collects the innerText of a node, additionally piercing open Shadow DOM // roots whose text is not reflected in Element.innerText. Nested shadow // roots are handled through the recursion. Closed shadow roots are // inaccessible from script and are skipped. const getInnerTextIncludingShadowDom = (node) => { if (!node) { return ''; } const texts = []; if (node.innerText) { texts.push(node.innerText); } // The node itself and its light-DOM descendants may be shadow hosts. // querySelectorAll does not cross shadow boundaries, so shadow hosts // nested inside another shadow tree are reached via the recursion. const hosts = safeQuerySelectorAllAsArray(node, '*'); if (safeShadowRoot(node)) { hosts.unshift(node); } for (const host of hosts) { const shadowRoot = safeShadowRoot(host); if (!shadowRoot) { continue; } // Iterate all child nodes (not just elements) so text authored // directly under the shadow root is not dropped. for (const shadowChild of shadowRoot.childNodes) { let shadowText = ''; if (shadowChild.nodeType === Node.ELEMENT_NODE) { shadowText = getInnerTextIncludingShadowDom(shadowChild); } else if (shadowChild.nodeType === Node.TEXT_NODE) { shadowText = (shadowChild.textContent || '').trim(); } if (shadowText) { texts.push(shadowText); } } } // `innerText` above collects all light-DOM text at once. Shadow-DOM // chunks are therefore appended after it rather than inserted at their // visual positions. return texts.join('\n'); }; // Recursively constructs the innerText tree for the passed node and its // children same-origin iframes. const constructSameOriginInnerTextTree = (node, frameURL, frameTitle, nonceAttributeValue) => { // Early return if the node is null or already processed. if (!node || node.getAttribute('data-__gCrWeb-innerText-processed') === nonceAttributeValue) { return null; } // Mark node as processed. node.setAttribute('data-__gCrWeb-innerText-processed', nonceAttributeValue); // Get all nested iframes within the current node. const nestedIframes = node.getElementsByTagName('iframe'); const childNodeInnerTexts = [...nestedIframes].map(iframe => { if (!iframe) { return null; } // Try to access the iframe's body, failure is possible (cross-origin iframes). let iframeBody; let iframeTitle; try { const contentDoc = iframe.contentDocument; iframeBody = contentDoc ? contentDoc.body : null; iframeTitle = contentDoc ? contentDoc.title : ''; } catch (error) { return null; } // Recursively construct the innerText tree for the iframe's body. return iframeBody ? constructSameOriginInnerTextTree(iframeBody, iframe.src, iframeTitle, nonceAttributeValue) : null; }); const result = { currentNodeInnerText: getInnerTextIncludingShadowDom(node), children: childNodeInnerTexts.filter(item => item !== null), sourceUrl: frameURL, title: frameTitle, }; // Anchor tag retrieval logic injected below. $2 return result; }; return constructSameOriginInnerTextTree(document.body, window.location.href, document.title, "$3"); })(); )DELIM"; // The JavaScript to be executed in each WebFrame which gets all of a frame's // anchor tags and adds them to an array with their corresponding URL and // textContent (which includes all text, including text that is not visually // rendered). Injected into the main script. constexpr const char16_t* kAnchorTagsJavaScript = uR"DELIM( // Add all the frame's anchor tags to a links array with their HREF/URL and // textContent. const linksArray = []; const anchorElements = node.querySelectorAll('a[href]'); anchorElements.forEach((anchor) => { linksArray.push({ href: (anchor instanceof SVGAElement) ? anchor.href.baseVal : anchor.href, linkText: anchor.textContent }); }); result.links = linksArray; )DELIM"; // We match Blue*'s PDF size limit of 64 MB. // TODO(crbug.com/485311221): make this configurable per-provider, e.g. // via Gemini's configurable params, instead of hardcoding it here. const NSUInteger kMaxPDFByteLimit = 64 * 1024 * 1024; } // namespace @implementation PageContextWrapper { base::WeakPtr _webState; // The amount of async tasks this specific instance of the PageContext wrapper // needs to complete before executing the `completionCallback`. NSInteger _asyncTasksToComplete; // The timer which keeps track of the overall execution timeout. base::OneShotTimer _timeoutTimer; // The timer which keeps track of the registration timeout. base::OneShotTimer _registrationTimeoutTimer; // The root node of the PageContext's AnnotatedPageContent (APC) tree. This // tree is constructed on the fly as values are returned from JavaScript. std::unique_ptr _rootAPCNode; // The string which aggregates all iframes' innerTexts. std::unique_ptr _innerText; // Whether the PageContext should be detached. Likely a protected page. BOOL _forceDetachPageContext; // Whether the PageContext is not extractable. BOOL _notExtractable; // Whether the PageContext was blocked because it's an unsafe/insecure page. BOOL _unsafePageBlocked; // The callback to execute once all async work is complete, whichs // relinquishes ownership of the PageContext proto to the callback's handler. base::OnceCallback _completionCallback; // Unique pointer to the PageContext proto. std::unique_ptr _pageContext; // The boxes to redact over the screenshot. std::vector _universalBoundingBoxesForRedaction; // Raw screenshot image before redaction masking is applied. UIImage* _rawScreenshotImage; // The current PageContext instance's metrics logger. Only created when async // tasks execution is started. PageContextWrapperMetrics* _pageContextMetrics; // Configuration for page context extraction. Using optional avoids using // the constructor. std::optional _config; // The SchemefulSite of the main frame, captured synchronously at the start // of extraction to prevent TOCTOU race conditions. net::SchemefulSite _mainFrameSite; // Graft frames across origins. FrameGrafter _grafter; // TODO(crbug.com/507879464): Add a metric to record the number of focused // frames per request to detect anomalies. // The identifiers of the traversed frames that reported being focused. std::vector _focusedFrameInfos; // Caches string representations of autofill sections to integers during // extraction to ensure consistency across the APC proto. base::flat_map _autofillSectionNumbers; // Whether the registration wait has completed or timed out. BOOL _registrationCompletedOrTimedOut; // Enforces that execution only happens once. BOOL _executionStarted; } - (instancetype)initWithWebState:(web::WebState*)webState config:(PageContextWrapperConfig)config completionCallback: (base::OnceCallback) completionCallback { CHECK(webState); self = [super init]; if (self) { _asyncTasksToComplete = 0; _webState = webState->GetWeakPtr(); _config = config; _mainFrameSite = net::SchemefulSite(webState->GetLastCommittedURL()); _completionCallback = std::move(completionCallback); // Create the PageContext proto/object. _pageContext = std::make_unique(); GURL url = _webState->GetLastCommittedURL(); if (url.SchemeIs(url::kDataScheme)) { _pageContext->set_url(kDataUrl); } else { _pageContext->set_url(url.spec()); } _pageContext->set_title(base::UTF16ToUTF8(_webState->GetTitle())); } return self; } - (instancetype)initWithWebState:(web::WebState*)webState completionCallback: (base::OnceCallback) completionCallback { return [self initWithWebState:webState config:PageContextWrapperConfigBuilder().Build() completionCallback:std::move(completionCallback)]; } - (void)dealloc { _timeoutTimer.Stop(); _registrationTimeoutTimer.Stop(); [self stopTextHighlighting]; } - (void)populatePageContextFieldsAsync { [self populatePageContextFieldsAsyncWithTimeout:kDefaultPageContextTimeout]; } - (void)populatePageContextFieldsAsyncWithTimeout:(base::TimeDelta)timeout { CHECK(!_executionStarted) << "A PageContextWrapper should only be used once."; _executionStarted = YES; if (_isLowPriorityExtraction) { __weak PageContextWrapper* weakSelf = self; base::SequencedTaskRunner::GetCurrentDefault()->PostTask( FROM_HERE, base::BindOnce(^{ [weakSelf populateAsyncFields:timeout]; })); return; } [self populateAsyncFields:timeout]; } #pragma mark - Setters // Sets the flag to enabled/disabled, and increments/decrements accordingly the // total amount of async tasks gating the completion callback. - (void)setShouldGetSnapshot:(BOOL)shouldGetSnapshot { if (_shouldGetSnapshot == shouldGetSnapshot) { return; } _shouldGetSnapshot = shouldGetSnapshot; _asyncTasksToComplete += shouldGetSnapshot ? 1 : -1; } // Sets the flag to enabled/disabled, and increments/decrements accordingly the // total amount of async tasks gating the completion callback. - (void)setShouldGetFullPagePDF:(BOOL)shouldGetFullPagePDF { if (_shouldGetFullPagePDF == shouldGetFullPagePDF) { return; } _shouldGetFullPagePDF = shouldGetFullPagePDF; _asyncTasksToComplete += shouldGetFullPagePDF ? 1 : -1; } // Sets the flag to enabled/disabled, and increments/decrements accordingly the // total amount of async tasks gating the completion callback. - (void)setShouldGetAnnotatedPageContent:(BOOL)shouldGetAnnotatedPageContent { if (_shouldGetAnnotatedPageContent == shouldGetAnnotatedPageContent) { return; } _shouldGetAnnotatedPageContent = shouldGetAnnotatedPageContent; // Only update `_asyncTasksToComplete` if `_shouldGetInnerText` is false, // since they both affect the same async task. if (!_shouldGetInnerText) { _asyncTasksToComplete += shouldGetAnnotatedPageContent ? 1 : -1; } } // Sets the flag to enabled/disabled, and increments/decrements accordingly the // total amount of async tasks gating the completion callback. - (void)setShouldGetInnerText:(BOOL)shouldGetInnerText { if (shouldGetInnerText == _shouldGetInnerText) { return; } _shouldGetInnerText = shouldGetInnerText; // Only update `_asyncTasksToComplete` if `_shouldGetAnnotatedPageContent` is // false, since they both affect the same async task. if (!_shouldGetAnnotatedPageContent) { _asyncTasksToComplete += shouldGetInnerText ? 1 : -1; } } #pragma mark - Private // Cancels the registration timeout. - (void)cancelRegistrationTimeout { _registrationTimeoutTimer.Stop(); } // Completes the registration wait then calls `barrier`. - (void)completeRegistrationWaitWithBarrier:(base::RepeatingClosure)barrier { if (_registrationCompletedOrTimedOut) { // Avoid double completion. return; } _registrationCompletedOrTimedOut = YES; [self cancelRegistrationTimeout]; barrier.Run(); } // Returns the registrar. - (autofill::ChildFrameRegistrar*)frameRegistrar { // Returns nullptr if the WebState has been destroyed during async operations. if (!_webState) { return nullptr; } return autofill::ChildFrameRegistrar::GetOrCreateForWebState(_webState.get()); } // Returns the WebFramesManager to use for executing Page Context script on // frames. - (web::WebFramesManager*)webFramesManager { web::ContentWorld world = base::FeatureList::IsEnabled(kPageContextExtractorRefactored) ? PageContextExtractorJavaScriptFeature::GetInstance() ->GetSupportedContentWorld() : web::ContentWorld::kPageContentWorld; return _webState->GetWebFramesManager(world); } // Populates the fields of the PageContext proto which necessitate async // calls. - (void)populateAsyncFields:(base::TimeDelta)timeout { CHECK_GE(_asyncTasksToComplete, 0); _pageContextMetrics = [[PageContextWrapperMetrics alloc] initWithAPCConfigVariant:_config->GetApcConfigVariant()]; if (!_webState || _asyncTasksToComplete == 0) { [self asyncWorkCompletedForPageContext]; return; } if (!CanExtractPageContextForWebState(_webState.get(), _shouldGetFullPagePDF)) { _notExtractable = YES; [self asyncWorkCompletedForPageContext]; return; } if (_config->block_unsafe_pages() && _webState->GetNavigationManager() && security_state::GetSecurityLevelForWebState(_webState.get()) == security_state::SecurityLevel::DANGEROUS) { _unsafePageBlocked = YES; [self asyncWorkCompletedForPageContext]; return; } __weak PageContextWrapper* weakSelf = self; // Start the timer. _timeoutTimer.Start(FROM_HERE, timeout, base::BindOnce(^{ [weakSelf onTimeout]; })); // Use a `BarrierClosure` to ensure all async tasks are completed before // executing the overall completion callback. The BarrierClosure will wait // until the `pageContextBarrier` callback is itself run // `_asyncTasksToComplete` times, then post the completion handler to // execute on the next loop of the current sequence. base::RepeatingClosure pageContextBarrier = base::BarrierClosure( _asyncTasksToComplete, base::BindPostTask(base::SequencedTaskRunner::GetCurrentDefault(), base::BindOnce(^{ [weakSelf asyncWorkCompletedForPageContext]; }))); // Asynchronous work. *IMPORTANT NOTES*: // When adding async tasks below, an accompanying setter should also be // created to follow the disabled-by-default pattern (which // increments/decrements `_asyncTasksToComplete` accordingly). Also, if a // given task is enabled, every code path for that task should eventually // execute the `pageContextBarrier` callback, otherwise the `BarrierClosure` // will never execute its completion block. if (_shouldGetSnapshot) { [self processSnapshotWithBarrier:pageContextBarrier]; } if (_shouldGetAnnotatedPageContent || _shouldGetInnerText) { [self processAnnotatedPageContentWithBarrier:pageContextBarrier]; } // Create full page PDF representation of the WebState, if enabled. if (_shouldGetFullPagePDF) { [_pageContextMetrics executionStartedForTask:PageContextTask::kPDF]; _webState->CreateFullPagePdf(base::BindOnce(^(NSData* PDFData) { [weakSelf encodeAndSetFullPagePDF:PDFData barrier:pageContextBarrier]; })); } } // Retrieve WebState snapshot. The barrier's callback will be executed for all // codepaths in this method. - (void)processSnapshotWithBarrier:(base::RepeatingClosure)barrier { [_pageContextMetrics executionStartedForTask:PageContextTask::kScreenshot]; __weak PageContextWrapper* weakSelf = self; auto callback = ^(UIImage* image) { __strong __typeof(weakSelf) strongSelf = weakSelf; if (!strongSelf) { return; } if ([strongSelf shouldUpdateSnapshotWithImage:image]) { [strongSelf updateSnapshotWithBarrier:barrier]; return; } [strongSelf stashRawScreenshotImage:image]; barrier.Run(); }; // If the WebState is currently visible, update the snapshot in case the // user was scrolling, otherwise retrieve the latest version in cache or on // disk. if (_webState->IsVisible()) { auto updateSnapshotCallback = base::BindOnce(^(std::optional result_matches) { // TODO(crbug.com/401282824): Log the matches count to measure text // highlighting precision. [weakSelf updateSnapshotWithCallback:callback]; }); // If there is text to highlight, do it before capturing the screenshot. if (_textToHighlight != nil) { web::FindInPageJavaScriptFeature* findInPageFeature = web::FindInPageJavaScriptFeature::GetInstance(); web::WebFrame* mainFrame = _webState ->GetWebFramesManager( findInPageFeature->GetSupportedContentWorld()) ->GetMainWebFrame(); findInPageFeature->Search(mainFrame, base::SysNSStringToUTF8(_textToHighlight), std::move(updateSnapshotCallback)); } else { std::move(updateSnapshotCallback).Run(std::nullopt); } } else { SnapshotTabHelper::FromWebState(_webState.get()) ->RetrieveColorSnapshot(callback); } } // Helper to extract page context for a given frame. // TODO(crbug.com/495446456): Clean up once the JSON experiment is done. - (void)extractPageContextForFrame:(web::WebFrame*)frame isMainFrame:(BOOL)isMainFrame nonce:(const std::string&)nonce annotatedPageContentBarrier: (base::RepeatingClosure)annotatedPageContentBarrier { PageContextExtractorJavaScriptFeature* extractorFeature = PageContextExtractorJavaScriptFeature::GetInstance(); __weak PageContextWrapper* weakSelf = self; // Use a timeout for the JS call larger than the wrapper's timer timeout // since this is the preferred way of timing out the dispatched jobs // (which will return a PageContextWrapperError::kTimeout error instead of // empty results). base::TimeDelta jsTimeout = _timeoutTimer.GetCurrentDelay() * 2; if (IsPageContextIPCOptimizationEnabled()) { // Callback to aggregate values from the JS execution via JSON string // parsing. auto callbackJson = [](PageContextWrapper* weakWrapper, base::RepeatingClosure barrier, BOOL isMainFrame, const url::Origin& securityOrigin, std::optional frameId, std::optional value) { // TODO(crbug.com/454261374): Remove `withError` from args once we // cleanup the old code. Can't provide an error object since the // javascript feature doesn't support that. [weakWrapper aggregateJavaScriptValue:value ? &value.value() : nullptr withError:nil isMainFrame:isMainFrame securityOrigin:securityOrigin localFrameToken:frameId]; barrier.Run(); // Defer destruction of the base::DictValue to a background thread to // prevent blocking the main thread. The object is self-contained and // thread-safe. if (value) { base::ThreadPool::PostTask( FROM_HERE, {base::TaskPriority::BEST_EFFORT}, base::BindOnce([](base::Value v) {}, std::move(*value))); } }; extractorFeature->ExtractPageContextJSON( frame, _config->graft_cross_origin_frame_content(), _config->use_rich_extraction(), _config->use_rich_extraction_with_actionable(), _config->extract_paid_content(), _config->attempt_paid_content_json_fixing(), _config->include_sensitive_payments_for_redaction(), _config->extract_autofill_otp_redactions(), _config->extract_password_screenshot_redactions(), nonce, jsTimeout, base::BindOnce( callbackJson, weakSelf, annotatedPageContentBarrier, isMainFrame, frame->GetSecurityOrigin(), DeserializeFrameIdAsLocalFrameToken(frame->GetFrameId()))); } else { // Callback to aggregate values from the JS execution via the base value // received directly from WebKit. auto callback = [](PageContextWrapper* weakWrapper, base::RepeatingClosure barrier, BOOL isMainFrame, const url::Origin& securityOrigin, std::optional frameId, const base::Value* value) { // TODO(crbug.com/454261374): Remove `withError` from args once we // cleanup the old code. Can't provide an error object since the // javascript feature doesn't support that. [weakWrapper aggregateJavaScriptValue:value withError:nil isMainFrame:isMainFrame securityOrigin:securityOrigin localFrameToken:frameId]; barrier.Run(); }; extractorFeature->ExtractPageContext( frame, _config->graft_cross_origin_frame_content(), _config->use_rich_extraction(), _config->use_rich_extraction_with_actionable(), _config->extract_paid_content(), _config->attempt_paid_content_json_fixing(), _config->include_sensitive_payments_for_redaction(), _config->extract_autofill_otp_redactions(), _config->extract_password_screenshot_redactions(), nonce, jsTimeout, base::BindOnce( callback, weakSelf, annotatedPageContentBarrier, isMainFrame, frame->GetSecurityOrigin(), DeserializeFrameIdAsLocalFrameToken(frame->GetFrameId()))); } } // Get the WebState's AnnotatedPageContent filled with innerTexts. The // barrier's callback will be executed for all codepaths in this method. - (void)processAnnotatedPageContentWithBarrier:(base::RepeatingClosure)barrier { if (_shouldGetAnnotatedPageContent) { [_pageContextMetrics executionStartedForTask:PageContextTask::kAnnotatedPageContent]; } if (_shouldGetInnerText) { [_pageContextMetrics executionStartedForTask:PageContextTask::kInnerText]; } web::WebFramesManager* manager = [self webFramesManager]; std::set webFrames = manager->GetAllWebFrames(); web::WebFrame* mainFrame = manager->GetMainWebFrame(); if (webFrames.empty() || !mainFrame) { if (_shouldGetAnnotatedPageContent) { [_pageContextMetrics executionFinishedForTask:PageContextTask::kAnnotatedPageContent withCompletionStatus:PageContextCompletionStatus::kFailure]; } if (_shouldGetInnerText) { [_pageContextMetrics executionFinishedForTask:PageContextTask::kInnerText withCompletionStatus:PageContextCompletionStatus::kFailure]; } barrier.Run(); return; } // Create the root node of the APC tree and its first root ContentNode. _rootAPCNode = std::make_unique(); _rootAPCNode->set_version( optimization_guide::proto::AnnotatedPageContentVersion:: ANNOTATED_PAGE_CONTENT_VERSION_1_0); if (_config->use_rich_extraction_with_actionable()) { _rootAPCNode->set_mode(optimization_guide::proto::AnnotatedPageContentMode:: ANNOTATED_PAGE_CONTENT_MODE_ACTIONABLE_ELEMENTS); } else { _rootAPCNode->set_mode(optimization_guide::proto::AnnotatedPageContentMode:: ANNOTATED_PAGE_CONTENT_MODE_DEFAULT); } _rootAPCNode->mutable_root_node() ->mutable_content_attributes() ->set_attribute_type(optimization_guide::proto::CONTENT_ATTRIBUTE_ROOT); CHECK(_webState); _rootAPCNode->set_tab_id(_webState->GetUniqueIdentifier().identifier()); // Create the aggregated innerText string. _innerText = std::make_unique(); // Use a `BarrierClosure` to ensure the JavaScript is done executing in // all WebFrames before executing the page context barrier `barrier`, // which in turn signals to the PageContextWrapper that the APC is done // being processed. The BarrierClosure will wait until the // `annotatedPageContentBarrier` callback is itself run once per WebFrame // (+1 since we execute the JS explicitly on the main frame first). __weak PageContextWrapper* weakSelf = self; base::RepeatingClosure annotatedPageContentBarrier = base::BarrierClosure( webFrames.size() + 1, base::BindOnce(^{ [weakSelf handlePageContentExtractionCompletedWithBarrier:barrier]; })); std::string nonce = base::Token::CreateRandom().ToString(); if (_config->use_refactored_extractor()) { // Use the new way for extracting context. // Autofill information is only needed when extracting detailed annotated // page content. It is not needed when extracting inner text. // TODO(crbug.com/493904351): Add kill switch by using autofill config bit. if (_shouldGetAnnotatedPageContent) { optimization_guide::proto::AutofillInformation* autofillInformation = _rootAPCNode->mutable_profile_information() ->mutable_autofill_information(); PopulateAutofillInformation(_webState.get(), autofillInformation); } if (ios::provider::IsProtectedUrl(mainFrame->GetUrl().spec())) { _forceDetachPageContext = YES; annotatedPageContentBarrier.Run(); } else { [self extractPageContextForFrame:mainFrame isMainFrame:YES nonce:nonce annotatedPageContentBarrier:annotatedPageContentBarrier]; } // Execute the JavaScript on each other WebFrame and pass in the callback // (which executes the barrier when run). for (web::WebFrame* webFrame : webFrames) { if (ios::provider::IsProtectedUrl(webFrame->GetUrl().spec())) { _forceDetachPageContext = YES; } // Skip if it's the main frame since it was already processed above, or // if Page Context should already be force detached. if (!webFrame || webFrame->IsMainFrame() || _forceDetachPageContext) { annotatedPageContentBarrier.Run(); continue; } if (_config->include_same_site_only()) { net::SchemefulSite childSite(webFrame->GetSecurityOrigin()); if (_mainFrameSite != childSite) { annotatedPageContentBarrier.Run(); continue; } } [self extractPageContextForFrame:webFrame isMainFrame:NO nonce:nonce annotatedPageContentBarrier:annotatedPageContentBarrier]; } } else { // Use the legacy way for extracting context. // Callback to aggregate values from the JS execution. auto callback = [](PageContextWrapper* weakWrapper, base::RepeatingClosure barrier, BOOL isMainFrame, const url::Origin& securityOrigin, std::optional frameId, const base::Value* value, NSError* error) { [weakWrapper aggregateJavaScriptValue:value withError:error isMainFrame:isMainFrame securityOrigin:securityOrigin localFrameToken:frameId]; barrier.Run(); }; // Construct the JavaScript script to be executed on each Web Frame with a // random token as nonce to differentiate between runs/executions. std::u16string script = base::ReplaceStringPlaceholders( kInnerTextTreeJavaScript, base::span( {ios::provider::GetPageContextShouldDetachScript(), kAnchorTagsJavaScript, base::UTF8ToUTF16(nonce)}), nullptr); // TODO(crbug.com/452568673): Refactor the force detach logic. // If the page is not protected, execute the JavaScript on the main // WebFrame first and pass in the callback (which executes the barrier // when run). if (ios::provider::IsProtectedUrl(mainFrame->GetUrl().spec())) { _forceDetachPageContext = YES; annotatedPageContentBarrier.Run(); } else { mainFrame->ExecuteJavaScript( script, base::BindOnce( callback, weakSelf, annotatedPageContentBarrier, /*isMainFrame=*/YES, mainFrame->GetSecurityOrigin(), DeserializeFrameIdAsLocalFrameToken(mainFrame->GetFrameId()))); } // Execute the JavaScript on each other WebFrame and pass in the callback // (which executes the barrier when run). for (web::WebFrame* webFrame : webFrames) { if (ios::provider::IsProtectedUrl(webFrame->GetUrl().spec())) { _forceDetachPageContext = YES; } // Skip if it's the main frame since it was already processed above, or // if Page Context should already be force detached. if (!webFrame || webFrame->IsMainFrame() || _forceDetachPageContext) { annotatedPageContentBarrier.Run(); continue; } if (_config->include_same_site_only()) { net::SchemefulSite childSite(webFrame->GetSecurityOrigin()); if (_mainFrameSite != childSite) { annotatedPageContentBarrier.Run(); continue; } } webFrame->ExecuteJavaScript( script, base::BindOnce( callback, weakSelf, annotatedPageContentBarrier, /*isMainFrame=*/NO, webFrame->GetSecurityOrigin(), DeserializeFrameIdAsLocalFrameToken(webFrame->GetFrameId()))); } } } // All async tasks are complete, execute the overall completion callback. // Relinquish ownership to the callback handler. - (void)asyncWorkCompletedForPageContext { _timeoutTimer.Stop(); if (!_completionCallback) { return; } [self stopTextHighlighting]; PageContextWrapperCallbackResponse response; PageContextCompletionStatus completionStatus; // Construct the response and completion status, either with the expected // value or an error. if (!_webState) { response = base::unexpected(PageContextWrapperError::kGenericError); completionStatus = PageContextCompletionStatus::kFailure; } else if (_notExtractable) { response = base::unexpected(PageContextWrapperError::kPageNotExtractableError); completionStatus = PageContextCompletionStatus::kNotExtractable; } else if (_unsafePageBlocked) { response = base::unexpected(PageContextWrapperError::kPageUnsafeError); completionStatus = PageContextCompletionStatus::kPageUnsafe; } else if (_forceDetachPageContext) { response = base::unexpected(PageContextWrapperError::kForceDetachError); completionStatus = PageContextCompletionStatus::kProtected; } else if (_shouldGetAnnotatedPageContent && !_pageContext->has_annotated_page_content()) { response = base::unexpected(PageContextWrapperError::kAPCError); completionStatus = PageContextCompletionStatus::kFailure; } else if (_shouldGetInnerText && !_pageContext->has_inner_text()) { response = base::unexpected(PageContextWrapperError::kInnerTextError); completionStatus = PageContextCompletionStatus::kFailure; } else if (_shouldGetFullPagePDF && !_pageContext->has_pdf_data()) { response = base::unexpected(PageContextWrapperError::kPDFDataError); completionStatus = PageContextCompletionStatus::kFailure; } else { // TODO(crbug.com/509589346): CHECK `registrar` since it is known here that // there is a `_webState`. // Set the focused frame based on the collected data on the same origin and // cross-origin frames. autofill::ChildFrameRegistrar* registrar = [self frameRegistrar]; if (_config->use_rich_extraction() && registrar) { ResolveFocusedFrame(_focusedFrameInfos, _grafter.GetRemoteFrames(), registrar, _pageContext->mutable_annotated_page_content()); } if (_config->graft_cross_origin_frame_content() && registrar) { ResolveCrossSiteFrameContent( _grafter, registrar, _config->include_same_site_only(), _pageContext->mutable_annotated_page_content()); } if (_rawScreenshotImage) { [self applyRedactionsAndEncodeScreenshot:_rawScreenshotImage]; } if (_shouldGetSnapshot && !_pageContext->has_tab_screenshot()) { response = base::unexpected(PageContextWrapperError::kScreenshotError); completionStatus = PageContextCompletionStatus::kFailure; } else { response = base::ok(std::move(_pageContext)); completionStatus = PageContextCompletionStatus::kSuccess; } } [_pageContextMetrics executionFinishedForTask:PageContextTask::kOverall withCompletionStatus:completionStatus]; std::move(_completionCallback).Run(std::move(response)); } // Returns YES if the image is nil and forcing the update of missing snapshots // is enabled. - (BOOL)shouldUpdateSnapshotWithImage:(UIImage*)image { return !image && _shouldForceUpdateMissingSnapshots; } // Updates the snapshot for the given WebState, and executes the `barrier` // callback when finished. - (void)updateSnapshotWithBarrier:(base::RepeatingClosure)barrier { if (!_webState) { barrier.Run(); return; } __weak PageContextWrapper* weakSelf = self; SnapshotTabHelper::FromWebState(_webState.get()) ->UpdateSnapshotWithCallback(^(UIImage* image) { __strong __typeof(weakSelf) strongSelf = weakSelf; if (!strongSelf) { return; } [strongSelf stashRawScreenshotImage:image]; barrier.Run(); }); } // Updates the current WebState's snapshot with the given callback. - (void)updateSnapshotWithCallback:(void (^)(UIImage*))callback { if (_webState) { SnapshotTabHelper::FromWebState(_webState.get()) ->UpdateSnapshotWithCallback(callback); } else { callback(nil); } } // Returns YES for any RedactionDecision enums that require screenshot // redaction. - (BOOL)shouldRedactDecisionForScreenshot: (optimization_guide::proto::RedactionDecision)decision { switch (decision) { case optimization_guide::proto:: REDACTION_DECISION_REDACTED_HAS_BEEN_PASSWORD: case optimization_guide::proto:: REDACTION_DECISION_REDACTED_CUSTOM_PASSWORD_CSS: case optimization_guide::proto:: REDACTION_DECISION_REDACTED_CUSTOM_PASSWORD_JS: return _config->extract_password_screenshot_redactions(); case optimization_guide::proto:: REDACTION_DECISION_REDACTED_IS_SENSITIVE_PAYMENT_FIELD: return _config->include_sensitive_payments_for_redaction(); case optimization_guide::proto::REDACTION_DECISION_REDACTED_IS_OTP: return _config->extract_autofill_otp_redactions(); case optimization_guide::proto::REDACTION_DECISION_NO_REDACTION_NECESSARY: case optimization_guide::proto:: REDACTION_DECISION_UNREDACTED_EMPTY_PASSWORD: case optimization_guide::proto:: REDACTION_DECISION_UNREDACTED_EMPTY_CUSTOM_PASSWORD: case optimization_guide::proto:: REDACTION_DECISION_UNREDACTED_EMPTY_PAYMENT_FIELD: case optimization_guide::proto:: REDACTION_DECISION_UNREDACTED_EMPTY_OTP_FIELD: return NO; default: return NO; } } - (void)applyRedactionsAndEncodeScreenshot:(UIImage*)image { if (_pageContext && _pageContext->has_annotated_page_content()) { _grafter.CollectFormControlRedactionBoxesFromTree( _pageContext->annotated_page_content().root_node()); } // If the page has a non-standard zoomScale, we scale the extracted DOM layout // coordinates by zoomScale so that they align 1:1 with the pixels in the // screenshot bitmap captured from WKWebView. CGFloat zoomScale = 1.0; if (_webState && _webState->GetWebViewProxy() && _webState->GetWebViewProxy().scrollViewProxy) { zoomScale = _webState->GetWebViewProxy().scrollViewProxy.zoomScale; } _universalBoundingBoxesForRedaction.clear(); for (const auto& entry : _grafter.universal_bounding_boxes_for_redaction()) { if ([self shouldRedactDecisionForScreenshot:entry.decision]) { _universalBoundingBoxesForRedaction.push_back( CGRectMake(entry.visible_box.origin.x * zoomScale, entry.visible_box.origin.y * zoomScale, entry.visible_box.size.width * zoomScale, entry.visible_box.size.height * zoomScale)); } } if (!_universalBoundingBoxesForRedaction.empty()) { UIGraphicsImageRendererFormat* format = [UIGraphicsImageRendererFormat defaultFormat]; format.scale = image.scale; format.opaque = NO; UIGraphicsImageRenderer* renderer = [[UIGraphicsImageRenderer alloc] initWithSize:image.size format:format]; image = [renderer imageWithActions:^(UIGraphicsImageRendererContext* context) { [image drawAtPoint:CGPointZero]; [[UIColor blackColor] setFill]; for (const CGRect& rect : _universalBoundingBoxesForRedaction) { [[UIBezierPath bezierPathWithRect:rect] fill]; } }]; } NSData* imageData = UIImagePNGRepresentation(image); if (!imageData) { [_pageContextMetrics executionFinishedForTask:PageContextTask::kScreenshot withCompletionStatus:PageContextCompletionStatus::kFailure]; DLOG(WARNING) << "Failed to convert the screenshot to PNG."; return; } NSString* base64String = [imageData base64EncodedStringWithOptions:0]; _pageContext->set_tab_screenshot(base::SysNSStringToUTF8(base64String)); _universalBoundingBoxesForRedaction.clear(); _rawScreenshotImage = nil; [_pageContextMetrics executionFinishedForTask:PageContextTask::kScreenshot withCompletionStatus:PageContextCompletionStatus::kSuccess]; } // Temporarily saves the raw screenshot image captured from the WebState so that // redactions can be applied and PNG encoding performed after APC extraction and // iframe coordinate grafting complete. - (void)stashRawScreenshotImage:(UIImage*)image { [self stopTextHighlighting]; if (!image) { [_pageContextMetrics executionFinishedForTask:PageContextTask::kScreenshot withCompletionStatus:PageContextCompletionStatus::kFailure]; DLOG(WARNING) << "Failed to fetch webpage screenshot."; return; } _rawScreenshotImage = image; } // If it exists, convert the PDF data to base64 encoded string and set it in // the PageContext proto. - (void)setEncodedPDFData:(std::string)base64PDFData { _pageContext->set_pdf_data(std::move(base64PDFData)); [_pageContextMetrics executionFinishedForTask:PageContextTask::kPDF withCompletionStatus:PageContextCompletionStatus::kSuccess]; } // If it exists, convert the PDF data to base64 encoded string and set it in // the PageContext proto. - (void)encodeAndSetFullPagePDF:(NSData*)PDFData barrier:(base::RepeatingClosure)barrier { if (!PDFData || PDFData.length > kMaxPDFByteLimit) { [_pageContextMetrics executionFinishedForTask:PageContextTask::kPDF withCompletionStatus:PageContextCompletionStatus::kFailure]; DLOG(WARNING) << "Failed to fetch webpage PDF data (or size limit exceeded)."; barrier.Run(); return; } __weak PageContextWrapper* weakSelf = self; base::ThreadPool::PostTaskAndReplyWithResult( FROM_HERE, {base::TaskPriority::USER_VISIBLE}, base::BindOnce( ^(NSData* data) { return base::Base64Encode(base::apple::NSDataToSpan(data)); }, PDFData), base::BindOnce(^(std::string base64_str) { PageContextWrapper* strongSelf = weakSelf; if (strongSelf) { [strongSelf setEncodedPDFData:std::move(base64_str)]; } barrier.Run(); })); } // Adds a frame's focus info to the flat array if it is focused. - (void)addFrameFocusInfo:(bool)isFocused documentId:(const std::string&)documentId { if (isFocused) { FrameFocusInfo info; info.document_id = documentId; _focusedFrameInfos.push_back(std::move(info)); } } // Helper to populate the Rich Extraction content for both the main frame and // iframes. - (void)populateForRichExtractionWithValue:(const base::DictValue&)value securityOrigin:(const url::Origin&)securityOrigin isMainFrame:(BOOL)isMainFrame localFrameToken: (std::optional) localFrameToken { // Populate the annotated page content for main frame including // frame data. const base::DictValue* rootNodeValue = value.FindDict("rootNode"); if (!rootNodeValue) { // There is no point in processing content if there is no root node. return; } const base::DictValue* frameDataValue = value.FindDict("frameData"); optimization_guide::proto::ContentNode* destinationContentNode = nullptr; optimization_guide::proto::FrameData* destinationFrameData = nullptr; // Pick the destinations for the APC node content and frame data. if (isMainFrame) { // Main frame: Use the root node as the destination. if (!_rootAPCNode) { _rootAPCNode = std::make_unique(); } destinationContentNode = _rootAPCNode->mutable_root_node(); destinationFrameData = _rootAPCNode->mutable_main_frame_data(); // Populate DocumentIdentifier for the Main Frame. It is not provided from // the renderer so we need to do it here. autofill::RemoteFrameToken token = static_cast( base::UnguessableToken::Create()); destinationFrameData->mutable_document_identifier()->set_serialized_token( token.ToString()); // Register the Main Frame token to allow lookups (e.g. for actions). autofill::ChildFrameRegistrar* registrar = [self frameRegistrar]; if (registrar && localFrameToken) { registrar->RegisterMapping(token, *localFrameToken); } } else if (localFrameToken) { // Grafting possible: Use the content node directly from the grafter. FrameGrafter::FrameContent* content = _grafter.DeclareContent(*localFrameToken); if (!content) { // Content already declared or invalid, skip processing. return; } destinationContentNode = &content->content; destinationFrameData = &content->frame_data; } else { // Grafting not possible: Drop the frame content to avoid duplicate // node_id collisions and untranslated local coordinate distortion. return; } // Destination placeholders must be set to something even if their content // won't be set. CHECK(destinationContentNode); CHECK(destinationFrameData); // Having root node content is a must at this point. CHECK(rootNodeValue); // Callback to collect the focus status of each traversed frame (same-origin // nested frames and the current frame) into the flat array. __weak __typeof(self) weakSelf = self; base::RepeatingCallback on_frame_extracted = base::BindRepeating( ^(bool isFocusedChild, const std::string& documentId) { [weakSelf addFrameFocusInfo:isFocusedChild documentId:documentId]; }); std::optional autofill_context; if (_config->extract_autofill() || _config->extract_autofill_credit_card_redactions() || _config->include_sensitive_payments_for_redaction() || _config->extract_autofill_otp_redactions()) { autofill_context.emplace(_webState, localFrameToken, _config->extract_autofill_credit_card_redactions(), _config->extract_autofill_otp_redactions(), &_autofillSectionNumbers); } PopulateAPCNodeFromContentTree( *rootNodeValue, securityOrigin, _grafter, autofill_context ? &*autofill_context : nullptr, destinationContentNode, on_frame_extracted); // Populate the frame data for this frame and determine whether it is focused. bool isFocused = false; if (frameDataValue) { isFocused = PopulateFrameDataNode(*frameDataValue, securityOrigin, destinationFrameData) .is_focused; } // Record focus status for later resolution if focused unless the frame is // the main frame which can be resolved right away. if (isFocused) { FrameFocusInfo info; info.local_token = localFrameToken; if (isMainFrame) { info.document_id = destinationFrameData->document_identifier().serialized_token(); } _focusedFrameInfos.push_back(std::move(info)); } // Populate the page data extracted from the main frame. if (isMainFrame) { const base::DictValue* pageInteractionInfoValue = value.FindDict(kPageInteractionInfoDictKey); if (pageInteractionInfoValue) { PopulatePageInteractionInfoNode( *pageInteractionInfoValue, _rootAPCNode->mutable_page_interaction_info()); } if (const base::DictValue* viewportGeometryValue = value.FindDict(kViewportGeometryDictKey)) { PopulateViewportGeometryNode(*viewportGeometryValue, _rootAPCNode->mutable_viewport_geometry()); } } } // Helper to populate the main frame's content for light extraction. - (void)populateMainFrameForLightExtractionWithValue: (const base::DictValue&)value securityOrigin: (const url::Origin&)securityOrigin { // Light Extraction: Populate the annotated page content for main frame. [self populateMainFrameSubtreeForLightExtractionWithValue:value origin:securityOrigin]; // Recursively populate the ContentNode subtree // for any of the main WebFrame's children iframes. Children can be remote // or same origin frames. const base::ListValue* childrenFrames = value.FindList(kChildrenFramesDictKey); if (childrenFrames && !childrenFrames->empty()) { for (const auto& childFrame : *childrenFrames) { if (!childFrame.is_dict()) { continue; } // Note: We need `node` to hold until extraction is done because it // is part of the root APC node content. optimization_guide::proto::ContentNode* node = _rootAPCNode->mutable_root_node()->add_children_nodes(); [self populateIframeSubtreeWithValue:childFrame.GetDict() origin:securityOrigin node:node]; } } } // Helper to populate an iframe's content for light extraction. - (void)populateIframeForLightExtractionWithValue:(const base::DictValue&)value securityOrigin: (const url::Origin&)securityOrigin localFrameToken: (std::optional) localFrameToken { optimization_guide::proto::ContentNode* destinationContentNode; // Pick a destination ContentNode to fill with the iframe content. if (_config->graft_cross_origin_frame_content()) { if (!localFrameToken) { return; } // Grafting possible: Populate iframe content by respecting the DOM // structure. // In ApcV1, all the frame data is stored in the `content` field. destinationContentNode = &_grafter.DeclareContent(*localFrameToken)->content; } else { // Grafting not possible: Add new child to the root node as the default // location for the iframe node. destinationContentNode = _rootAPCNode->mutable_root_node()->add_children_nodes(); } CHECK(destinationContentNode); [self populateIframeSubtreeWithValue:value origin:securityOrigin node:destinationContentNode]; } // If it exists, parse the returned JavaScript value from the WebFrame, // construct its ContentNode subtree and insert it into the APC tree. - (void)aggregateJavaScriptValue:(const base::Value*)value withError:(NSError*)error isMainFrame:(BOOL)isMainFrame securityOrigin:(const url::Origin&)securityOrigin localFrameToken: (std::optional)localFrameToken { if (error || !value || !value->is_dict() || !_webState) { if (error) { // TODO(crbug.com/401282824): Log the failure rate of aggregation. DLOG(WARNING) << "Failed to fetch frame's innerText tree." << base::SysNSStringToUTF8([error localizedDescription]); } return; } if (_forceDetachPageContext) { return; } const base::DictValue& valueAsDict = value->GetDict(); // Check if PageContext should be force detached. // TODO(crbug.com/471244309): Force detaching PageContext shouldn't depend on // fetching innerText/APC, it should always be enabled. std::optional shouldDetachPageContext = valueAsDict.FindBool(kShouldDetachPageContext); if (shouldDetachPageContext.has_value() && shouldDetachPageContext.value()) { _forceDetachPageContext = YES; if (_shouldGetAnnotatedPageContent) { [_pageContextMetrics executionFinishedForTask:PageContextTask::kAnnotatedPageContent withCompletionStatus:PageContextCompletionStatus::kProtected]; } if (_shouldGetInnerText) { [_pageContextMetrics executionFinishedForTask:PageContextTask::kInnerText withCompletionStatus:PageContextCompletionStatus::kProtected]; } return; } if (_config->graft_cross_origin_frame_content() && valueAsDict.FindString(kRemoteFrameTokenKey)) { // Do not aggregate the results if they contain a remote frame token as the // root content from frame extraction which is unexpected and shouldn't be // handled. // TODO(crbug.com/464473686): Measure this via a metric. return; } // Create a special subtree for the mainframe, and then recursively populate // its children iframe subtrees. Else, recursively populate cross-origin // iframes. if (isMainFrame) { // Populate main frame root node. if (_config->use_rich_extraction()) { [self populateForRichExtractionWithValue:valueAsDict securityOrigin:securityOrigin isMainFrame:YES localFrameToken:localFrameToken]; } else { [self populateMainFrameForLightExtractionWithValue:valueAsDict securityOrigin:securityOrigin]; } return; } // Populate iframes nodes from the root. if (_config->use_rich_extraction()) { [self populateForRichExtractionWithValue:valueAsDict securityOrigin:securityOrigin isMainFrame:NO localFrameToken:localFrameToken]; } else { [self populateIframeForLightExtractionWithValue:valueAsDict securityOrigin:securityOrigin localFrameToken:localFrameToken]; } } // Called when all JS extraction tasks are completed. It handles the remote // frame registration waiting logic before finalizing the APC tree. - (void)handlePageContentExtractionCompletedWithBarrier: (base::RepeatingClosure)barrier { [self webFramesAnnotatedPageContentFetchCompleted]; if (!_config->graft_cross_origin_frame_content()) { barrier.Run(); return; } std::vector discoveredRemoteTokens = _grafter.GetRemoteFrames(); if (discoveredRemoteTokens.empty()) { barrier.Run(); return; } // TODO(crbug.com/475858687): Measure registration time. autofill::ChildFrameRegistrar* registrar = [self frameRegistrar]; if (!registrar) { barrier.Run(); return; } // Reset the flag for the new wait cycle. _registrationCompletedOrTimedOut = NO; // TODO(crbug.com/475854782): Measure registration timeouts. __weak PageContextWrapper* weakSelf = self; auto run_barrier_once = base::BindRepeating( [](base::RepeatingClosure barrier, __weak PageContextWrapper* weakSelf) { [weakSelf completeRegistrationWaitWithBarrier:barrier]; }, std::move(barrier), weakSelf); // Start the global registration timer. _registrationTimeoutTimer.Start(FROM_HERE, kRegistrationTimeout, run_barrier_once); // The barrier that waits for all registrations. base::RepeatingClosure registrationBarrier = base::BarrierClosure(discoveredRemoteTokens.size(), run_barrier_once); for (const auto& token : discoveredRemoteTokens) { registrar->DeclareNewRemoteToken( token, base::BindOnce([](base::RepeatingClosure barrier, autofill::LocalFrameToken) { barrier.Run(); }, registrationBarrier)); } } // Set the constructed APC tree on the PageContext proto. - (void)webFramesAnnotatedPageContentFetchCompleted { if (_shouldGetInnerText) { _pageContext->set_allocated_inner_text(_innerText.release()); [_pageContextMetrics executionFinishedForTask:PageContextTask::kInnerText withCompletionStatus:PageContextCompletionStatus::kSuccess]; } if (_shouldGetAnnotatedPageContent) { size_t sizeInBytes = _rootAPCNode->ByteSizeLong(); _pageContext->set_allocated_annotated_page_content(_rootAPCNode.release()); [_pageContextMetrics executionFinishedForTask:PageContextTask::kAnnotatedPageContent withCompletionStatus:PageContextCompletionStatus::kSuccess]; int sizeInBytesInt = base::saturated_cast(sizeInBytes); [_pageContextMetrics logAnnotatedPageContentSize:sizeInBytesInt]; [_pageContextMetrics logAnnotatedPageContentHighRangeSizeInKb:sizeInBytesInt]; } } // Populate the main frame's ContentNode subtree with the correct nodes and // their values. Adds Main Frame data and the root text ContentNode. - (void)populateMainFrameSubtreeForLightExtractionWithValue: (const base::DictValue&)value origin:(const url::Origin&) origin { // Set the main frame's security origin. PopulateFrameDataNode(value, origin, _rootAPCNode->mutable_main_frame_data()); // Set its child text node. [self populateTextInfoNodeWithValue:value origin:origin parentNode:_rootAPCNode->mutable_root_node()]; // Set its children anchor nodes. [self populateAnchorNodeChildrenWithValue:value parentNode:_rootAPCNode->mutable_root_node()]; } // Populate a ContentNode with a TextInfo node and its correct values. - (void)populateTextInfoNodeWithValue:(const base::DictValue&)value origin:(const url::Origin&)origin parentNode:(optimization_guide::proto::ContentNode*) parentNode { if (!parentNode) { return; } // Early return if there is no text to add. const std::string* innerTextPtr = value.FindString(kCurrentNodeInnerTextDictKey); if (!innerTextPtr) { return; } std::string_view trimmedText = base::TrimWhitespaceASCII(*innerTextPtr, base::TRIM_ALL); if (trimmedText.empty()) { return; } // Create and add the text node. optimization_guide::proto::ContentNode* childTextNode = parentNode->add_children_nodes(); childTextNode->mutable_content_attributes()->set_attribute_type( optimization_guide::proto::CONTENT_ATTRIBUTE_TEXT); childTextNode->mutable_content_attributes() ->mutable_text_data() ->set_text_content(trimmedText); if (_shouldGetInnerText) { _innerText->append(trimmedText); } } // Populate the ContentNode subtree for an iframe with the correct values. // Also recursively populates the subtrees for all of this iframe's children. - (void)populateIframeSubtreeWithValue:(const base::DictValue&)value origin:(const url::Origin&)origin node:(optimization_guide::proto::ContentNode*) node { const std::string* remoteTokenString = _config->graft_cross_origin_frame_content() ? value.FindString(kRemoteFrameTokenKey) : nullptr; if (remoteTokenString) { // Register a placeholder for the frame content when the frame is on a // different origin from the parent frame. // // TODO(crbug.com/460823916): The grafter should re-attempt registering the // placeholder upon registration completion via DeclareNewRemoteToken(). We // could use extra barriers for that. std::optional remote = DeserializeFrameIdAsRemoteFrameToken(*remoteTokenString); if (remote) { _grafter.RegisterPlaceholder(*remote, node); } return; } // Create the child iframe node. node->mutable_content_attributes()->set_attribute_type( optimization_guide::proto::CONTENT_ATTRIBUTE_IFRAME); // Set its FrameData values. optimization_guide::proto::FrameData* nodeFrameData = node->mutable_content_attributes() ->mutable_iframe_data() ->mutable_frame_data(); PopulateFrameDataNode(value, origin, nodeFrameData); // Create the nested root child ContentNode. optimization_guide::proto::ContentNode* childRootNode = node->add_children_nodes(); childRootNode->mutable_content_attributes()->set_attribute_type( optimization_guide::proto::CONTENT_ATTRIBUTE_ROOT); // Create the nested text node. [self populateTextInfoNodeWithValue:value origin:origin parentNode:childRootNode]; // Create the children anchor nodes. [self populateAnchorNodeChildrenWithValue:value parentNode:childRootNode]; // Recursively populate the ContentNode subtree for any children iframes. // Child frame content will either be filled immediately or claimed for // later. const base::ListValue* childrenFrames = value.FindList(kChildrenFramesDictKey); if (childrenFrames && !childrenFrames->empty()) { for (const auto& childFrame : *childrenFrames) { if (childFrame.is_dict()) { optimization_guide::proto::ContentNode* childNode = childRootNode->add_children_nodes(); [self populateIframeSubtreeWithValue:childFrame.GetDict() origin:origin node:childNode]; } } } return; } // Populate all anchor tags as AnchorData nodes which are direct children // of `parentNode`. - (void)populateAnchorNodeChildrenWithValue:(const base::DictValue&)value parentNode: (optimization_guide::proto::ContentNode*) parentNode { if (!parentNode) { return; } const base::ListValue* links = value.FindList(kFrameLinksDictKey); if (!links || links->empty()) { return; } for (const auto& linkValue : *links) { [self populateAnchorNodeWithValue:&linkValue parentNode:parentNode]; } } // Creates an AnchorData node (with the corresponding URL) with one child // TextInfo node (with the corresponding innerText). Set the AnchorData node // as direct child of `parentNode`. - (void)populateAnchorNodeWithValue:(const base::Value*)linkData parentNode:(optimization_guide::proto::ContentNode*) parentNode { if (!linkData || !linkData->is_dict() || !parentNode) { return; } const std::string* href = linkData->GetDict().FindString(kLinkHREFDictKey); if (!href || href->empty()) { return; } // Create the anchor node. optimization_guide::proto::ContentNode* anchorNode = parentNode->add_children_nodes(); anchorNode->mutable_content_attributes()->set_attribute_type( optimization_guide::proto::CONTENT_ATTRIBUTE_ANCHOR); // Set the anchor data (the HREF). anchorNode->mutable_content_attributes()->mutable_anchor_data()->set_url( *href); // Create a child text node for the anchor's innerText. const std::string* linkText = linkData->GetDict().FindString(kLinkTextDictKey); if (!linkText || linkText->empty() || base::TrimWhitespaceASCII(*linkText, base::TRIM_ALL).empty()) { return; } // Set the child text node's text value. optimization_guide::proto::ContentNode* textNode = anchorNode->add_children_nodes(); textNode->mutable_content_attributes()->set_attribute_type( optimization_guide::proto::CONTENT_ATTRIBUTE_TEXT); textNode->mutable_content_attributes()->mutable_text_data()->set_text_content( *linkText); } // Stop the highlighting of text. - (void)stopTextHighlighting { if (!_textToHighlight) { return; } web::WebState* webState = _webState.get(); if (!webState) { return; } web::FindInPageJavaScriptFeature* find_in_page_feature = web::FindInPageJavaScriptFeature::GetInstance(); if (!find_in_page_feature) { return; } web::WebFramesManager* framesManager = webState->GetWebFramesManager( find_in_page_feature->GetSupportedContentWorld()); if (!framesManager) { return; } web::WebFrame* mainFrame = framesManager->GetMainWebFrame(); if (!mainFrame) { return; } find_in_page_feature->Stop(mainFrame); } // Called when the overall execution times out. Cancels the timer and executes // the completion callback with `kTimeout`. - (void)onTimeout { if (!_completionCallback) { return; } [self stopTextHighlighting]; DLOG(WARNING) << "PageContextWrapper execution timed out."; [_pageContextMetrics executionFinishedForTask:PageContextTask::kOverall withCompletionStatus:PageContextCompletionStatus::kTimeout]; std::move(_completionCallback) .Run(base::unexpected(PageContextWrapperError::kTimeout)); } #pragma mark - Testing // Test-only hooks to allow synchronous unit testing of the screenshot masking // and redaction bounding box collection engine without requiring an // asynchronous WKWebView E2E snapshot round-trip. - (void)encodeImageAndSetTabScreenshotForTesting:(UIImage*)image { [self stashRawScreenshotImage:image]; if (_rawScreenshotImage) { [self applyRedactionsAndEncodeScreenshot:_rawScreenshotImage]; } } - (void)setBoxesToRedactForTesting:(const std::vector&)boxes { _universalBoundingBoxesForRedaction = boxes; } - (const std::vector&)boxesToRedactForTesting { return _universalBoundingBoxesForRedaction; } - (optimization_guide::proto::PageContext*)pageContextForTesting { if (_rootAPCNode) { _pageContext->mutable_annotated_page_content()->CopyFrom(*_rootAPCNode); } return _pageContext.get(); } - (optimization_guide::proto::AnnotatedPageContent*)rootAPCNodeForTesting { if (!_rootAPCNode) { _rootAPCNode = std::make_unique(); } return _rootAPCNode.get(); } - (FrameGrafter&)grafterForTesting { return _grafter; } @end