From a4c96130e64cec1ce8307e7e504300ab153a0652 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 29 Aug 2023 20:55:27 -0400
Subject: [PATCH] Fixing a couple security concerns in `raft-dask` nccl unique
 id generation (#1785)

Authors:
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Divye Gala (https://github.com/divyegala)

URL: https://github.com/rapidsai/raft/pull/1785
---
 cpp/include/raft/comms/std_comms.hpp       |  8 ++++----
 python/raft-dask/raft_dask/common/nccl.pyx | 12 ++++++------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/cpp/include/raft/comms/std_comms.hpp b/cpp/include/raft/comms/std_comms.hpp
index 165f721708..df1687f93d 100644
--- a/cpp/include/raft/comms/std_comms.hpp
+++ b/cpp/include/raft/comms/std_comms.hpp
@@ -143,17 +143,17 @@ void build_comms_nccl_ucx(
  * @}
  */
 
-inline void nccl_unique_id_from_char(ncclUniqueId* id, char* uniqueId, int size)
+inline void nccl_unique_id_from_char(ncclUniqueId* id, char* uniqueId)
 {
-  memcpy(id->internal, uniqueId, size);
+  memcpy(id->internal, uniqueId, NCCL_UNIQUE_ID_BYTES);
 }
 
-inline void get_unique_id(char* uid, int size)
+inline void get_nccl_unique_id(char* uid)
 {
   ncclUniqueId id;
   ncclGetUniqueId(&id);
 
-  memcpy(uid, id.internal, size);
+  memcpy(uid, id.internal, NCCL_UNIQUE_ID_BYTES);
 }
 };  // namespace comms
 };  // end namespace raft
diff --git a/python/raft-dask/raft_dask/common/nccl.pyx b/python/raft-dask/raft_dask/common/nccl.pyx
index a4d59610d3..a2802d8d17 100644
--- a/python/raft-dask/raft_dask/common/nccl.pyx
+++ b/python/raft-dask/raft_dask/common/nccl.pyx
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2020-2022, NVIDIA CORPORATION.
+# Copyright (c) 2020-2023, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -26,10 +26,9 @@ from libcpp cimport bool
 
 
 cdef extern from "raft/comms/std_comms.hpp" namespace "raft::comms":
-    void get_unique_id(char *uid, int size) except +
+    void get_nccl_unique_id(char *uid) except +
     void nccl_unique_id_from_char(ncclUniqueId *id,
-                                  char *uniqueId,
-                                  int size) except +
+                                  char *uniqueId) except +
 
 cdef extern from "nccl.h":
 
@@ -80,8 +79,9 @@ def unique_id():
     128-byte unique id : str
     """
     cdef char *uid = <char *> malloc(NCCL_UNIQUE_ID_BYTES * sizeof(char))
-    get_unique_id(uid, NCCL_UNIQUE_ID_BYTES)
+    get_nccl_unique_id(uid)
     c_str = uid[:NCCL_UNIQUE_ID_BYTES-1]
+    c_str
     free(uid)
     return c_str
 
@@ -132,7 +132,7 @@ cdef class nccl:
         self.rank = rank
 
         cdef ncclUniqueId *ident = <ncclUniqueId*>malloc(sizeof(ncclUniqueId))
-        nccl_unique_id_from_char(ident, commId, NCCL_UNIQUE_ID_BYTES)
+        nccl_unique_id_from_char(ident, commId)
 
         comm_ = <ncclComm_t*>self.comm