reset existing (takeover) socket when rebinding to different address
MODIFY
quic/server/QuicServer.cpp
quic/server/QuicServer.cpp
@@ -298,7 +298,10 @@ folly::SocketAddress QuicServer::overrideTakeoverHandlerAddress(
worker->getEventBase()->runInEventBaseThreadAndWait([&] {
std::lock_guard<std::mutex> guard(startMutex_);
CHECK(initialized_);
- boundAddress = worker->overrideTakeoverHandlerAddress(addr);
+ auto workerEvb = worker->getEventBase();
+ auto localListenSocket = listenerSocketFactory_->make(workerEvb, -1);
+ boundAddress = worker->overrideTakeoverHandlerAddress(
+ std::move(localListenSocket), addr);
});
}
return boundAddress;

MODIFY
quic/server/QuicServerPacketRouter.cpp
quic/server/QuicServerPacketRouter.cpp
@@ -39,6 +39,20 @@ TakeoverHandlerCallback::~TakeoverHandlerCallback() {
void TakeoverHandlerCallback::bind(const folly::SocketAddress& addr) {
CHECK(socket_);
+ // first reset existing socket if any
+ socket_->bind(addr);
+ socket_->resumeRead(this);
+}
+
+void TakeoverHandlerCallback::rebind(
+ std::unique_ptr<folly::AsyncUDPSocket> socket,
+ const folly::SocketAddress& addr) {
+ if (socket_) {
+ // first reset existing socket if any
+ socket_->pauseRead();
+ socket_.reset();
+ }
+ socket_ = std::move(socket);
socket_->bind(addr);
socket_->resumeRead(this);
}

MODIFY
quic/server/QuicServerPacketRouter.h
quic/server/QuicServerPacketRouter.h
@@ -123,6 +123,14 @@ class TakeoverHandlerCallback : public folly::AsyncUDPSocket::ReadCallback,
void bind(const folly::SocketAddress& addr);
+ /**
+ * Rebinds the given socket to the given address
+ * Frees existing socket if any
+ */
+ void rebind(
+ std::unique_ptr<folly::AsyncUDPSocket> socket,
+ const folly::SocketAddress& addr);
+
void pause();
int getSocketFD();

MODIFY
quic/server/QuicServerWorker.cpp
quic/server/QuicServerWorker.cpp
@@ -520,9 +520,10 @@ void QuicServerWorker::allowBeingTakenOver(
}
const folly::SocketAddress& QuicServerWorker::overrideTakeoverHandlerAddress(
+ std::unique_ptr<folly::AsyncUDPSocket> socket,
const folly::SocketAddress& address) {
CHECK(takeoverCB_);
- takeoverCB_->bind(address);
+ takeoverCB_->rebind(std::move(socket), address);
return takeoverCB_->getAddress();
}

MODIFY
quic/server/QuicServerWorker.h
quic/server/QuicServerWorker.h
@@ -102,6 +102,7 @@ class QuicServerWorker : public folly::AsyncUDPSocket::ReadCallback,
* Returns const ref to SocketAddress representing the address it is bound to.
*/
const folly::SocketAddress& overrideTakeoverHandlerAddress(
+ std::unique_ptr<folly::AsyncUDPSocket> socket,
const folly::SocketAddress& address);
/**

MODIFY
quic/server/test/QuicServerTest.cpp
quic/server/test/QuicServerTest.cpp
@@ -788,6 +788,21 @@ class QuicServerWorkerTakeoverTest : public Test {
uint16_t clientHostId_{25};
};
+TEST_F(QuicServerWorkerTakeoverTest, QuicServerTakeoverReInitHandler) {
+ auto takeoverSock = std::make_unique<folly::test::MockAsyncUDPSocket>(&evb_);
+ folly::SocketAddress takeoverAddr;
+ EXPECT_CALL(*takeoverSocket_, pauseRead());
+
+ EXPECT_CALL(*takeoverSock, bind(_));
+ EXPECT_CALL(*takeoverSock, resumeRead(_));
+ EXPECT_CALL(*takeoverSock, address()).WillOnce(Invoke([&]() {
+ return takeoverAddr;
+ }));
+ takeoverWorker_->overrideTakeoverHandlerAddress(
+ std::move(takeoverSock), takeoverAddr);
+ takeoverSocket_ = takeoverSock.get();
+}
+
TEST_F(QuicServerWorkerTakeoverTest, QuicServerTakeoverNoForwarding) {
ConnectionId connId = createConnIdForServer(ProcessId::ONE),
clientConnId = getTestConnectionId(clientHostId_);
@@ -1369,6 +1384,24 @@ TEST_F(QuicServerTest, RouteDataFromDifferentThread) {
t.join();
}
+TEST_F(QuicServerTest, OverrideTakeoverAddressTest) {
+ folly::ScopedEventBaseThread evbThread;
+ std::vector<folly::EventBase*> evbs;
+ evbs.emplace_back(evbThread.getEventBase());
+ auto serverAddr = initializeServer(evbs);
+ folly::SocketAddress takeoverAddr("::1", 0);
+ server_->allowBeingTakenOver(takeoverAddr);
+ uint8_t bindAttempt = 0;
+ folly::SocketAddress boundAddr;
+ while (bindAttempt < 5) {
+ boundAddr = server_->overrideTakeoverHandlerAddress(takeoverAddr);
+ bindAttempt++;
+ }
+ EXPECT_TRUE(boundAddr.isInitialized());
+ std::thread t([&] { server_->shutdown(); });
+ t.join();
+}
+
class QuicServerTakeoverTest : public Test {
public:
void SetUp() override {

reset existing (takeover) socket when rebinding to different address

Reviewed By: afrind

Differential Revision: D18925452

fbshipit-source-id: dc431cbf899eb8e21a18d421a640e4f0649743cb