fix: address PR #491 review issues (5 critical fixes)

Author: MC952-arch

flagcx/core/flagcx_p2p.cc

@@ -1642,6 +1642,14 @@ static int bootstrapExchangeDescTable(struct bootstrapState *bsState,
                         &remoteCount, sizeof(remoteCount)) != flagcxSuccess)
     return -1;
 
+  // Sanity check: reject absurdly large counts to prevent OOM or overflow
+  const uint32_t MAX_REMOTE_REGIONS = 65536;
+  if (remoteCount > MAX_REMOTE_REGIONS) {
+    WARN("bootstrapExchangeDescTable: remote count %u exceeds limit %u",
+         remoteCount, MAX_REMOTE_REGIONS);
+    return -1;
+  }
+
   std::vector<FlagcxP2pMemRegWire> remoteTable(remoteCount);
   if (bootstrapExchange(
           bsState, 0, 3, localTable.data(),
@@ -2647,15 +2655,15 @@ int flagcxP2pEngineGetMetadata(FlagcxP2pEngine *engine, char **metadataStr) {
   if (engine == NULL || metadataStr == NULL)
     return -1;
 
-  const int netDev = chooseEngineNetDev(engine);
-  if (engine->listeners[netDev].listenComm == NULL)
+  // After bootstrap P2P integration, metadata must expose the bootstrap listen
+  // port (used by flagcxP2pEngineConnect for the initial handshake), not the
+  // RDMA listen port (which is now exchanged during the bootstrap handshake).
+  if (engine->bsListenState == NULL || engine->bsListenPort <= 0)
     return -1;
 
-  FlagcxP2pListenHandleView *listenHandle =
-      reinterpret_cast<FlagcxP2pListenHandleView *>(
-          engine->listeners[netDev].handle);
-  const std::string rdmaAddr =
-      socketAddrToHostPortString(&listenHandle->connectAddr);
+  union flagcxSocketAddress bsAddr;
+  flagcxSocketGetAddr(&engine->bsListenState->p2p->sock, &bsAddr);
+  const std::string rdmaAddr = socketAddrToHostPortString(&bsAddr);
   if (rdmaAddr.empty())
     return -1;
 

flagcx/service/bootstrap.cc

@@ -174,6 +174,7 @@ static void *bootstrapRoot(void *rargs) {
   flagcxResult_t res = flagcxSuccess;
 
   /* Receive addresses from all ranks */
+  int checkedIn = 0;
   do {
     struct flagcxSocket sock;
     FLAGCXCHECKGOTO(flagcxSocketInit(&sock), res, fail);
@@ -184,12 +185,20 @@ static void *bootstrapRoot(void *rargs) {
     if (nranks == 0) {
       nranks = info.nranks;
       FLAGCXCHECKGOTO(flagcxCalloc(&rankInfo, nranks), res, fail);
+    } else if (info.nranks != nranks) {
+      WARN("Bootstrap Root: rank %d nranks mismatch: expected %d, got %d",
+           info.rank, nranks, info.nranks);
+      res = flagcxInvalidArgument;
+      goto fail;
     }
     if (info.rank < 0 || info.rank >= nranks) {
+      res = flagcxInvalidArgument;
       goto fail;
     }
+    if (rankInfo[info.rank].nranks == 0)
+      checkedIn++;
     rankInfo[info.rank] = info;
-  } while (rankInfo[nranks - 1].nranks == 0);
+  } while (checkedIn < nranks);
 
   /* Send everyone info about their "next" rank in the ring */
   for (int i = 0; i < nranks; ++i) {
@@ -234,6 +243,9 @@ flagcxResult_t bootstrapCollCreateRoot(struct flagcxBootstrapHandle *handle,
   args->magic = handle->magic;
   if (pthread_create(&thread, NULL, bootstrapRoot, (void *)args) != 0) {
     WARN("bootstrapCollCreateRoot: pthread_create failed");
+    free(args);
+    flagcxSocketClose(listenSock);
+    free(listenSock);
     return flagcxSystemError;
   }
   flagcxSetThreadName(thread, "FlagCX BootRoot");
@@ -889,11 +901,18 @@ bootstrapRingAllReduce(struct flagcxSocket *prevSocket,
                                          sendbuff, recvbuff, offset.data(),
                                          length.data(), datatype, op));
 
-  // AllGather the results
+  // AllGather the results with variable chunk sizes
   // Copy my chunk into correct position
-  memcpy(recvbuff + offset[rank], recvbuff, length[rank]);
-  FLAGCXCHECK(bootstrapRingAllGather(prevSocket, nextSocket, rank, nranks,
-                                     recvbuff, ChunkBytes));
+  memmove(recvbuff + offset[rank], recvbuff, length[rank]);
+
+  // Ring AllGather with per-slice variable sizes
+  for (int i = 0; i < nranks - 1; i++) {
+    size_t sslice = (rank - i + nranks) % nranks;
+    size_t rslice = (rank - i - 1 + nranks) % nranks;
+    FLAGCXCHECK(bootstrapNetSendRecv(
+        nextSocket, recvbuff + offset[sslice], (int)length[sslice], prevSocket,
+        recvbuff + offset[rslice], (int)length[rslice]));
+  }
   return flagcxSuccess;
 }
 
@@ -930,6 +949,9 @@ bootstrapRingReduce(struct bootstrapState *commState,
                                          sendbuff, recvbuff, offset.data(),
                                          length.data(), datatype, op));
 
+  // Move my reduced chunk to its final position before gather to root
+  memmove(recvbuff + offset[rank], recvbuff, length[rank]);
+
   // gather to root
   const int bootstrapTag = BOOTSTRAP_TAG_REDUCE;
   if (rank == root) {

flagcx/service/include/bootstrap.h

@@ -12,6 +12,7 @@ extern "C" {
 #endif
 
 #include "flagcx.h"
+#include "flagcx_net.h"
 #include "socket.h"
 
 struct flagcxBootstrapHandle {