// Copyright 2015 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 <string>

#include "base/callback_forward.h"
#include "base/memory/weak_ptr.h"
#include "base/message_loop/message_loop.h"
#include "base/run_loop.h"
#include "base/strings/string_util.h"
#include "base/strings/stringprintf.h"
#include "base/thread_task_runner_handle.h"
#include "net/cookies/cookie_store.h"
#include "net/socket/socket_test_util.h"
#include "net/websockets/websocket_stream_create_test_base.h"
#include "net/websockets/websocket_test_util.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"

namespace net {
namespace {

using ::testing::TestWithParam;
using ::testing::ValuesIn;

const char kNoCookieHeader[] = "";

class TestBase : public WebSocketStreamCreateTestBase {
 public:
  void CreateAndConnect(const GURL& url,
                        const std::string& origin,
                        const std::string& cookie_header,
                        const std::string& response_body) {
    // We assume cookie_header ends with CRLF if not empty, as
    // WebSocketStandardRequestWithCookies requires. Use AddCRLFIfNotEmpty
    // in a call site.
    CHECK(cookie_header.empty() ||
          base::EndsWith(cookie_header, "\r\n", base::CompareCase::SENSITIVE));

    url_request_context_host_.SetExpectations(
        WebSocketStandardRequestWithCookies(url.path(), url.host(), origin,
                                            cookie_header, std::string()),
        response_body);
    CreateAndConnectStream(url.spec(), NoSubProtocols(), origin, nullptr);
  }

  std::string AddCRLFIfNotEmpty(const std::string& s) {
    return s.empty() ? s : s + "\r\n";
  }
};

struct ClientUseCookieParameter {
  // The URL for the WebSocket connection.
  const char* const url;
  // The URL for the previously set cookies.
  const char* const cookie_url;
  // The previously set cookies contents.
  const char* const cookie_line;
  // The Cookie: HTTP header expected to appear in the WS request. An empty
  // string means there is no Cookie: header.
  const char* const cookie_header;
};

class WebSocketStreamClientUseCookieTest
    : public TestBase,
      public TestWithParam<ClientUseCookieParameter> {
 public:
  ~WebSocketStreamClientUseCookieTest() override {
    // Permit any endpoint locks to be released.
    stream_request_.reset();
    stream_.reset();
    base::RunLoop().RunUntilIdle();
  }

  static void SetCookieHelperFunction(const base::Closure& task,
                                      base::WeakPtr<bool> weak_is_called,
                                      base::WeakPtr<bool> weak_result,
                                      bool success) {
    *weak_is_called = true;
    *weak_result = success;
    base::ThreadTaskRunnerHandle::Get()->PostTask(FROM_HERE, task);
  }
};

struct ServerSetCookieParameter {
  // The URL for the WebSocket connection.
  const char* const url;
  // The URL used to query cookies after the response received.
  const char* const cookie_url;
  // The cookies expected to appear for |cookie_url| inquiry.
  const char* const cookie_line;
  // The Set-Cookie: HTTP header attached to the response.
  const char* const cookie_header;
};

class WebSocketStreamServerSetCookieTest
    : public TestBase,
      public TestWithParam<ServerSetCookieParameter> {
 public:
  ~WebSocketStreamServerSetCookieTest() override {
    // Permit any endpoint locks to be released.
    stream_request_.reset();
    stream_.reset();
    base::RunLoop().RunUntilIdle();
  }

  static void GetCookiesHelperFunction(const base::Closure& task,
                                       base::WeakPtr<bool> weak_is_called,
                                       base::WeakPtr<std::string> weak_result,
                                       const std::string& cookies) {
    *weak_is_called = true;
    *weak_result = cookies;
    base::ThreadTaskRunnerHandle::Get()->PostTask(FROM_HERE, task);
  }
};

TEST_P(WebSocketStreamClientUseCookieTest, ClientUseCookie) {
  // For wss tests.
  ssl_data_.push_back(new SSLSocketDataProvider(ASYNC, OK));

  CookieStore* store =
      url_request_context_host_.GetURLRequestContext()->cookie_store();

  const GURL url(GetParam().url);
  const GURL cookie_url(GetParam().cookie_url);
  const std::string origin("http://www.example.com");
  const std::string cookie_line(GetParam().cookie_line);
  const std::string cookie_header(AddCRLFIfNotEmpty(GetParam().cookie_header));

  bool is_called = false;
  bool set_cookie_result = false;
  base::WeakPtrFactory<bool> weak_is_called(&is_called);
  base::WeakPtrFactory<bool> weak_set_cookie_result(&set_cookie_result);

  base::RunLoop run_loop;
  store->SetCookieWithOptionsAsync(
      cookie_url, cookie_line, CookieOptions(),
      base::Bind(&SetCookieHelperFunction, run_loop.QuitClosure(),
                 weak_is_called.GetWeakPtr(),
                 weak_set_cookie_result.GetWeakPtr()));
  run_loop.Run();
  ASSERT_TRUE(is_called);
  ASSERT_TRUE(set_cookie_result);

  CreateAndConnect(url, origin, cookie_header, WebSocketStandardResponse(""));
  WaitUntilConnectDone();
  EXPECT_FALSE(has_failed());
}

TEST_P(WebSocketStreamServerSetCookieTest, ServerSetCookie) {
  // For wss tests.
  ssl_data_.push_back(new SSLSocketDataProvider(ASYNC, OK));

  const GURL url(GetParam().url);
  const GURL cookie_url(GetParam().cookie_url);
  const std::string origin("http://www.example.com");
  const std::string cookie_line(GetParam().cookie_line);
  const std::string cookie_header(AddCRLFIfNotEmpty(GetParam().cookie_header));

  const std::string response = base::StringPrintf(
      "HTTP/1.1 101 Switching Protocols\r\n"
      "Upgrade: websocket\r\n"
      "Connection: Upgrade\r\n"
      "%s"
      "Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n"
      "\r\n",
      cookie_header.c_str());

  CookieStore* store =
      url_request_context_host_.GetURLRequestContext()->cookie_store();

  CreateAndConnect(url, origin, "", response);
  WaitUntilConnectDone();
  EXPECT_FALSE(has_failed());

  bool is_called = false;
  std::string get_cookies_result;
  base::WeakPtrFactory<bool> weak_is_called(&is_called);
  base::WeakPtrFactory<std::string> weak_get_cookies_result(
      &get_cookies_result);
  base::RunLoop run_loop;
  store->GetCookiesWithOptionsAsync(
      cookie_url, CookieOptions(),
      base::Bind(&GetCookiesHelperFunction, run_loop.QuitClosure(),
                 weak_is_called.GetWeakPtr(),
                 weak_get_cookies_result.GetWeakPtr()));
  run_loop.Run();
  EXPECT_TRUE(is_called);
  EXPECT_EQ(cookie_line, get_cookies_result);
}

// Test parameters definitions follow...

const ClientUseCookieParameter kClientUseCookieParameters[] = {
    // Non-secure cookies for ws
    {"ws://www.example.com",
     "http://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    {"ws://www.example.com",
     "https://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    {"ws://www.example.com",
     "ws://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    {"ws://www.example.com",
     "wss://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    // Non-secure cookies for wss
    {"wss://www.example.com",
     "http://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "https://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "ws://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "wss://www.example.com",
     "test-cookie",
     "Cookie: test-cookie"},

    // Secure-cookies for ws
    {"ws://www.example.com",
     "https://www.example.com",
     "test-cookie; secure",
     kNoCookieHeader},

    {"ws://www.example.com",
     "wss://www.example.com",
     "test-cookie; secure",
     kNoCookieHeader},

    // Secure-cookies for wss
    {"wss://www.example.com",
     "https://www.example.com",
     "test-cookie; secure",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "wss://www.example.com",
     "test-cookie; secure",
     "Cookie: test-cookie"},

    // Non-secure cookies for ws (sharing domain)
    {"ws://www.example.com",
     "http://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    {"ws://www.example.com",
     "https://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    {"ws://www.example.com",
     "ws://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    {"ws://www.example.com",
     "wss://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    // Non-secure cookies for wss (sharing domain)
    {"wss://www.example.com",
     "http://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "https://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "ws://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "wss://www2.example.com",
     "test-cookie; Domain=example.com",
     "Cookie: test-cookie"},

    // Secure-cookies for ws (sharing domain)
    {"ws://www.example.com",
     "https://www2.example.com",
     "test-cookie; Domain=example.com; secure",
     kNoCookieHeader},

    {"ws://www.example.com",
     "wss://www2.example.com",
     "test-cookie; Domain=example.com; secure",
     kNoCookieHeader},

    // Secure-cookies for wss (sharing domain)
    {"wss://www.example.com",
     "https://www2.example.com",
     "test-cookie; Domain=example.com; secure",
     "Cookie: test-cookie"},

    {"wss://www.example.com",
     "wss://www2.example.com",
     "test-cookie; Domain=example.com; secure",
     "Cookie: test-cookie"},

    // Non-matching cookies for ws
    {"ws://www.example.com",
     "http://www2.example.com",
     "test-cookie",
     kNoCookieHeader},

    {"ws://www.example.com",
     "https://www2.example.com",
     "test-cookie",
     kNoCookieHeader},

    {"ws://www.example.com",
     "ws://www2.example.com",
     "test-cookie",
     kNoCookieHeader},

    {"ws://www.example.com",
     "wss://www2.example.com",
     "test-cookie",
     kNoCookieHeader},

    // Non-matching cookies for wss
    {"wss://www.example.com",
     "http://www2.example.com",
     "test-cookie",
     kNoCookieHeader},

    {"wss://www.example.com",
     "https://www2.example.com",
     "test-cookie",
     kNoCookieHeader},

    {"wss://www.example.com",
     "ws://www2.example.com",
     "test-cookie",
     kNoCookieHeader},

    {"wss://www.example.com",
     "wss://www2.example.com",
     "test-cookie",
     kNoCookieHeader},
};

INSTANTIATE_TEST_CASE_P(WebSocketStreamClientUseCookieTest,
                        WebSocketStreamClientUseCookieTest,
                        ValuesIn(kClientUseCookieParameters));

const ServerSetCookieParameter kServerSetCookieParameters[] = {
    // Cookies coming from ws
    {"ws://www.example.com",
     "http://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    {"ws://www.example.com",
     "https://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    {"ws://www.example.com",
     "ws://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    {"ws://www.example.com",
     "wss://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    // Cookies coming from wss
    {"wss://www.example.com",
     "http://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    {"wss://www.example.com",
     "https://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    {"wss://www.example.com",
     "ws://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    {"wss://www.example.com",
     "wss://www.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie"},

    // cookies coming from ws (sharing domain)
    {"ws://www.example.com",
     "http://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    {"ws://www.example.com",
     "https://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    {"ws://www.example.com",
     "ws://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    {"ws://www.example.com",
     "wss://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    // cookies coming from wss (sharing domain)
    {"wss://www.example.com",
     "http://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    {"wss://www.example.com",
     "https://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    {"wss://www.example.com",
     "ws://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    {"wss://www.example.com",
     "wss://www2.example.com",
     "test-cookie",
     "Set-Cookie: test-cookie; Domain=example.com"},

    // Non-matching cookies coming from ws
    {"ws://www.example.com",
     "http://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},

    {"ws://www.example.com",
     "https://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},

    {"ws://www.example.com",
     "ws://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},

    {"ws://www.example.com",
     "wss://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},

    // Non-matching cookies coming from wss
    {"wss://www.example.com",
     "http://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},

    {"wss://www.example.com",
     "https://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},

    {"wss://www.example.com",
     "ws://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},

    {"wss://www.example.com",
     "wss://www2.example.com",
     "",
     "Set-Cookie: test-cookie"},
};

INSTANTIATE_TEST_CASE_P(WebSocketStreamServerSetCookieTest,
                        WebSocketStreamServerSetCookieTest,
                        ValuesIn(kServerSetCookieParameters));

}  // namespace
}  // namespace net