AMPI: Improve safety of function pointer shim transport

Author: evan-charmworks

src/libs/ck-libs/ampi/ampi_funcptr_loader.C

@@ -9,8 +9,11 @@
 #include <stdio.h>
 #include <string.h>
 
-void AMPI_FuncPtr_Pack(struct AMPI_FuncPtr_Transport * funcptrs)
+int AMPI_FuncPtr_Pack(struct AMPI_FuncPtr_Transport * funcptrs, size_t size)
 {
+  if (sizeof(*funcptrs) != size)
+    return 1;
+
 #define AMPI_CUSTOM_FUNC(return_type, function_name, ...) \
     funcptrs->function_name = &function_name;
 #if AMPI_HAVE_PMPI
@@ -27,6 +30,8 @@ void AMPI_FuncPtr_Pack(struct AMPI_FuncPtr_Transport * funcptrs)
 #undef AMPI_FUNC
 #undef AMPI_FUNC_NOIMPL
 #undef AMPI_CUSTOM_FUNC
+
+  return 0;
 }
 
 AMPI_FuncPtr_Unpack_t AMPI_FuncPtr_Unpack_Locate(SharedObject myexe)

src/libs/ck-libs/ampi/ampi_funcptr_loader.h

@@ -4,9 +4,9 @@
 #include "ampiimpl.h"
 #include "ampi_funcptr.h"
 
-void AMPI_FuncPtr_Pack(struct AMPI_FuncPtr_Transport *);
+int AMPI_FuncPtr_Pack(struct AMPI_FuncPtr_Transport *, size_t);
 
-typedef int (*AMPI_FuncPtr_Unpack_t)(const struct AMPI_FuncPtr_Transport *);
+typedef int (*AMPI_FuncPtr_Unpack_t)(const struct AMPI_FuncPtr_Transport *, size_t);
 AMPI_FuncPtr_Unpack_t AMPI_FuncPtr_Unpack_Locate(SharedObject);
 
 #endif /* AMPI_FUNCPTR_LOADER_H_ */

src/libs/ck-libs/ampi/ampi_funcptr_loader_stub.C

@@ -1,13 +1,12 @@
 
 #include "ampi_funcptr_loader.h"
 
-void AMPI_FuncPtr_Pack(struct AMPI_FuncPtr_Transport * funcptrs)
+int AMPI_FuncPtr_Pack(struct AMPI_FuncPtr_Transport *, size_t)
 {
-  (void)funcptrs;
+  return 0;
 }
 
-AMPI_FuncPtr_Unpack_t AMPI_FuncPtr_Unpack_Locate(SharedObject myexe)
+AMPI_FuncPtr_Unpack_t AMPI_FuncPtr_Unpack_Locate(SharedObject)
 {
-  (void)myexe;
   return nullptr;
 }

src/libs/ck-libs/ampi/ampi_funcptr_shim.C

@@ -21,8 +21,11 @@
 
 // Provide an interface to link the function pointers at runtime.
 
-extern "C" CMI_EXPORT void AMPI_FuncPtr_Unpack(const struct AMPI_FuncPtr_Transport * funcptrs)
+extern "C" CMI_EXPORT int AMPI_FuncPtr_Unpack(const struct AMPI_FuncPtr_Transport * funcptrs, size_t size)
 {
+  if (sizeof(*funcptrs) != size)
+    return 1;
+
 #define AMPI_CUSTOM_FUNC(return_type, function_name, ...) \
   function_name = funcptrs->function_name;
 #if AMPI_HAVE_PMPI
@@ -39,4 +42,6 @@ extern "C" CMI_EXPORT void AMPI_FuncPtr_Unpack(const struct AMPI_FuncPtr_Transpo
 #undef AMPI_FUNC
 #undef AMPI_FUNC_NOIMPL
 #undef AMPI_CUSTOM_FUNC
+
+  return 0;
 }

src/libs/ck-libs/ampi/ampi_globals_fs.C

@@ -84,7 +84,8 @@ void AMPI_Node_Setup(int numranks)
     CmiPrintf("AMPI> Using fsglobals privatization method.\n");
 
   AMPI_FuncPtr_Transport funcptrs{};
-  AMPI_FuncPtr_Pack(&funcptrs);
+  if (AMPI_FuncPtr_Pack(&funcptrs, sizeof(funcptrs)))
+    CkAbort("Globals runtime linking pack failed due to mismatch!");
 
   static const char exe_suffix[] = STRINGIFY(CMK_POST_EXE);
   static const char user_suffix[] = STRINGIFY(CMK_USER_SUFFIX);
@@ -124,7 +125,10 @@ void AMPI_Node_Setup(int numranks)
 
     auto unpack = AMPI_FuncPtr_Unpack_Locate(myexe);
     if (unpack != nullptr)
-      unpack(&funcptrs);
+    {
+      if (unpack(&funcptrs, sizeof(funcptrs)))
+        CkAbort("Globals runtime linking unpack failed due to mismatch!");
+    }
 
     rankdata[myrank].exe = myexe;
     rankdata[myrank].mainstruct = AMPI_Main_Get(myexe);

src/libs/ck-libs/ampi/ampi_globals_pip.C

@@ -19,7 +19,8 @@ void AMPI_Node_Setup(int numranks)
     CmiPrintf("AMPI> Using pipglobals privatization method.\n");
 
   AMPI_FuncPtr_Transport funcptrs{};
-  AMPI_FuncPtr_Pack(&funcptrs);
+  if (AMPI_FuncPtr_Pack(&funcptrs, sizeof(funcptrs)))
+    CkAbort("Globals runtime linking pack failed due to mismatch!");
 
   static const char exe_suffix[] = STRINGIFY(CMK_POST_EXE);
   static const char suffix[] = STRINGIFY(CMK_USER_SUFFIX) "." STRINGIFY(CMK_SHARED_SUF);
@@ -49,7 +50,10 @@ void AMPI_Node_Setup(int numranks)
 
     auto unpack = AMPI_FuncPtr_Unpack_Locate(myexe);
     if (unpack != nullptr)
-      unpack(&funcptrs);
+    {
+      if (unpack(&funcptrs, sizeof(funcptrs)))
+        CkAbort("Globals runtime linking unpack failed due to mismatch!");
+    }
 
     rankdata[myrank].exe = myexe;
     rankdata[myrank].mainstruct = AMPI_Main_Get(myexe);