1// Copyright 2014 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#import "ios/chrome/browser/snapshots/snapshot_generator.h"
6
7#include <algorithm>
8
9// TODO(crbug.com/636188): required to implement ViewHierarchyContainsWKWebView
10// for -drawViewHierarchyInRect:afterScreenUpdates:, remove once the workaround
11// is no longer needed.
12#import <WebKit/WebKit.h>
13
14#include "base/bind.h"
15#include "base/check_op.h"
16#include "base/task/post_task.h"
17#import "ios/chrome/browser/snapshots/snapshot_cache.h"
18#import "ios/chrome/browser/snapshots/snapshot_generator_delegate.h"
19#include "ios/chrome/browser/ui/ui_feature_flags.h"
20#import "ios/chrome/browser/ui/util/uikit_ui_util.h"
21#include "ios/web/public/thread/web_task_traits.h"
22#include "ios/web/public/thread/web_thread.h"
23#import "ios/web/public/web_client.h"
24#import "ios/web/public/web_state.h"
25#import "ios/web/public/web_state_observer_bridge.h"
26#include "ui/gfx/geometry/rect_f.h"
27#include "ui/gfx/image/image.h"
28
29#if !defined(__has_feature) || !__has_feature(objc_arc)
30#error "This file requires ARC support."
31#endif
32
33namespace {
34
35// Contains information needed for snapshotting.
36struct SnapshotInfo {
37  UIView* baseView;
38  CGRect snapshotFrameInBaseView;
39  CGRect snapshotFrameInWindow;
40  NSArray<UIView*>* overlays;
41};
42
43// Returns YES if |view| or any view it contains is a WKWebView.
44BOOL ViewHierarchyContainsWKWebView(UIView* view) {
45  if ([view isKindOfClass:[WKWebView class]])
46    return YES;
47  for (UIView* subview in view.subviews) {
48    if (ViewHierarchyContainsWKWebView(subview))
49      return YES;
50  }
51  return NO;
52}
53
54}  // namespace
55
56@interface SnapshotGenerator ()<CRWWebStateObserver>
57
58// The unique ID for the web state.
59@property(nonatomic, copy) NSString* tabID;
60
61// The associated web state.
62@property(nonatomic, assign) web::WebState* webState;
63
64@end
65
66@implementation SnapshotGenerator {
67  std::unique_ptr<web::WebStateObserver> _webStateObserver;
68}
69
70- (instancetype)initWithWebState:(web::WebState*)webState
71                           tabID:(NSString*)tabID {
72  if ((self = [super init])) {
73    DCHECK(webState);
74    DCHECK(tabID);
75    _webState = webState;
76    _tabID = tabID;
77
78    _webStateObserver = std::make_unique<web::WebStateObserverBridge>(self);
79    _webState->AddObserver(_webStateObserver.get());
80  }
81  return self;
82}
83
84- (void)dealloc {
85  if (_webState) {
86    _webState->RemoveObserver(_webStateObserver.get());
87    _webStateObserver.reset();
88    _webState = nullptr;
89  }
90}
91
92- (void)retrieveSnapshot:(void (^)(UIImage*))callback {
93  DCHECK(callback);
94  if (self.snapshotCache) {
95    [self.snapshotCache retrieveImageForSnapshotID:self.tabID
96                                          callback:callback];
97  } else {
98    callback(nil);
99  }
100}
101
102- (void)retrieveGreySnapshot:(void (^)(UIImage*))callback {
103  DCHECK(callback);
104
105  __weak SnapshotGenerator* weakSelf = self;
106  void (^wrappedCallback)(UIImage*) = ^(UIImage* image) {
107    if (!image) {
108      image = [weakSelf updateSnapshot];
109      if (image)
110        image = GreyImage(image);
111    }
112    callback(image);
113  };
114
115  SnapshotCache* snapshotCache = self.snapshotCache;
116  if (snapshotCache) {
117    [snapshotCache retrieveGreyImageForSnapshotID:self.tabID
118                                         callback:wrappedCallback];
119  } else {
120    wrappedCallback(nil);
121  }
122}
123
124- (UIImage*)updateSnapshot {
125  UIImage* snapshot = [self generateSnapshotWithOverlays:YES];
126  [self updateSnapshotCacheWithImage:snapshot];
127  return snapshot;
128}
129
130- (void)updateWebViewSnapshotWithCompletion:(void (^)(UIImage*))completion {
131  DCHECK(!web::GetWebClient()->IsAppSpecificURL(
132      self.webState->GetLastCommittedURL()));
133
134  if (![self canTakeSnapshot]) {
135    if (completion) {
136      base::PostTask(FROM_HERE, {web::WebThread::UI}, base::BindOnce(^{
137                       completion(nil);
138                     }));
139    }
140    return;
141  }
142  SnapshotInfo snapshotInfo = [self snapshotInfo];
143  CGRect snapshotFrameInWebView =
144      [self.webState->GetView() convertRect:snapshotInfo.snapshotFrameInBaseView
145                                   fromView:snapshotInfo.baseView];
146  [self.delegate snapshotGenerator:self
147      willUpdateSnapshotForWebState:self.webState];
148  __weak SnapshotGenerator* weakSelf = self;
149  self.webState->TakeSnapshot(
150      gfx::RectF(snapshotFrameInWebView),
151      base::BindRepeating(^(const gfx::Image& image) {
152        UIImage* snapshot = nil;
153        if (!image.IsEmpty()) {
154          snapshot = [weakSelf
155              snapshotWithOverlays:snapshotInfo.overlays
156                         baseImage:image.ToUIImage()
157                     frameInWindow:snapshotInfo.snapshotFrameInWindow];
158        }
159        [weakSelf updateSnapshotCacheWithImage:snapshot];
160        if (completion)
161          completion(snapshot);
162      }));
163}
164
165- (UIImage*)generateSnapshotWithOverlays:(BOOL)shouldAddOverlay {
166  if (![self canTakeSnapshot])
167    return nil;
168  SnapshotInfo snapshotInfo = [self snapshotInfo];
169  [self.delegate snapshotGenerator:self
170      willUpdateSnapshotForWebState:self.webState];
171  UIImage* baseImage =
172      [self snapshotBaseView:snapshotInfo.baseView
173             frameInBaseView:snapshotInfo.snapshotFrameInBaseView];
174  return [self
175      snapshotWithOverlays:(shouldAddOverlay ? snapshotInfo.overlays : nil)
176                 baseImage:baseImage
177             frameInWindow:snapshotInfo.snapshotFrameInWindow];
178}
179
180- (void)removeSnapshot {
181  [self.snapshotCache removeImageWithSnapshotID:self.tabID];
182}
183
184#pragma mark - Private methods
185
186// Returns NO if WebState or the view is not ready for snapshot.
187- (BOOL)canTakeSnapshot {
188  // This allows for easier unit testing of classes that use SnapshotGenerator.
189  if (!self.delegate)
190    return NO;
191
192  // Do not generate a snapshot if web usage is disabled (as the WebState's
193  // view is blank in that case).
194  if (!self.webState->IsWebUsageEnabled())
195    return NO;
196
197  return [self.delegate snapshotGenerator:self
198               canTakeSnapshotForWebState:self.webState];
199}
200
201// Returns a snapshot of |baseView| with |frameInBaseView|.
202- (UIImage*)snapshotBaseView:(UIView*)baseView
203             frameInBaseView:(CGRect)frameInBaseView {
204  DCHECK(baseView);
205  DCHECK(!CGRectIsEmpty(frameInBaseView));
206  // Note: When not using device scale, the output image size may slightly
207  // differ from the input size due to rounding.
208  const CGFloat kScale =
209      std::max<CGFloat>(1.0, [self.snapshotCache snapshotScaleForDevice]);
210  UIGraphicsBeginImageContextWithOptions(frameInBaseView.size, YES, kScale);
211  CGContext* context = UIGraphicsGetCurrentContext();
212  // This shifts the origin of the context to be the origin of the snapshot
213  // frame.
214  CGContextTranslateCTM(context, -frameInBaseView.origin.x,
215                        -frameInBaseView.origin.y);
216  BOOL snapshotSuccess = YES;
217
218  // TODO(crbug.com/636188): |-drawViewHierarchyInRect:afterScreenUpdates:| is
219  // buggy on iOS 8/9/10 (and state is unknown for iOS 11) causing GPU glitches,
220  // screen redraws during animations, broken pinch to dismiss on tablet, etc.
221  // Ensure iOS 11 is not affected by these issues before turning on
222  // |kSnapshotDrawView| experiment. On the other hand, |-renderInContext:| is
223  // buggy for WKWebView, which is used for some Chromium pages such as "No
224  // internet" or "Site can't be reached".
225  BOOL useDrawViewHierarchy = ViewHierarchyContainsWKWebView(baseView) ||
226                              base::FeatureList::IsEnabled(kSnapshotDrawView);
227  // |drawViewHierarchyInRect:| has undefined behavior when the view is not
228  // in the visible view hierarchy. In practice, when this method is called
229  // on a view that is part of view controller containment and not in the view
230  // hierarchy, an UIViewControllerHierarchyInconsistency exception will be
231  // thrown.
232  if (useDrawViewHierarchy && baseView.window) {
233    snapshotSuccess = [baseView drawViewHierarchyInRect:baseView.bounds
234                                     afterScreenUpdates:YES];
235  } else {
236    [[baseView layer] renderInContext:context];
237  }
238  UIImage* image = nil;
239  if (snapshotSuccess)
240    image = UIGraphicsGetImageFromCurrentImageContext();
241  UIGraphicsEndImageContext();
242  return image;
243}
244
245// Returns an image of the |baseImage| overlaid with |overlays| with the given
246// |frameInWindow|.
247- (UIImage*)snapshotWithOverlays:(NSArray<UIView*>*)overlays
248                       baseImage:(UIImage*)baseImage
249                   frameInWindow:(CGRect)frameInWindow {
250  DCHECK(!CGRectIsEmpty(frameInWindow));
251  if (!baseImage)
252    return nil;
253  // Note: If the baseImage scale differs from device scale, the baseImage size
254  // may slightly differ from frameInWindow size due to rounding. Do not attempt
255  // to compare the baseImage size and frameInWindow size.
256  if (overlays.count == 0)
257    return baseImage;
258  const CGFloat kScale =
259      std::max<CGFloat>(1.0, [self.snapshotCache snapshotScaleForDevice]);
260  UIGraphicsBeginImageContextWithOptions(frameInWindow.size, YES, kScale);
261  CGContext* context = UIGraphicsGetCurrentContext();
262  // The base image is already a cropped snapshot so it is drawn at the origin
263  // of the new image.
264  [baseImage drawAtPoint:CGPointZero];
265  // This shifts the origin of the context so that future drawings can be in
266  // window coordinates. For example, suppose that the desired snapshot area is
267  // at (0, 99) in the window coordinate space. Drawing at (0, 99) will appear
268  // as (0, 0) in the resulting image.
269  CGContextTranslateCTM(context, -frameInWindow.origin.x,
270                        -frameInWindow.origin.y);
271  [self drawOverlays:overlays context:context];
272  UIImage* snapshot = UIGraphicsGetImageFromCurrentImageContext();
273  UIGraphicsEndImageContext();
274  return snapshot;
275}
276
277// Updates the snapshot cache with |snapshot|.
278- (void)updateSnapshotCacheWithImage:(UIImage*)snapshot {
279  if (snapshot) {
280    [self.snapshotCache setImage:snapshot withSnapshotID:self.tabID];
281  } else {
282    // Remove any stale snapshot since the snapshot failed.
283    [self.snapshotCache removeImageWithSnapshotID:self.tabID];
284  }
285}
286
287// Draws |overlays| onto |context| at offsets relative to the window.
288- (void)drawOverlays:(NSArray<UIView*>*)overlays context:(CGContext*)context {
289  for (UIView* overlay in overlays) {
290    CGContextSaveGState(context);
291    CGRect frameInWindow = [overlay.superview convertRect:overlay.frame
292                                                   toView:nil];
293    // This shifts the context so that drawing starts at the overlay's offset.
294    CGContextTranslateCTM(context, frameInWindow.origin.x,
295                          frameInWindow.origin.y);
296    // |drawViewHierarchyInRect:| has undefined behavior when the view is not
297    // in the visible view hierarchy. In practice, when this method is called
298    // on a view that is part of view controller containment, an
299    // UIViewControllerHierarchyInconsistency exception will be thrown.
300    if (base::FeatureList::IsEnabled(kSnapshotDrawView) && overlay.window) {
301      // The rect's origin is ignored. Only size is used.
302      [overlay drawViewHierarchyInRect:overlay.bounds afterScreenUpdates:YES];
303    } else {
304      [[overlay layer] renderInContext:context];
305    }
306    CGContextRestoreGState(context);
307  }
308}
309
310// Retrieves information needed for snapshotting.
311- (SnapshotInfo)snapshotInfo {
312  SnapshotInfo snapshotInfo;
313  snapshotInfo.baseView = [self.delegate snapshotGenerator:self
314                                       baseViewForWebState:self.webState];
315  DCHECK(snapshotInfo.baseView);
316  UIEdgeInsets baseViewInsets = [self.delegate snapshotGenerator:self
317                                   snapshotEdgeInsetsForWebState:self.webState];
318  snapshotInfo.snapshotFrameInBaseView =
319      UIEdgeInsetsInsetRect(snapshotInfo.baseView.bounds, baseViewInsets);
320  DCHECK(!CGRectIsEmpty(snapshotInfo.snapshotFrameInBaseView));
321  snapshotInfo.snapshotFrameInWindow =
322      [snapshotInfo.baseView convertRect:snapshotInfo.snapshotFrameInBaseView
323                                  toView:nil];
324  DCHECK(!CGRectIsEmpty(snapshotInfo.snapshotFrameInWindow));
325  snapshotInfo.overlays = [self.delegate snapshotGenerator:self
326                               snapshotOverlaysForWebState:self.webState];
327  return snapshotInfo;
328}
329
330#pragma mark - CRWWebStateObserver
331
332- (void)webStateDestroyed:(web::WebState*)webState {
333  DCHECK_EQ(_webState, webState);
334  _webState->RemoveObserver(_webStateObserver.get());
335  _webStateObserver.reset();
336  _webState = nullptr;
337}
338
339@end
340