libnbc: Fix int overflow when handling count parameters

 * In a reduce_scatter operation if the count array adds up to a
   value greater than INT_MAX then the count passed around is negative
   leading to an invalid buffer bring passed around often resulting in
   a segv crash.
 * The fix is to preserve the true count size as a `size_t` at all
   levels in the schedule (thus why there is a change to the protocol
   structures).
   - Instead of changing the count parameter of `ompi_op_reduce` we
     iterate over INT_MAX chunks of the buffer reducing each in turn.

Signed-off-by: Joshua Hursey <jhursey@us.ibm.com>

Author: jjhursey

ompi/mca/coll/libnbc/nbc.c

@@ -16,7 +16,7 @@
  * Author(s): Torsten Hoefler <htor@cs.indiana.edu>
  *
  * Copyright (c) 2012      Oracle and/or its affiliates.  All rights reserved.
- * Copyright (c) 2016      IBM Corporation.  All rights reserved.
+ * Copyright (c) 2016-2021 IBM Corporation.  All rights reserved.
  * Copyright (c) 2017      Ian Bradley Morgan and Anthony Skjellum. All
  *                         rights reserved.
  * Copyright (c) 2018      FUJITSU LIMITED.  All rights reserved.
@@ -119,7 +119,7 @@ static int nbc_schedule_round_append (NBC_Schedule *schedule, void *data, int da
 }
 
 /* this function puts a send into the schedule */
-static int NBC_Sched_send_internal (const void* buf, char tmpbuf, int count, MPI_Datatype datatype, int dest, bool local, NBC_Schedule *schedule, bool barrier) {
+static int NBC_Sched_send_internal (const void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int dest, bool local, NBC_Schedule *schedule, bool barrier) {
   NBC_Args_send send_args;
   int ret;
 
@@ -143,16 +143,16 @@ static int NBC_Sched_send_internal (const void* buf, char tmpbuf, int count, MPI
   return OMPI_SUCCESS;
 }
 
-int NBC_Sched_send (const void* buf, char tmpbuf, int count, MPI_Datatype datatype, int dest, NBC_Schedule *schedule, bool barrier) {
+int NBC_Sched_send (const void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int dest, NBC_Schedule *schedule, bool barrier) {
   return NBC_Sched_send_internal (buf, tmpbuf, count, datatype, dest, false, schedule, barrier);
 }
 
-int NBC_Sched_local_send (const void* buf, char tmpbuf, int count, MPI_Datatype datatype, int dest, NBC_Schedule *schedule, bool barrier) {
+int NBC_Sched_local_send (const void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int dest, NBC_Schedule *schedule, bool barrier) {
   return NBC_Sched_send_internal (buf, tmpbuf, count, datatype, dest, true, schedule, barrier);
 }
 
 /* this function puts a receive into the schedule */
-static int NBC_Sched_recv_internal (void* buf, char tmpbuf, int count, MPI_Datatype datatype, int source, bool local, NBC_Schedule *schedule, bool barrier) {
+static int NBC_Sched_recv_internal (void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int source, bool local, NBC_Schedule *schedule, bool barrier) {
   NBC_Args_recv recv_args;
   int ret;
 
@@ -176,16 +176,16 @@ static int NBC_Sched_recv_internal (void* buf, char tmpbuf, int count, MPI_Datat
   return OMPI_SUCCESS;
 }
 
-int NBC_Sched_recv (void* buf, char tmpbuf, int count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier) {
+int NBC_Sched_recv (void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier) {
   return NBC_Sched_recv_internal(buf, tmpbuf, count, datatype, source, false, schedule, barrier);
 }
 
-int NBC_Sched_local_recv (void* buf, char tmpbuf, int count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier) {
+int NBC_Sched_local_recv (void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier) {
   return NBC_Sched_recv_internal(buf, tmpbuf, count, datatype, source, true, schedule, barrier);
 }
 
 /* this function puts an operation into the schedule */
-int NBC_Sched_op (const void* buf1, char tmpbuf1, void* buf2, char tmpbuf2, int count, MPI_Datatype datatype,
+int NBC_Sched_op (const void* buf1, char tmpbuf1, void* buf2, char tmpbuf2, size_t count, MPI_Datatype datatype,
                   MPI_Op op, NBC_Schedule *schedule, bool barrier) {
   NBC_Args_op op_args;
   int ret;
@@ -212,7 +212,8 @@ int NBC_Sched_op (const void* buf1, char tmpbuf1, void* buf2, char tmpbuf2, int
 }
 
 /* this function puts a copy into the schedule */
-int NBC_Sched_copy (void *src, char tmpsrc, int srccount, MPI_Datatype srctype, void *tgt, char tmptgt, int tgtcount,
+int NBC_Sched_copy (void *src, char tmpsrc, size_t srccount, MPI_Datatype srctype,
+                    void *tgt, char tmptgt, size_t tgtcount,
                     MPI_Datatype tgttype, NBC_Schedule *schedule, bool barrier) {
   NBC_Args_copy copy_args;
   int ret;
@@ -240,7 +241,7 @@ int NBC_Sched_copy (void *src, char tmpsrc, int srccount, MPI_Datatype srctype,
 }
 
 /* this function puts a unpack into the schedule */
-int NBC_Sched_unpack (void *inbuf, char tmpinbuf, int count, MPI_Datatype datatype, void *outbuf, char tmpoutbuf,
+int NBC_Sched_unpack (void *inbuf, char tmpinbuf, size_t count, MPI_Datatype datatype, void *outbuf, char tmpoutbuf,
                       NBC_Schedule *schedule, bool barrier) {
   NBC_Args_unpack unpack_args;
   int ret;
@@ -534,7 +535,31 @@ static inline int NBC_Start_round(NBC_Handle *handle) {
         } else {
           buf2=opargs.buf2;
         }
-        ompi_op_reduce(opargs.op, buf1, buf2, opargs.count, opargs.datatype);
+
+        /* If the count is > INT_MAX then we need to call ompi_op_reduce()
+         * in iterations of counts <= INT_MAX since it has an `int count`
+         * parameter.
+         */
+        if( OPAL_UNLIKELY(opargs.count > INT_MAX) ) {
+          size_t done_count = 0, shift;
+          int iter_count;
+          ptrdiff_t ext, lb;
+
+          ompi_datatype_get_extent (opargs.datatype, &lb, &ext);
+
+          while(done_count < opargs.count) {
+            if( done_count + INT_MAX > opargs.count ) {
+              iter_count = opargs.count - done_count;
+            } else {
+              iter_count = INT_MAX;
+            }
+            shift = done_count * ext;
+            ompi_op_reduce(opargs.op, buf1 + shift, buf2 + shift, iter_count, opargs.datatype);
+            done_count += iter_count;
+          }
+        } else {
+          ompi_op_reduce(opargs.op, buf1, buf2, opargs.count, opargs.datatype);
+        }
         break;
       case COPY:
         NBC_DEBUG(5, "  COPY   (offset %li) ", offset);

ompi/mca/coll/libnbc/nbc_internal.h

@@ -15,6 +15,7 @@
  * Copyright (c) 2015      Los Alamos National Security, LLC. All rights
  *                         reserved.
  * Copyright (c) 2018      FUJITSU LIMITED.  All rights reserved.
+ * Copyright (c) 2021      IBM Corporation.  All rights reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -90,7 +91,7 @@ typedef enum {
 /* the send argument struct */
 typedef struct {
   NBC_Fn_type type;
-  int count;
+  size_t count;
   const void *buf;
   MPI_Datatype datatype;
   int dest;
@@ -101,7 +102,7 @@ typedef struct {
 /* the receive argument struct */
 typedef struct {
   NBC_Fn_type type;
-  int count;
+  size_t count;
   void *buf;
   MPI_Datatype datatype;
   char tmpbuf;
@@ -118,26 +119,26 @@ typedef struct {
   void *buf2;
   MPI_Op op;
   MPI_Datatype datatype;
-  int count;
+  size_t count;
 } NBC_Args_op;
 
 /* the copy argument struct */
 typedef struct {
   NBC_Fn_type type;
-  int srccount;
+  size_t srccount;
   void *src;
   void *tgt;
   MPI_Datatype srctype;
   MPI_Datatype tgttype;
-  int tgtcount;
+  size_t tgtcount;
   char tmpsrc;
   char tmptgt;
 } NBC_Args_copy;
 
 /* unpack operation arguments */
 typedef struct {
   NBC_Fn_type type;
-  int count;
+  size_t count;
   void *inbuf;
   void *outbuf;
   MPI_Datatype datatype;
@@ -146,15 +147,15 @@ typedef struct {
 } NBC_Args_unpack;
 
 /* internal function prototypes */
-int NBC_Sched_send (const void* buf, char tmpbuf, int count, MPI_Datatype datatype, int dest, NBC_Schedule *schedule, bool barrier);
-int NBC_Sched_local_send (const void* buf, char tmpbuf, int count, MPI_Datatype datatype, int dest,NBC_Schedule *schedule, bool barrier);
-int NBC_Sched_recv (void* buf, char tmpbuf, int count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier);
-int NBC_Sched_local_recv (void* buf, char tmpbuf, int count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier);
-int NBC_Sched_op (const void* buf1, char tmpbuf1, void* buf2, char tmpbuf2, int count, MPI_Datatype datatype,
+int NBC_Sched_send (const void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int dest, NBC_Schedule *schedule, bool barrier);
+int NBC_Sched_local_send (const void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int dest,NBC_Schedule *schedule, bool barrier);
+int NBC_Sched_recv (void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier);
+int NBC_Sched_local_recv (void* buf, char tmpbuf, size_t count, MPI_Datatype datatype, int source, NBC_Schedule *schedule, bool barrier);
+int NBC_Sched_op (const void* buf1, char tmpbuf1, void* buf2, char tmpbuf2, size_t count, MPI_Datatype datatype,
                   MPI_Op op, NBC_Schedule *schedule, bool barrier);
-int NBC_Sched_copy (void *src, char tmpsrc, int srccount, MPI_Datatype srctype, void *tgt, char tmptgt, int tgtcount,
+int NBC_Sched_copy (void *src, char tmpsrc, size_t srccount, MPI_Datatype srctype, void *tgt, char tmptgt, size_t tgtcount,
                     MPI_Datatype tgttype, NBC_Schedule *schedule, bool barrier);
-int NBC_Sched_unpack (void *inbuf, char tmpinbuf, int count, MPI_Datatype datatype, void *outbuf, char tmpoutbuf,
+int NBC_Sched_unpack (void *inbuf, char tmpinbuf, size_t count, MPI_Datatype datatype, void *outbuf, char tmpoutbuf,
                       NBC_Schedule *schedule, bool barrier);
 
 int NBC_Sched_barrier (NBC_Schedule *schedule);

ompi/mca/coll/libnbc/nbc_ireduce_scatter.c

@@ -45,7 +45,8 @@
 static int nbc_reduce_scatter_init(const void* sendbuf, void* recvbuf, const int *recvcounts, MPI_Datatype datatype,
                                    MPI_Op op, struct ompi_communicator_t *comm, ompi_request_t ** request,
                                    mca_coll_base_module_t *module, bool persistent) {
-  int peer, rank, maxr, p, res, count;
+  int peer, rank, maxr, p, res;
+  size_t count;
   MPI_Aint ext;
   ptrdiff_t gap, span, span_align;
   char *sbuf, inplace;
@@ -230,7 +231,8 @@ int ompi_coll_libnbc_ireduce_scatter (const void* sendbuf, void* recvbuf, const
 static int nbc_reduce_scatter_inter_init (const void* sendbuf, void* recvbuf, const int *recvcounts, MPI_Datatype datatype,
                                           MPI_Op op, struct ompi_communicator_t *comm, ompi_request_t ** request,
                                           mca_coll_base_module_t *module, bool persistent) {
-  int rank, res, count, lsize, rsize;
+  int rank, res, lsize, rsize;
+  size_t count;
   MPI_Aint ext;
   ptrdiff_t gap, span, span_align;
   NBC_Schedule *schedule;

ompi/mca/coll/libnbc/nbc_ireduce_scatter_block.c

@@ -43,7 +43,8 @@
 static int nbc_reduce_scatter_block_init(const void* sendbuf, void* recvbuf, int recvcount, MPI_Datatype datatype,
                                          MPI_Op op, struct ompi_communicator_t *comm, ompi_request_t ** request,
                                          mca_coll_base_module_t *module, bool persistent) {
-  int peer, rank, maxr, p, res, count;
+  int peer, rank, maxr, p, res;
+  size_t count;
   MPI_Aint ext;
   ptrdiff_t gap, span;
   char *redbuf, *sbuf, inplace;
@@ -229,7 +230,8 @@ int ompi_coll_libnbc_ireduce_scatter_block(const void* sendbuf, void* recvbuf, i
 static int nbc_reduce_scatter_block_inter_init(const void *sendbuf, void *recvbuf, int rcount, struct ompi_datatype_t *dtype,
                                                struct ompi_op_t *op, struct ompi_communicator_t *comm, ompi_request_t **request,
                                                mca_coll_base_module_t *module, bool persistent) {
-  int rank, res, count, lsize, rsize;
+  int rank, res, lsize, rsize;
+  size_t count;
   MPI_Aint ext;
   ptrdiff_t gap, span, span_align;
   NBC_Schedule *schedule;