1 // Copyright 2021 The Chromium Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #ifndef NET_TEST_EMBEDDED_TEST_SERVER_CONNECTION_TRACKER_H_ 6 #define NET_TEST_EMBEDDED_TEST_SERVER_CONNECTION_TRACKER_H_ 7 8 #include <stdint.h> 9 10 #include <map> 11 12 #include "base/memory/raw_ptr.h" 13 #include "base/run_loop.h" 14 #include "base/task/single_thread_task_runner.h" 15 #include "net/base/net_errors.h" 16 #include "net/socket/stream_socket.h" 17 #include "net/test/embedded_test_server/embedded_test_server.h" 18 #include "net/test/embedded_test_server/embedded_test_server_connection_listener.h" 19 20 namespace net::test_server { 21 22 // Keeps track of incoming connections being accepted or read from and exposes 23 // that info to the tests. 24 // A port being reused is currently considered an error. 25 // If a test needs to verify multiple connections are opened in sequence, that 26 // will need to be changed. 27 class ConnectionTracker { 28 public: 29 explicit ConnectionTracker(EmbeddedTestServer* test_server); 30 31 ~ConnectionTracker(); 32 33 ConnectionTracker(const ConnectionTracker&) = delete; 34 ConnectionTracker& operator=(const ConnectionTracker&) = delete; 35 36 // Returns the number of sockets that were accepted by the server. 37 size_t GetAcceptedSocketCount() const; 38 39 // Returns the number of sockets that were read from by the server. 40 size_t GetReadSocketCount() const; 41 42 // Waits until one connection is read. 43 void WaitUntilConnectionRead(); 44 45 // Waits until exactly `num_connections` have been made since construction or 46 // since ResetCounts() has been invoked. Fails if more connections are made. 47 // `num_connections` must be greater than 0. 48 void WaitForAcceptedConnections(size_t num_connections); 49 50 // Helper function to stop the waiting for sockets to be accepted for 51 // WaitForAcceptedConnections. |num_accepted_connections_loop_| spins 52 // until |num_accepted_connections_needed_| sockets are accepted by the test 53 // server. The values will be null/0 if the loop is not running. 54 void CheckAccepted(); 55 56 // This clears all state and counters. If any socket connected before 57 // `ResetCounts` is invoked is later read from, the test fails. 58 void ResetCounts(); 59 60 private: 61 // Gets notified by the EmbeddedTestServer on incoming connections being 62 // accepted or read from and transfers this information to ConnectionTracker. 63 class ConnectionListener : public EmbeddedTestServerConnectionListener { 64 public: 65 explicit ConnectionListener(ConnectionTracker* tracker); 66 67 ConnectionListener(const ConnectionListener&) = delete; 68 ConnectionListener operator=(const ConnectionListener&) = delete; 69 70 ~ConnectionListener() override; 71 72 // Gets called from the EmbeddedTestServer thread to be notified that 73 // a connection was accepted. 74 std::unique_ptr<net::StreamSocket> AcceptedSocket( 75 std::unique_ptr<net::StreamSocket> connection) override; 76 77 // Gets called from the EmbeddedTestServer thread to be notified that 78 // a connection was read from. 79 void ReadFromSocket(const net::StreamSocket& connection, int rv) override; 80 81 private: 82 // Task runner on which the connection tracker `tracker_` will be accessed. 83 scoped_refptr<base::SingleThreadTaskRunner> task_runner_; 84 85 // This pointer should be only accessed on the `task_runner_` thread. 86 raw_ptr<ConnectionTracker> tracker_; 87 }; 88 89 void AcceptedSocketWithPort(uint16_t port); 90 void ReadFromSocketWithPort(uint16_t port); 91 92 enum class SocketStatus { kAccepted, kReadFrom }; 93 94 ConnectionListener connection_listener_; 95 96 raw_ptr<base::RunLoop> read_loop_ = nullptr; 97 98 // Port -> SocketStatus. 99 using SocketContainer = std::map<uint16_t, SocketStatus>; 100 SocketContainer sockets_; 101 102 size_t num_connected_sockets_ = 0; 103 size_t num_read_sockets_ = 0; 104 105 // If |num_accepted_connections_needed_| is non zero, then the object is 106 // waiting for |num_accepted_connections_needed_| sockets to be accepted 107 // before quitting the |num_accepted_connections_loop_|. 108 size_t num_accepted_connections_needed_ = 0; 109 raw_ptr<base::RunLoop> num_accepted_connections_loop_ = nullptr; 110 }; 111 112 } // namespace net::test_server 113 114 #endif // NET_TEST_EMBEDDED_TEST_SERVER_SIMPLE_CONNECTION_TRACKER_H_ 115