blob: 09f1a059baad972a44216d77bb94d059e98d564d [file] [log] [blame]
// Copyright 2021 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 "net/test/embedded_test_server/connection_tracker.h"
#include "base/containers/contains.h"
#include "base/run_loop.h"
#include "net/test/embedded_test_server/embedded_test_server.h"
#include "testing/gtest/include/gtest/gtest.h"
namespace {
bool GetPort(const net::StreamSocket& connection, uint16_t* port) {
// Gets the remote port of the peer, since the local port will always be
// the port the test server is listening on. This isn't strictly correct -
// it's possible for multiple peers to connect with the same remote port
// but different remote IPs - but the tests here assume that connections
// to the test server (running on localhost) will always come from
// localhost, and thus the peer port is all that's needed to distinguish
// two connections. This also would be problematic if the OS reused ports,
// but that's not something to worry about for these tests.
net::IPEndPoint address;
int result = connection.GetPeerAddress(&address);
if (result != net::OK)
return false;
*port = address.port();
return true;
}
} // namespace
namespace net::test_server {
ConnectionTracker::ConnectionTracker(EmbeddedTestServer* test_server)
: connection_listener_(this) {
test_server->SetConnectionListener(&connection_listener_);
}
ConnectionTracker::~ConnectionTracker() = default;
void ConnectionTracker::AcceptedSocketWithPort(uint16_t port) {
num_connected_sockets_++;
sockets_[port] = SocketStatus::kAccepted;
CheckAccepted();
}
void ConnectionTracker::ReadFromSocketWithPort(uint16_t port) {
EXPECT_TRUE(base::Contains(sockets_, port));
if (sockets_[port] == SocketStatus::kAccepted)
num_read_sockets_++;
sockets_[port] = SocketStatus::kReadFrom;
if (read_loop_) {
read_loop_->Quit();
read_loop_ = nullptr;
}
}
// Returns the number of sockets that were accepted by the server.
size_t ConnectionTracker::GetAcceptedSocketCount() const {
return num_connected_sockets_;
}
// Returns the number of sockets that were read from by the server.
size_t ConnectionTracker::GetReadSocketCount() const {
return num_read_sockets_;
}
void ConnectionTracker::WaitUntilConnectionRead() {
base::RunLoop run_loop;
read_loop_ = &run_loop;
read_loop_->Run();
}
// This will wait for exactly |num_connections| items in |sockets_|. This method
// expects the server will not accept more than |num_connections| connections.
// |num_connections| must be greater than 0.
void ConnectionTracker::WaitForAcceptedConnections(size_t num_connections) {
DCHECK(!num_accepted_connections_loop_);
DCHECK_GT(num_connections, 0u);
base::RunLoop run_loop;
EXPECT_GE(num_connections, num_connected_sockets_);
num_accepted_connections_loop_ = &run_loop;
num_accepted_connections_needed_ = num_connections;
CheckAccepted();
// Note that the previous call to CheckAccepted can quit this run loop
// before this call, which will make this call a no-op.
run_loop.Run();
EXPECT_EQ(num_connections, num_connected_sockets_);
}
// Helper function to stop the waiting for sockets to be accepted for
// WaitForAcceptedConnections. |num_accepted_connections_loop_| spins
// until |num_accepted_connections_needed_| sockets are accepted by the test
// server. The values will be null/0 if the loop is not running.
void ConnectionTracker::CheckAccepted() {
// |num_accepted_connections_loop_| null implies
// |num_accepted_connections_needed_| == 0.
DCHECK(num_accepted_connections_loop_ ||
num_accepted_connections_needed_ == 0);
if (!num_accepted_connections_loop_ ||
num_accepted_connections_needed_ != num_connected_sockets_) {
return;
}
num_accepted_connections_loop_->Quit();
num_accepted_connections_needed_ = 0;
num_accepted_connections_loop_ = nullptr;
}
void ConnectionTracker::ResetCounts() {
sockets_.clear();
num_connected_sockets_ = 0;
num_read_sockets_ = 0;
}
ConnectionTracker::ConnectionListener::ConnectionListener(
ConnectionTracker* tracker)
: task_runner_(base::ThreadTaskRunnerHandle::Get()), tracker_(tracker) {}
ConnectionTracker::ConnectionListener::~ConnectionListener() = default;
// Gets called from the EmbeddedTestServer thread to be notified that
// a connection was accepted.
std::unique_ptr<net::StreamSocket>
ConnectionTracker::ConnectionListener::AcceptedSocket(
std::unique_ptr<net::StreamSocket> connection) {
uint16_t port;
if (GetPort(*connection, &port)) {
task_runner_->PostTask(
FROM_HERE, base::BindOnce(&ConnectionTracker::AcceptedSocketWithPort,
base::Unretained(tracker_), port));
}
return connection;
}
// Gets called from the EmbeddedTestServer thread to be notified that
// a connection was read from.
void ConnectionTracker::ConnectionListener::ReadFromSocket(
const net::StreamSocket& connection,
int rv) {
// Don't log a read if no data was transferred. This case often happens if
// the sockets of the test server are being flushed and disconnected.
if (rv <= 0)
return;
uint16_t port;
if (GetPort(connection, &port)) {
task_runner_->PostTask(
FROM_HERE, base::BindOnce(&ConnectionTracker::ReadFromSocketWithPort,
base::Unretained(tracker_), port));
}
}
} // namespace net::test_server