// Copyright 2013 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 "content/browser/websockets/websocket_manager.h"

#include <algorithm>
#include <string>
#include <vector>

#include "base/callback.h"
#include "base/logging.h"
#include "base/numerics/safe_conversions.h"
#include "base/rand_util.h"
#include "content/public/browser/browser_thread.h"
#include "content/public/browser/render_process_host.h"
#include "content/public/browser/render_process_host_observer.h"
#include "content/public/browser/storage_partition.h"
#include "services/service_manager/public/cpp/interface_registry.h"

namespace content {

namespace {

const char kWebSocketManagerKeyName[] = "web_socket_manager";

// Max number of pending connections per WebSocketManager used for per-renderer
// WebSocket throttling.
const int kMaxPendingWebSocketConnections = 255;

}  // namespace

class WebSocketManager::Handle : public base::SupportsUserData::Data,
                                 public RenderProcessHostObserver {
 public:
  explicit Handle(WebSocketManager* manager) : manager_(manager) {}

  ~Handle() override {
    DCHECK(!manager_) << "Should have received RenderProcessHostDestroyed";
  }

  WebSocketManager* manager() const { return manager_; }

  // The network stack could be shutdown after this notification, so be sure to
  // stop using it before then.
  void RenderProcessHostDestroyed(RenderProcessHost* host) override {
    DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));
    BrowserThread::DeleteSoon(BrowserThread::IO, FROM_HERE, manager_);
    manager_ = nullptr;
  }

 private:
  WebSocketManager* manager_;
};

// static
void WebSocketManager::CreateWebSocket(int process_id, int frame_id,
                                       blink::mojom::WebSocketRequest request) {
  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::UI));

  RenderProcessHost* host = RenderProcessHost::FromID(process_id);
  DCHECK(host);

  // Maintain a WebSocketManager per RenderProcessHost. While the instance of
  // WebSocketManager is allocated on the UI thread, it must only be used and
  // deleted from the IO thread.

  Handle* handle =
      static_cast<Handle*>(host->GetUserData(kWebSocketManagerKeyName));
  if (!handle) {
    handle = new Handle(
        new WebSocketManager(process_id, host->GetStoragePartition()));
    host->SetUserData(kWebSocketManagerKeyName, handle);
    host->AddObserver(handle);
  } else {
    DCHECK(handle->manager());
  }

  BrowserThread::PostTask(
      BrowserThread::IO,
      FROM_HERE,
      base::Bind(&WebSocketManager::DoCreateWebSocket,
                 base::Unretained(handle->manager()),
                 frame_id,
                 base::Passed(&request)));
}

WebSocketManager::WebSocketManager(int process_id,
                                   StoragePartition* storage_partition)
    : process_id_(process_id),
      storage_partition_(storage_partition),
      num_pending_connections_(0),
      num_current_succeeded_connections_(0),
      num_previous_succeeded_connections_(0),
      num_current_failed_connections_(0),
      num_previous_failed_connections_(0),
      context_destroyed_(false) {
  if (storage_partition_) {
    url_request_context_getter_ = storage_partition_->GetURLRequestContext();
    // This unretained pointer is safe because we destruct a WebSocketManager
    // only via WebSocketManager::Handle::RenderProcessHostDestroyed which
    // posts a deletion task to the IO thread.
    BrowserThread::PostTask(
        BrowserThread::IO,
        FROM_HERE,
        base::Bind(
            &WebSocketManager::ObserveURLRequestContextGetter,
            base::Unretained(this)));
  }
}

WebSocketManager::~WebSocketManager() {
  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));

  if (!context_destroyed_ && url_request_context_getter_)
    url_request_context_getter_->RemoveObserver(this);

  for (auto impl : impls_) {
    impl->GoAway();
    delete impl;
  }
}

void WebSocketManager::DoCreateWebSocket(
    int frame_id,
    blink::mojom::WebSocketRequest request) {
  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));

  if (num_pending_connections_ >= kMaxPendingWebSocketConnections) {
    // Too many websockets!
    request.ResetWithReason(
        blink::mojom::WebSocket::kInsufficientResources,
        "Error in connection establishment: net::ERR_INSUFFICIENT_RESOURCES");
    return;
  }
  if (context_destroyed_) {
    request.ResetWithReason(
        blink::mojom::WebSocket::kInsufficientResources,
        "Error in connection establishment: net::ERR_UNEXPECTED");
    return;
  }

  // Keep all WebSocketImpls alive until either the client drops its
  // connection (see OnLostConnectionToClient) or we need to shutdown.

  impls_.insert(CreateWebSocketImpl(this, std::move(request), process_id_,
                                    frame_id, CalculateDelay()));
  ++num_pending_connections_;

  if (!throttling_period_timer_.IsRunning()) {
    throttling_period_timer_.Start(
        FROM_HERE,
        base::TimeDelta::FromMinutes(2),
        this,
        &WebSocketManager::ThrottlingPeriodTimerCallback);
  }
}

// Calculate delay as described in the per-renderer WebSocket throttling
// design doc: https://goo.gl/tldFNn
base::TimeDelta WebSocketManager::CalculateDelay() const {
  int64_t f = num_previous_failed_connections_ +
              num_current_failed_connections_;
  int64_t s = num_previous_succeeded_connections_ +
              num_current_succeeded_connections_;
  int p = num_pending_connections_;
  return base::TimeDelta::FromMilliseconds(
      base::RandInt(1000, 5000) *
      (1 << std::min(p + f / (s + 1), INT64_C(16))) / 65536);
}

void WebSocketManager::ThrottlingPeriodTimerCallback() {
  num_previous_failed_connections_ = num_current_failed_connections_;
  num_current_failed_connections_ = 0;

  num_previous_succeeded_connections_ = num_current_succeeded_connections_;
  num_current_succeeded_connections_ = 0;

  if (num_pending_connections_ == 0 &&
      num_previous_failed_connections_ == 0 &&
      num_previous_succeeded_connections_ == 0) {
    throttling_period_timer_.Stop();
  }
}

WebSocketImpl* WebSocketManager::CreateWebSocketImpl(
    WebSocketImpl::Delegate* delegate,
    blink::mojom::WebSocketRequest request,
    int child_id,
    int frame_id,
    base::TimeDelta delay) {
  return new WebSocketImpl(delegate, std::move(request), child_id, frame_id,
                           delay);
}

int WebSocketManager::GetClientProcessId() {
  return process_id_;
}

StoragePartition* WebSocketManager::GetStoragePartition() {
  return storage_partition_;
}

void WebSocketManager::OnReceivedResponseFromServer(WebSocketImpl* impl) {
  // The server accepted this WebSocket connection.
  impl->OnHandshakeSucceeded();
  --num_pending_connections_;
  DCHECK_GE(num_pending_connections_, 0);
  ++num_current_succeeded_connections_;
}

void WebSocketManager::OnLostConnectionToClient(WebSocketImpl* impl) {
  // The client is no longer interested in this WebSocket.
  if (!impl->handshake_succeeded()) {
    // Update throttling counters (failure).
    --num_pending_connections_;
    DCHECK_GE(num_pending_connections_, 0);
    ++num_current_failed_connections_;
  }
  impl->GoAway();
  impls_.erase(impl);
  delete impl;
}

void WebSocketManager::OnContextShuttingDown() {
  context_destroyed_ = true;
  url_request_context_getter_ = nullptr;
  for (const auto& impl : impls_) {
    impl->GoAway();
    delete impl;
  }
  impls_.clear();
}

void WebSocketManager::ObserveURLRequestContextGetter() {
  DCHECK(BrowserThread::CurrentlyOn(BrowserThread::IO));
  if (!url_request_context_getter_->GetURLRequestContext()) {
    context_destroyed_ = true;
    url_request_context_getter_ = nullptr;
    return;
  }
  url_request_context_getter_->AddObserver(this);
}

}  // namespace content