1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
|
// 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 "net/test/embedded_test_server/connection_tracker.h"
#include "base/containers/contains.h"
#include "base/run_loop.h"
#include "base/task/single_thread_task_runner.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::SingleThreadTaskRunner::GetCurrentDefault()),
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
|
