From b64b928b534bd2472fd39d5effaf8b755d5c9d75 Mon Sep 17 00:00:00 2001
From: Mykola Mokhnach <nick+1@wire.com>
Date: Sat, 31 Mar 2018 22:03:44 +0200
Subject: [PATCH] Implement better pasteboard data verification

---
 WebDriverAgentLib/Utilities/FBPasteboard.m    | 22 ++++++++++++++++---
 .../IntegrationTests/FBPasteboardTests.m      | 13 +++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/WebDriverAgentLib/Utilities/FBPasteboard.m b/WebDriverAgentLib/Utilities/FBPasteboard.m
index aeb2e89e2..2e678b2ea 100644
--- a/WebDriverAgentLib/Utilities/FBPasteboard.m
+++ b/WebDriverAgentLib/Utilities/FBPasteboard.m
@@ -19,10 +19,26 @@ + (BOOL)setData:(NSData *)data forType:(NSString *)type error:(NSError **)error
   if ([type.lowercaseString isEqualToString:@"plaintext"]) {
     pb.string = [[NSString alloc] initWithData:data encoding:NSUTF8StringEncoding];
   } else if ([type.lowercaseString isEqualToString:@"image"]) {
-    pb.image = [UIImage imageWithData:data];
+    UIImage *image = [UIImage imageWithData:data];
+    if (nil == image) {
+      NSString *description = @"No image can be parsed from the given pasteboard data";
+      if (error) {
+        *error = [[FBErrorBuilder.builder withDescription:description] build];
+      }
+      return NO;
+    }
+    pb.image = image;
   } else if ([type.lowercaseString isEqualToString:@"url"]) {
     NSString *urlString = [[NSString alloc] initWithData:data encoding:NSUTF8StringEncoding];
-    pb.URL = [[NSURL alloc] initWithString:urlString];
+    NSURL *url = [[NSURL alloc] initWithString:urlString];
+    if (nil == url) {
+      NSString *description = @"No URL can be parsed from the given pasteboard data";
+      if (error) {
+        *error = [[FBErrorBuilder.builder withDescription:description] build];
+      }
+      return NO;
+    }
+    pb.URL = url;
   } else {
     NSString *description = [NSString stringWithFormat:@"Unsupported content type: %@", type];
     if (error) {
@@ -46,7 +62,7 @@ + (NSData *)dataForType:(NSString *)type error:(NSError **)error
     }
   } else if ([type.lowercaseString isEqualToString:@"url"]) {
     if (pb.hasURLs) {
-      return [NSData dataWithContentsOfURL:(NSURL *)pb.URL];
+      return [pb.URL.absoluteString dataUsingEncoding:NSUTF8StringEncoding];
     }
   } else {
     NSString *description = [NSString stringWithFormat:@"Unsupported content type: %@", type];
diff --git a/WebDriverAgentTests/IntegrationTests/FBPasteboardTests.m b/WebDriverAgentTests/IntegrationTests/FBPasteboardTests.m
index ca752950a..35f8a008f 100644
--- a/WebDriverAgentTests/IntegrationTests/FBPasteboardTests.m
+++ b/WebDriverAgentTests/IntegrationTests/FBPasteboardTests.m
@@ -69,6 +69,19 @@ - (void)testGetPasteboard
   XCTAssertEqualObjects(textField.value, [[NSString alloc] initWithData:result encoding:NSUTF8StringEncoding]);
 }
 
+- (void)testUrlCopyPaste
+{
+  NSString *urlString = @"http://appium.io?some=value";
+  NSError *error;
+  XCTAssertTrue([FBPasteboard setData:(NSData *)[urlString dataUsingEncoding:NSUTF8StringEncoding]
+                              forType:@"url"
+                              error:&error]);
+  XCTAssertNil(error);
+  NSData *result = [FBPasteboard dataForType:@"url" error:&error];
+  XCTAssertNil(error);
+  XCTAssertEqualObjects(urlString, [[NSString alloc] initWithData:result encoding:NSUTF8StringEncoding]);
+}
+
 @end