chrome/browser/chromeos/policy/remote_commands/device_command_screenshot_job_unittest.cc

   1 // Copyright 2015 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 #include <map>
   6 #include <vector>
   7
   8 #include "ash/test/ash_test_base.h"
   9 #include "base/json/json_writer.h"
  10 #include "base/run_loop.h"
  11 #include "base/single_thread_task_runner.h"
  12 #include "base/test/test_mock_time_task_runner.h"
  13 #include "base/thread_task_runner_handle.h"
  14 #include "base/threading/sequenced_worker_pool.h"
  15 #include "base/time/time.h"
  16 #include "base/values.h"
  17 #include "chrome/browser/chromeos/policy/remote_commands/device_command_screenshot_job.h"
  18 #include "content/public/browser/browser_thread.h"
  19 #include "content/public/test/test_browser_thread_bundle.h"
  20 #include "policy/proto/device_management_backend.pb.h"
  21 #include "testing/gtest/include/gtest/gtest.h"
  22 #include "third_party/skia/include/core/SkBitmap.h"
  23 #include "ui/gfx/codec/png_codec.h"
  24
  25 namespace policy {
  26
  27 namespace em = enterprise_management;
  28
  29 namespace {
  30
  31 // String constant identifying the result field in the result payload.
  32 const char* const kResultFieldName = "result";
  33
  34 const char* const kMockUploadUrl = "http://example.com/upload";
  35
  36 const RemoteCommandJob::UniqueIDType kUniqueID = 123456789;
  37
  38 // String constant identifying the upload url field in the command payload.
  39 const char* const kUploadUrlFieldName = "fileUploadUrl";
  40
  41 em::RemoteCommand GenerateScreenshotCommandProto(
  42     RemoteCommandJob::UniqueIDType unique_id,
  43     base::TimeDelta age_of_command,
  44     const std::string upload_url) {
  45   em::RemoteCommand command_proto;
  46   command_proto.set_type(
  47       enterprise_management::RemoteCommand_Type_DEVICE_SCREENSHOT);
  48   command_proto.set_unique_id(unique_id);
  49   command_proto.set_age_of_command(age_of_command.InMilliseconds());
  50   std::string payload;
  51   base::DictionaryValue root_dict;
  52   root_dict.SetString(kUploadUrlFieldName, upload_url);
  53   base::JSONWriter::Write(root_dict, &payload);
  54   command_proto.set_payload(payload);
  55   return command_proto;
  56 }
  57
  58 class MockUploadJob : public policy::UploadJob {
  59  public:
  60   // If |error_code| is a null pointer OnSuccess() will be invoked when the
  61   // Start() method is called, otherwise OnFailure() will be invoked with the
  62   // respective |error_code|.
  63   MockUploadJob(const GURL& upload_url,
  64                 UploadJob::Delegate* delegate,
  65                 scoped_ptr<UploadJob::ErrorCode> error_code);
  66   ~MockUploadJob() override;
  67
  68   // policy::UploadJob:
  69   void AddDataSegment(const std::string& name,
  70                       const std::string& filename,
  71                       const std::map<std::string, std::string>& header_entries,
  72                       scoped_ptr<std::string> data) override;
  73   void Start() override;
  74
  75   const GURL& GetUploadUrl() const;
  76
  77  protected:
  78   const GURL upload_url_;
  79   UploadJob::Delegate* delegate_;
  80   scoped_ptr<UploadJob::ErrorCode> error_code_;
  81   bool add_datasegment_succeeds_;
  82 };
  83
  84 MockUploadJob::MockUploadJob(const GURL& upload_url,
  85                              UploadJob::Delegate* delegate,
  86                              scoped_ptr<UploadJob::ErrorCode> error_code)
  87     : upload_url_(upload_url),
  88       delegate_(delegate),
  89       error_code_(error_code.Pass()) {
  90 }
  91
  92 MockUploadJob::~MockUploadJob() {
  93 }
  94
  95 void MockUploadJob::AddDataSegment(
  96     const std::string& name,
  97     const std::string& filename,
  98     const std::map<std::string, std::string>& header_entries,
  99     scoped_ptr<std::string> data) {
 100 }
 101
 102 void MockUploadJob::Start() {
 103   DCHECK(delegate_);
 104   EXPECT_EQ(kMockUploadUrl, upload_url_.spec());
 105   if (error_code_) {
 106     base::ThreadTaskRunnerHandle::Get()->PostTask(
 107         FROM_HERE, base::Bind(&UploadJob::Delegate::OnFailure,
 108                               base::Unretained(delegate_), *error_code_));
 109     return;
 110   }
 111   base::ThreadTaskRunnerHandle::Get()->PostTask(
 112       FROM_HERE,
 113       base::Bind(&UploadJob::Delegate::OnSuccess, base::Unretained(delegate_)));
 114 }
 115
 116 scoped_refptr<base::RefCountedBytes> GenerateTestPNG(const int& width,
 117                                                      const int& height) {
 118   const SkColor background_color = SK_ColorBLUE;
 119   SkBitmap bmp;
 120   bmp.allocN32Pixels(width, height);
 121   for (int y = 0; y < height; ++y) {
 122     for (int x = 0; x < width; ++x)
 123       *bmp.getAddr32(x, y) = background_color;
 124   }
 125   SkAutoLockPixels lock(bmp);
 126   scoped_refptr<base::RefCountedBytes> png_bytes(new base::RefCountedBytes());
 127   gfx::PNGCodec::ColorFormat color_format = gfx::PNGCodec::FORMAT_RGBA;
 128   if (!gfx::PNGCodec::Encode(
 129           reinterpret_cast<const unsigned char*>(bmp.getPixels()), color_format,
 130           gfx::Size(bmp.width(), bmp.height()),
 131           static_cast<int>(bmp.rowBytes()), false,
 132           std::vector<gfx::PNGCodec::Comment>(), &png_bytes->data())) {
 133     LOG(ERROR) << "Failed to encode image";
 134   }
 135   return png_bytes;
 136 }
 137
 138 class MockScreenshotDelegate : public DeviceCommandScreenshotJob::Delegate {
 139  public:
 140   MockScreenshotDelegate(scoped_ptr<UploadJob::ErrorCode> upload_job_error_code,
 141                          bool screenshot_allowed);
 142   ~MockScreenshotDelegate() override;
 143
 144   bool IsScreenshotAllowed() override;
 145   void TakeSnapshot(
 146       gfx::NativeWindow window,
 147       const gfx::Rect& source_rect,
 148       const ui::GrabWindowSnapshotAsyncPNGCallback& callback) override;
 149   scoped_ptr<UploadJob> CreateUploadJob(const GURL&,
 150                                         UploadJob::Delegate*) override;
 151
 152  private:
 153   scoped_ptr<UploadJob::ErrorCode> upload_job_error_code_;
 154   bool screenshot_allowed_;
 155 };
 156
 157 MockScreenshotDelegate::MockScreenshotDelegate(
 158     scoped_ptr<UploadJob::ErrorCode> upload_job_error_code,
 159     bool screenshot_allowed)
 160     : upload_job_error_code_(upload_job_error_code.Pass()),
 161       screenshot_allowed_(screenshot_allowed) {
 162 }
 163
 164 MockScreenshotDelegate::~MockScreenshotDelegate() {
 165 }
 166
 167 bool MockScreenshotDelegate::IsScreenshotAllowed() {
 168   return screenshot_allowed_;
 169 }
 170
 171 void MockScreenshotDelegate::TakeSnapshot(
 172     gfx::NativeWindow window,
 173     const gfx::Rect& source_rect,
 174     const ui::GrabWindowSnapshotAsyncPNGCallback& callback) {
 175   const int width = source_rect.width();
 176   const int height = source_rect.height();
 177   scoped_refptr<base::RefCountedBytes> test_png =
 178       GenerateTestPNG(width, height);
 179   base::ThreadTaskRunnerHandle::Get()->PostTask(FROM_HERE,
 180                                                 base::Bind(callback, test_png));
 181 }
 182
 183 scoped_ptr<UploadJob> MockScreenshotDelegate::CreateUploadJob(
 184     const GURL& upload_url,
 185     UploadJob::Delegate* delegate) {
 186   return make_scoped_ptr(
 187       new MockUploadJob(upload_url, delegate, upload_job_error_code_.Pass()));
 188 }
 189
 190 }  // namespace
 191
 192 class DeviceCommandScreenshotTest : public ash::test::AshTestBase {
 193  public:
 194   void VerifyResults(RemoteCommandJob* job,
 195                      RemoteCommandJob::Status expected_status,
 196                      std::string expected_payload);
 197
 198  protected:
 199   DeviceCommandScreenshotTest();
 200
 201   // ash::test::AshTestBase:
 202   void SetUp() override;
 203
 204   void InitializeScreenshotJob(RemoteCommandJob* job,
 205                                RemoteCommandJob::UniqueIDType unique_id,
 206                                base::TimeTicks issued_time,
 207                                const std::string& upload_url);
 208
 209   std::string CreatePayloadFromResultCode(
 210       DeviceCommandScreenshotJob::ResultCode result_code);
 211
 212   base::RunLoop run_loop_;
 213   base::TimeTicks test_start_time_;
 214
 215  private:
 216   scoped_refptr<base::TestMockTimeTaskRunner> task_runner_;
 217
 218   DISALLOW_COPY_AND_ASSIGN(DeviceCommandScreenshotTest);
 219 };
 220
 221 DeviceCommandScreenshotTest::DeviceCommandScreenshotTest()
 222     : task_runner_(new base::TestMockTimeTaskRunner()) {
 223 }
 224
 225 void DeviceCommandScreenshotTest::SetUp() {
 226   ash::test::AshTestBase::SetUp();
 227   test_start_time_ = base::TimeTicks::Now();
 228 }
 229
 230 void DeviceCommandScreenshotTest::InitializeScreenshotJob(
 231     RemoteCommandJob* job,
 232     RemoteCommandJob::UniqueIDType unique_id,
 233     base::TimeTicks issued_time,
 234     const std::string& upload_url) {
 235   EXPECT_TRUE(job->Init(
 236       base::TimeTicks::Now(),
 237       GenerateScreenshotCommandProto(
 238           unique_id, base::TimeTicks::Now() - issued_time, upload_url)));
 239   EXPECT_EQ(unique_id, job->unique_id());
 240   EXPECT_EQ(RemoteCommandJob::NOT_STARTED, job->status());
 241 }
 242
 243 std::string DeviceCommandScreenshotTest::CreatePayloadFromResultCode(
 244     DeviceCommandScreenshotJob::ResultCode result_code) {
 245   std::string payload;
 246   base::DictionaryValue root_dict;
 247   if (result_code != DeviceCommandScreenshotJob::SUCCESS)
 248     root_dict.Set(kResultFieldName, new base::FundamentalValue(result_code));
 249   base::JSONWriter::Write(root_dict, &payload);
 250   return payload;
 251 }
 252
 253 void DeviceCommandScreenshotTest::VerifyResults(
 254     RemoteCommandJob* job,
 255     RemoteCommandJob::Status expected_status,
 256     std::string expected_payload) {
 257   EXPECT_EQ(expected_status, job->status());
 258   if (job->status() == RemoteCommandJob::SUCCEEDED) {
 259     scoped_ptr<std::string> payload = job->GetResultPayload();
 260     EXPECT_TRUE(payload);
 261     EXPECT_EQ(expected_payload, *payload);
 262   }
 263   run_loop_.Quit();
 264 }
 265
 266 TEST_F(DeviceCommandScreenshotTest, Success) {
 267   scoped_ptr<RemoteCommandJob> job(new DeviceCommandScreenshotJob(
 268       make_scoped_ptr(new MockScreenshotDelegate(nullptr, true))));
 269   InitializeScreenshotJob(job.get(), kUniqueID, test_start_time_,
 270                           kMockUploadUrl);
 271   bool success = job->Run(
 272       base::TimeTicks::Now(),
 273       base::Bind(
 274           &DeviceCommandScreenshotTest::VerifyResults, base::Unretained(this),
 275           base::Unretained(job.get()), RemoteCommandJob::SUCCEEDED,
 276           CreatePayloadFromResultCode(DeviceCommandScreenshotJob::SUCCESS)));
 277   EXPECT_TRUE(success);
 278   run_loop_.Run();
 279 }
 280
 281 TEST_F(DeviceCommandScreenshotTest, FailureUserInput) {
 282   scoped_ptr<RemoteCommandJob> job(new DeviceCommandScreenshotJob(
 283       make_scoped_ptr(new MockScreenshotDelegate(nullptr, false))));
 284   InitializeScreenshotJob(job.get(), kUniqueID, test_start_time_,
 285                           kMockUploadUrl);
 286   bool success =
 287       job->Run(base::TimeTicks::Now(),
 288                base::Bind(&DeviceCommandScreenshotTest::VerifyResults,
 289                           base::Unretained(this), base::Unretained(job.get()),
 290                           RemoteCommandJob::FAILED,
 291                           CreatePayloadFromResultCode(
 292                               DeviceCommandScreenshotJob::FAILURE_USER_INPUT)));
 293   EXPECT_TRUE(success);
 294   run_loop_.Run();
 295 }
 296
 297 TEST_F(DeviceCommandScreenshotTest, Failure) {
 298   using ErrorCode = UploadJob::ErrorCode;
 299   scoped_ptr<ErrorCode> error_code(
 300       new ErrorCode(UploadJob::AUTHENTICATION_ERROR));
 301   scoped_ptr<RemoteCommandJob> job(new DeviceCommandScreenshotJob(
 302       make_scoped_ptr(new MockScreenshotDelegate(error_code.Pass(), true))));
 303   InitializeScreenshotJob(job.get(), kUniqueID, test_start_time_,
 304                           kMockUploadUrl);
 305   bool success = job->Run(
 306       base::TimeTicks::Now(),
 307       base::Bind(&DeviceCommandScreenshotTest::VerifyResults,
 308                  base::Unretained(this), base::Unretained(job.get()),
 309                  RemoteCommandJob::FAILED,
 310                  CreatePayloadFromResultCode(
 311                      DeviceCommandScreenshotJob::FAILURE_AUTHENTICATION)));
 312   EXPECT_TRUE(success);
 313   run_loop_.Run();
 314 }
 315
 316 }  // namespace policy