// Copyright 2021 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "services/network/websocket_interceptor.h"
#include "base/functional/bind.h"
#include "net/base/net_errors.h"
#include "services/network/throttling/throttling_controller.h"
#include "services/network/throttling/throttling_network_interceptor.h"

namespace network {

WebSocketInterceptor::WebSocketInterceptor(
    uint32_t net_log_source_id,
    const absl::optional<base::UnguessableToken>& throttling_profile_id)
    : net_log_source_id_(net_log_source_id),
      throttling_token_(
          network::ScopedThrottlingToken::MaybeCreate(net_log_source_id_,
                                                      throttling_profile_id)) {
  throttle_callbacks_[kIncoming] =
      base::BindRepeating(&WebSocketInterceptor::ThrottleCallback,
                          base::Unretained(this), kIncoming);
  throttle_callbacks_[kOutgoing] =
      base::BindRepeating(&WebSocketInterceptor::ThrottleCallback,
                          base::Unretained(this), kOutgoing);
}

WebSocketInterceptor::~WebSocketInterceptor() {
  auto* throttling_interceptor =
      ThrottlingController::GetInterceptor(net_log_source_id_);
  if (throttling_interceptor) {
    throttling_interceptor->StopThrottle(throttle_callbacks_[kIncoming]);
    throttling_interceptor->StopThrottle(throttle_callbacks_[kOutgoing]);
  }
}

WebSocketInterceptor::InterceptResult WebSocketInterceptor::Intercept(
    FrameDirection direction,
    size_t size,
    base::OnceClosure retry_callback) {
  DCHECK(!pending_callbacks_[direction]);

  auto* throttling_interceptor =
      ThrottlingController::GetInterceptor(net_log_source_id_);

  if (!throttling_interceptor)
    return kContinue;

  throttling_interceptor->SetSuspendWhenOffline(true);
  int start_throttle_result = throttling_interceptor->StartThrottle(
      /*result=*/0, size, /*send_end=*/base::TimeTicks(), /*start=*/false,
      /*is_upload=*/direction == kOutgoing, throttle_callbacks_[direction]);
  if (start_throttle_result == net::ERR_IO_PENDING) {
    pending_callbacks_[direction] = std::move(retry_callback);
    return kShouldWait;
  }
  return kContinue;
}

void WebSocketInterceptor::ThrottleCallback(FrameDirection direction,
                                            int result,
                                            int64_t bytes) {
  if (pending_callbacks_[direction])
    std::move(pending_callbacks_[direction]).Run();
}

}  // namespace network