// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <algorithm>

#include "base/pickle.h"
#include "base/time.h"
#include "chrome/browser/profiles/profile.h"
#include "chrome/browser/safe_browsing/malware_details.h"
#include "chrome/browser/safe_browsing/report.pb.h"
#include "chrome/common/render_messages.h"
#include "chrome/common/safe_browsing/safebrowsing_messages.h"
#include "chrome/test/test_url_request_context_getter.h"
#include "chrome/test/testing_profile.h"
#include "content/browser/browser_thread.h"
#include "content/browser/renderer_host/test_render_view_host.h"
#include "content/browser/tab_contents/navigation_entry.h"
#include "content/browser/tab_contents/test_tab_contents.h"
#include "net/base/io_buffer.h"
#include "net/base/test_completion_callback.h"
#include "net/disk_cache/disk_cache.h"
#include "net/http/http_cache.h"
#include "net/http/http_response_headers.h"
#include "net/http/http_response_info.h"
#include "net/http/http_util.h"
#include "net/url_request/url_request_context.h"
#include "net/url_request/url_request_context_getter.h"

static const char* kOriginalLandingURL = "http://www.originallandingpage.com/";
static const char* kHttpsURL = "https://www.url.com/";
static const char* kDOMChildURL = "http://www.domparent.com/";
static const char* kDOMParentURL = "http://www.domchild.com/";
static const char* kFirstRedirectURL = "http://redirectone.com/";
static const char* kSecondRedirectURL = "http://redirecttwo.com/";

static const char* kMalwareURL = "http://www.malware.com/";
static const char* kMalwareHeaders =
    "HTTP/1.1 200 OK\n"
    "Content-Type: image/jpeg\n";
static const char* kMalwareData = "exploit();";

static const char* kLandingURL = "http://www.landingpage.com/";
static const char* kLandingHeaders =
    "HTTP/1.1 200 OK\n"
    "Content-Type: text/html\n"
    "Content-Length: 1024\n"
    "Set-Cookie: tastycookie\n";  // This header is stripped.
static const char* kLandingData = "<iframe src='http://www.malware.com'>";

using safe_browsing::ClientMalwareReportRequest;

namespace {

void WriteHeaders(disk_cache::Entry* entry, const std::string headers) {
  net::HttpResponseInfo responseinfo;
  std::string raw_headers = net::HttpUtil::AssembleRawHeaders(
      headers.c_str(), headers.size());
  responseinfo.headers = new net::HttpResponseHeaders(raw_headers);

  Pickle pickle;
  responseinfo.Persist(&pickle, false, false);

  scoped_refptr<net::WrappedIOBuffer> buf(new net::WrappedIOBuffer(
      reinterpret_cast<const char*>(pickle.data())));
  int len = static_cast<int>(pickle.size());

  TestCompletionCallback cb;
  int rv = entry->WriteData(0, 0, buf, len, &cb, true);
  ASSERT_EQ(len, cb.GetResult(rv));
}

void WriteData(disk_cache::Entry* entry, const std::string data) {
  if (data.empty())
    return;

  int len = data.length();
  scoped_refptr<net::IOBuffer> buf(new net::IOBuffer(len));
  memcpy(buf->data(), data.data(), data.length());

  TestCompletionCallback cb;
  int rv = entry->WriteData(1, 0, buf, len, &cb, true);
  ASSERT_EQ(len, cb.GetResult(rv));
}

void WriteToEntry(disk_cache::Backend* cache, const std::string key,
                  const std::string headers, const std::string data) {
  TestCompletionCallback cb;
  disk_cache::Entry* entry;
  int rv = cache->CreateEntry(key, &entry, &cb);
  rv = cb.GetResult(rv);
  if (rv != net::OK) {
    rv = cache->OpenEntry(key, &entry, &cb);
    ASSERT_EQ(net::OK, cb.GetResult(rv));
  }

  WriteHeaders(entry, headers);
  WriteData(entry, data);

  entry->Close();
}

void FillCache(net::URLRequestContext* context) {
  TestCompletionCallback cb;
  disk_cache::Backend* cache;
  int rv =
      context->http_transaction_factory()->GetCache()->GetBackend(&cache, &cb);
  ASSERT_EQ(net::OK, cb.GetResult(rv));

  std::string empty;
  WriteToEntry(cache, kMalwareURL, kMalwareHeaders, kMalwareData);
  WriteToEntry(cache, kLandingURL, kLandingHeaders, kLandingData);
}

void QuitUIMessageLoop() {
  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
  BrowserThread::PostTask(BrowserThread::UI,
                          FROM_HERE,
                          new MessageLoop::QuitTask());
}

// Lets us provide a MockURLRequestContext with an HTTP Cache we pre-populate.
// Also exposes the constructor.
class MalwareDetailsWrap : public MalwareDetails {
 public:
  MalwareDetailsWrap(SafeBrowsingService* sb_service,
                     TabContents* tab_contents,
                     const SafeBrowsingService::UnsafeResource& unsafe_resource,
                     net::URLRequestContextGetter* request_context_getter)
      : MalwareDetails(sb_service, tab_contents, unsafe_resource) {
    request_context_getter_ = request_context_getter;
  }

  virtual ~MalwareDetailsWrap() {}
};

class MockSafeBrowsingService : public SafeBrowsingService {
 public:
  MockSafeBrowsingService() {}
  virtual ~MockSafeBrowsingService() {}

  // When the MalwareDetails is done, this is called.
  virtual void SendSerializedMalwareDetails(const std::string& serialized) {
    DVLOG(1) << "SendSerializedMalwareDetails";
    // Notify WaitForSerializedReport.
    BrowserThread::PostTask(BrowserThread::IO,
                            FROM_HERE,
                            NewRunnableFunction(&QuitUIMessageLoop));
    serialized_ = serialized;
  }

  const std::string& GetSerialized() {
    return serialized_;
  }

 private:
  std::string serialized_;
  DISALLOW_COPY_AND_ASSIGN(MockSafeBrowsingService);
};

}  // namespace.

class MalwareDetailsTest : public RenderViewHostTestHarness {
 public:
  MalwareDetailsTest()
      : ui_thread_(BrowserThread::UI, &message_loop_),
        io_thread_(BrowserThread::IO),
        sb_service_(new MockSafeBrowsingService()) {
  }

  virtual void SetUp() {
    RenderViewHostTestHarness::SetUp();
    // request_context_getter_ = new TestURLRequestContextGetter();

    // The URLFetcher checks that the messageloop type is IO.
    ASSERT_TRUE(io_thread_.StartWithOptions(
        base::Thread::Options(MessageLoop::TYPE_IO, 0)));
  }

  virtual void TearDown() {
    io_thread_.Stop();
    RenderViewHostTestHarness::TearDown();
  }

  static bool ResourceLessThan(
      const ClientMalwareReportRequest::Resource* lhs,
      const ClientMalwareReportRequest::Resource* rhs) {
    return lhs->id() < rhs->id();
  }

  std::string WaitForSerializedReport(MalwareDetails* report) {
    BrowserThread::PostTask(
        BrowserThread::IO,
        FROM_HERE,
        NewRunnableMethod(
            report, &MalwareDetails::FinishCollection));
    // Wait for the callback (SendSerializedMalwareDetails).
    DVLOG(1) << "Waiting for SendSerializedMalwareDetails";
    MessageLoop::current()->Run();
    return sb_service_->GetSerialized();
  }

 protected:
  void InitResource(SafeBrowsingService::UnsafeResource* resource,
                    ResourceType::Type resource_type,
                    const GURL& url) {
    resource->client = NULL;
    resource->url = url;
    resource->resource_type = resource_type;
    resource->threat_type = SafeBrowsingService::URL_MALWARE;
    resource->render_process_host_id = contents()->GetRenderProcessHost()->id();
    resource->render_view_id = contents()->render_view_host()->routing_id();
  }

  void VerifyResults(const ClientMalwareReportRequest& report_pb,
                     const ClientMalwareReportRequest& expected_pb) {
    EXPECT_EQ(expected_pb.malware_url(), report_pb.malware_url());
    EXPECT_EQ(expected_pb.page_url(), report_pb.page_url());
    EXPECT_EQ(expected_pb.referrer_url(), report_pb.referrer_url());

    ASSERT_EQ(expected_pb.resources_size(), report_pb.resources_size());
    // Sort the resources, to make the test deterministic
    std::vector<const ClientMalwareReportRequest::Resource*> resources;
    for (int i = 0; i < report_pb.resources_size(); ++i) {
      const ClientMalwareReportRequest::Resource& resource =
          report_pb.resources(i);
      resources.push_back(&resource);
    }
    std::sort(resources.begin(), resources.end(),
              &MalwareDetailsTest::ResourceLessThan);

    std::vector<const ClientMalwareReportRequest::Resource*> expected;
    for (int i = 0; i < report_pb.resources_size(); ++i) {
      const ClientMalwareReportRequest::Resource& resource =
          expected_pb.resources(i);
      expected.push_back(&resource);
    }
    std::sort(expected.begin(), expected.end(),
              &MalwareDetailsTest::ResourceLessThan);

    for (uint32 i = 0; i < expected.size(); ++i) {
      VerifyResource(resources[i], expected[i]);
    }

    EXPECT_EQ(expected_pb.complete(), report_pb.complete());
  }

  void VerifyResource(const ClientMalwareReportRequest::Resource* resource,
                      const ClientMalwareReportRequest::Resource* expected) {
    EXPECT_EQ(expected->id(), resource->id());
    EXPECT_EQ(expected->url(), resource->url());
    EXPECT_EQ(expected->parent_id(), resource->parent_id());
    ASSERT_EQ(expected->child_ids_size(), resource->child_ids_size());
    for (int i = 0; i < expected->child_ids_size(); i++) {
      EXPECT_EQ(expected->child_ids(i), resource->child_ids(i));
    }

    // Verify HTTP Responses
    if (expected->has_response()) {
      ASSERT_TRUE(resource->has_response());
      EXPECT_EQ(expected->response().firstline().code(),
                resource->response().firstline().code());

      ASSERT_EQ(expected->response().headers_size(),
                resource->response().headers_size());
      for (int i = 0; i < expected->response().headers_size(); ++i) {
        EXPECT_EQ(expected->response().headers(i).name(),
                  resource->response().headers(i).name());
        EXPECT_EQ(expected->response().headers(i).value(),
                  resource->response().headers(i).value());
      }

      EXPECT_EQ(expected->response().body(), resource->response().body());
      EXPECT_EQ(expected->response().bodylength(),
                resource->response().bodylength());
      EXPECT_EQ(expected->response().bodydigest(),
                resource->response().bodydigest());
    }
  }

  BrowserThread ui_thread_;
  BrowserThread io_thread_;
  scoped_refptr<MockSafeBrowsingService> sb_service_;
};

// Tests creating a simple malware report.
TEST_F(MalwareDetailsTest, MalwareSubResource) {
  // Start a load.
  controller().LoadURL(GURL(kLandingURL), GURL(), PageTransition::TYPED);

  SafeBrowsingService::UnsafeResource resource;
  InitResource(&resource, ResourceType::SUB_RESOURCE, GURL(kMalwareURL));

  scoped_refptr<MalwareDetailsWrap> report = new MalwareDetailsWrap(
      sb_service_, contents(), resource, NULL);

  std::string serialized = WaitForSerializedReport(report);

  ClientMalwareReportRequest actual;
  actual.ParseFromString(serialized);

  ClientMalwareReportRequest expected;
  expected.set_malware_url(kMalwareURL);
  expected.set_page_url(kLandingURL);
  expected.set_referrer_url("");

  ClientMalwareReportRequest::Resource* pb_resource = expected.add_resources();
  pb_resource->set_id(0);
  pb_resource->set_url(kLandingURL);
  pb_resource = expected.add_resources();
  pb_resource->set_id(1);
  pb_resource->set_url(kMalwareURL);

  VerifyResults(actual, expected);
}

// Tests creating a simple malware report where the subresource has a
// different original_url.
TEST_F(MalwareDetailsTest, MalwareSubResourceWithOriginalUrl) {
  controller().LoadURL(GURL(kLandingURL), GURL(), PageTransition::TYPED);

  SafeBrowsingService::UnsafeResource resource;
  InitResource(&resource, ResourceType::SUB_RESOURCE, GURL(kMalwareURL));
  resource.original_url = GURL(kOriginalLandingURL);

  scoped_refptr<MalwareDetailsWrap> report = new MalwareDetailsWrap(
      sb_service_.get(), contents(), resource, NULL);

  std::string serialized = WaitForSerializedReport(report);

  ClientMalwareReportRequest actual;
  actual.ParseFromString(serialized);

  ClientMalwareReportRequest expected;
  expected.set_malware_url(kMalwareURL);
  expected.set_page_url(kLandingURL);
  expected.set_referrer_url("");

  ClientMalwareReportRequest::Resource* pb_resource = expected.add_resources();
  pb_resource->set_id(0);
  pb_resource->set_url(kLandingURL);

  pb_resource = expected.add_resources();
  pb_resource->set_id(1);
  pb_resource->set_url(kOriginalLandingURL);

  pb_resource = expected.add_resources();
  pb_resource->set_id(2);
  pb_resource->set_url(kMalwareURL);
  // The Resource for kMmalwareUrl should have the Resource for
  // kOriginalLandingURL (with id 1) as parent.
  pb_resource->set_parent_id(1);

  VerifyResults(actual, expected);
}

// Tests creating a malware report with data from the renderer.
TEST_F(MalwareDetailsTest, MalwareDOMDetails) {
  controller().LoadURL(GURL(kLandingURL), GURL(), PageTransition::TYPED);

  SafeBrowsingService::UnsafeResource resource;
  InitResource(&resource, ResourceType::SUB_RESOURCE, GURL(kMalwareURL));

  scoped_refptr<MalwareDetailsWrap> report = new MalwareDetailsWrap(
      sb_service_.get(), contents(), resource, NULL);

  // Send a message from the DOM, with 2 nodes, a parent and a child.
  std::vector<SafeBrowsingHostMsg_MalwareDOMDetails_Node> params;
  SafeBrowsingHostMsg_MalwareDOMDetails_Node child_node;
  child_node.url = GURL(kDOMChildURL);
  child_node.tag_name = "iframe";
  child_node.parent = GURL(kDOMParentURL);
  params.push_back(child_node);
  SafeBrowsingHostMsg_MalwareDOMDetails_Node parent_node;
  parent_node.url = GURL(kDOMParentURL);
  parent_node.children.push_back(GURL(kDOMChildURL));
  params.push_back(parent_node);
  report->OnReceivedMalwareDOMDetails(params);

  MessageLoop::current()->RunAllPending();

  std::string serialized = WaitForSerializedReport(report);
  ClientMalwareReportRequest actual;
  actual.ParseFromString(serialized);

  ClientMalwareReportRequest expected;
  expected.set_malware_url(kMalwareURL);
  expected.set_page_url(kLandingURL);
  expected.set_referrer_url("");

  ClientMalwareReportRequest::Resource* pb_resource = expected.add_resources();
  pb_resource->set_id(0);
  pb_resource->set_url(kLandingURL);

  pb_resource = expected.add_resources();
  pb_resource->set_id(1);
  pb_resource->set_url(kMalwareURL);

  pb_resource = expected.add_resources();
  pb_resource->set_id(2);
  pb_resource->set_url(kDOMChildURL);
  pb_resource->set_parent_id(3);

  pb_resource = expected.add_resources();
  pb_resource->set_id(3);
  pb_resource->set_url(kDOMParentURL);
  pb_resource->add_child_ids(2);
  expected.set_complete(false);  // Since the cache was missing.

  VerifyResults(actual, expected);
}

// Verify that https:// urls are dropped.
TEST_F(MalwareDetailsTest, NotPublicUrl) {
  controller().LoadURL(GURL(kHttpsURL), GURL(), PageTransition::TYPED);
  SafeBrowsingService::UnsafeResource resource;
  InitResource(&resource, ResourceType::SUB_RESOURCE, GURL(kMalwareURL));
  scoped_refptr<MalwareDetailsWrap> report = new MalwareDetailsWrap(
      sb_service_.get(), contents(), resource, NULL);

  std::string serialized = WaitForSerializedReport(report);
  ClientMalwareReportRequest actual;
  actual.ParseFromString(serialized);

  ClientMalwareReportRequest expected;
  expected.set_malware_url(kMalwareURL);  // No page_url
  expected.set_referrer_url("");

  ClientMalwareReportRequest::Resource* pb_resource = expected.add_resources();
  pb_resource->set_url(kMalwareURL);  // Only one resource

  VerifyResults(actual, expected);
}

// Tests creating a malware report where there are redirect urls to an unsafe
// resource url
TEST_F(MalwareDetailsTest, MalwareWithRedirectUrl) {
  controller().LoadURL(GURL(kLandingURL), GURL(), PageTransition::TYPED);

  SafeBrowsingService::UnsafeResource resource;
  InitResource(&resource, ResourceType::SUB_RESOURCE, GURL(kMalwareURL));
  resource.original_url = GURL(kOriginalLandingURL);

  // add some redirect urls
  resource.redirect_urls.push_back(GURL(kFirstRedirectURL));
  resource.redirect_urls.push_back(GURL(kSecondRedirectURL));
  resource.redirect_urls.push_back(GURL(kMalwareURL));

  scoped_refptr<MalwareDetailsWrap> report = new MalwareDetailsWrap(
      sb_service_.get(), contents(), resource, NULL);

  std::string serialized = WaitForSerializedReport(report);
  ClientMalwareReportRequest actual;
  actual.ParseFromString(serialized);

  ClientMalwareReportRequest expected;
  expected.set_malware_url(kMalwareURL);
  expected.set_page_url(kLandingURL);
  expected.set_referrer_url("");

  ClientMalwareReportRequest::Resource* pb_resource = expected.add_resources();
  pb_resource->set_id(0);
  pb_resource->set_url(kLandingURL);

  pb_resource = expected.add_resources();
  pb_resource->set_id(1);
  pb_resource->set_url(kOriginalLandingURL);

  pb_resource = expected.add_resources();
  pb_resource->set_id(2);
  pb_resource->set_url(kMalwareURL);
  pb_resource->set_parent_id(4);

  pb_resource = expected.add_resources();
  pb_resource->set_id(3);
  pb_resource->set_url(kFirstRedirectURL);
  pb_resource->set_parent_id(1);

  pb_resource = expected.add_resources();
  pb_resource->set_id(4);
  pb_resource->set_url(kSecondRedirectURL);
  pb_resource->set_parent_id(3);

  VerifyResults(actual, expected);
}

// Tests the interaction with the HTTP cache.
TEST_F(MalwareDetailsTest, HTTPCache) {
  controller().LoadURL(GURL(kLandingURL), GURL(), PageTransition::TYPED);

  SafeBrowsingService::UnsafeResource resource;
  InitResource(&resource, ResourceType::SUB_RESOURCE, GURL(kMalwareURL));

  profile()->CreateRequestContext();
  scoped_refptr<MalwareDetailsWrap> report = new MalwareDetailsWrap(
      sb_service_.get(), contents(), resource
      , profile()->GetRequestContext());

  FillCache(profile()->GetRequestContext()->GetURLRequestContext());

  // The cache collection starts after the IPC from the DOM is fired.
  std::vector<SafeBrowsingHostMsg_MalwareDOMDetails_Node> params;
  report->OnReceivedMalwareDOMDetails(params);

  // Let the cache callbacks complete
  MessageLoop::current()->RunAllPending();

  DVLOG(1) << "Getting serialized report";
  std::string serialized = WaitForSerializedReport(report);
  ClientMalwareReportRequest actual;
  actual.ParseFromString(serialized);

  ClientMalwareReportRequest expected;
  expected.set_malware_url(kMalwareURL);
  expected.set_page_url(kLandingURL);
  expected.set_referrer_url("");

  ClientMalwareReportRequest::Resource* pb_resource = expected.add_resources();
  pb_resource->set_id(0);
  pb_resource->set_url(kLandingURL);
  safe_browsing::ClientMalwareReportRequest::HTTPResponse* pb_response =
      pb_resource->mutable_response();
  pb_response->mutable_firstline()->set_code(200);
  safe_browsing::ClientMalwareReportRequest::HTTPHeader* pb_header =
      pb_response->add_headers();
  pb_header->set_name("Content-Type");
  pb_header->set_value("text/html");
  pb_header = pb_response->add_headers();
  pb_header->set_name("Content-Length");
  pb_header->set_value("1024");
  pb_header = pb_response->add_headers();
  pb_header->set_name("Set-Cookie");
  pb_header->set_value("");  // The cookie is dropped.
  pb_response->set_body(kLandingData);
  pb_response->set_bodylength(37);
  pb_response->set_bodydigest("9ca97475598a79bc1e8fc9bd6c72cd35");

  pb_resource = expected.add_resources();
  pb_resource->set_id(1);
  pb_resource->set_url(kMalwareURL);
  pb_response = pb_resource->mutable_response();
  pb_response->mutable_firstline()->set_code(200);
  pb_header = pb_response->add_headers();
  pb_header->set_name("Content-Type");
  pb_header->set_value("image/jpeg");
  pb_response->set_body(kMalwareData);
  pb_response->set_bodylength(10);
  pb_response->set_bodydigest("581373551c43d4cf33bfb3b26838ff95");
  expected.set_complete(true);

  VerifyResults(actual, expected);
}

// Tests the interaction with the HTTP cache (where the cache is empty).
TEST_F(MalwareDetailsTest, HTTPCacheNoEntries) {
  controller().LoadURL(GURL(kLandingURL), GURL(), PageTransition::TYPED);

  SafeBrowsingService::UnsafeResource resource;
  InitResource(&resource, ResourceType::SUB_RESOURCE, GURL(kMalwareURL));

  profile()->CreateRequestContext();
  scoped_refptr<MalwareDetailsWrap> report = new MalwareDetailsWrap(
      sb_service_.get(), contents(), resource,
      profile()->GetRequestContext());

  // No call to FillCache

  // The cache collection starts after the IPC from the DOM is fired.
  std::vector<SafeBrowsingHostMsg_MalwareDOMDetails_Node> params;
  report->OnReceivedMalwareDOMDetails(params);

  // Let the cache callbacks complete
  MessageLoop::current()->RunAllPending();

  DVLOG(1) << "Getting serialized report";
  std::string serialized = WaitForSerializedReport(report);
  ClientMalwareReportRequest actual;
  actual.ParseFromString(serialized);

  ClientMalwareReportRequest expected;
  expected.set_malware_url(kMalwareURL);
  expected.set_page_url(kLandingURL);
  expected.set_referrer_url("");

  ClientMalwareReportRequest::Resource* pb_resource = expected.add_resources();
  pb_resource->set_id(0);
  pb_resource->set_url(kLandingURL);
  pb_resource = expected.add_resources();
  pb_resource->set_id(1);
  pb_resource->set_url(kMalwareURL);
  expected.set_complete(true);

  VerifyResults(actual, expected);
}