adding in public routing steps

Enclave/obliv.cpp

@@ -69,3 +69,76 @@ void obliv_pad_stg(uint8_t *buf, uint32_t msg_size,
         --tot_padding;
     }
 }
+
+// Determine the number of messages exceeding the maximum that can be sent to a 
+// storage server. Oblivious to the contents of tally vector.
+std::vector<uint32_t> obliv_excess_stg(std::vector<uint32_t> &tally,
+    nodenum_t num_storage_nodes, uint32_t msgs_per_stg)
+{
+    std::vector<uint32_t> excess(num_storage_nodes, 0);
+    for (nodenum_t i=0; i<num_storage_nodes; ++i) {
+        bool exceeds = tally[i] > msgs_per_stg;
+        uint32_t diff = tally[i] - msgs_per_stg; // nonsensical if !exceeds
+        excess[i] = oselect_uint32_t(0, diff, exceeds);
+    }
+
+    return excess;
+}
+
+// Determine the number of messages under the maximum that can be sent to a 
+// storage server. Oblivious to the contents of tally vector.
+std::vector<uint32_t> obliv_padding_stg(std::vector<uint32_t> &tally,
+    nodenum_t num_storage_nodes, uint32_t msgs_per_stg)
+{
+    std::vector<uint32_t> padding(num_storage_nodes, 0);
+    for (nodenum_t i=0; i<num_storage_nodes; ++i) {
+        bool under = tally[i] < msgs_per_stg;
+        uint32_t diff = msgs_per_stg - tally[i]; // nonsensical if !under
+        padding[i] = oselect_uint32_t(0, diff, under);
+    }
+
+    return padding;
+}
+
+// For each excess messages, convert into padding for nodes that will need some.
+// Oblivious to contents of excess, padding, and tally vectors. May modify
+// excess, padding, and tally vectors.
+void obliv_excess_to_padding(uint8_t *buf, uint32_t msg_size, uint32_t num_msgs,
+    std::vector<uint32_t> &excess, std::vector<uint32_t> &padding,
+    std::vector<uint32_t> &tally, nodenum_t num_storage_nodes)
+{
+    uint8_t *cur_msg = buf + ((num_msgs-1)*msg_size);
+    uint32_t pad_user = (1<<DEST_UID_BITS)-1;
+    for (uint32_t i=0; i<num_msgs; ++i) {
+        // Determine if storage node for current node has excess messages.
+        // Also, decrement excess count and tally if so.
+        uint32_t storage_node_id = (*(const uint32_t*)cur_msg) >> DEST_UID_BITS;
+        bool stg_node_excess = false;
+        for (uint32_t j=0; j<num_storage_nodes; ++j) {
+            bool at_msg_node = (storage_node_id == j);
+            bool cur_node_excess = (excess[j] > 0);
+            stg_node_excess = oselect_uint32_t(stg_node_excess,
+                cur_node_excess, at_msg_node);
+            excess[j] -= (at_msg_node & cur_node_excess);
+            tally[j] -= (at_msg_node & cur_node_excess);
+        }
+        // Find first node that needs padding. Also, decrement padding count
+        // and increment tally if current node has excess messages.
+        bool found_padding = false;
+        nodenum_t found_padding_node = 0;
+        for (uint32_t j=0; j<num_storage_nodes; ++j) {
+            bool found_padding_here = (!found_padding) & (!!padding[j]);
+            found_padding_node = oselect_uint32_t(found_padding_node, j,
+                found_padding_here);
+            found_padding = found_padding | found_padding_here;
+            padding[j] -= (found_padding_here & stg_node_excess);
+            tally[j] += (found_padding_here & stg_node_excess);
+        }
+        // Convert to padding if excess
+        uint32_t pad = ((found_padding_node<<DEST_UID_BITS) | pad_user);
+        *(uint32_t*)cur_msg = oselect_uint32_t(*(uint32_t*)cur_msg, pad,
+            stg_node_excess);
+        // Go to previous message for backwards iteration through messages
+        cur_msg -= msg_size;
+    }
+}

Enclave/obliv.hpp

@@ -27,4 +27,17 @@ std::vector<uint32_t> obliv_tally_stg(const uint8_t *buf,
 void obliv_pad_stg(uint8_t *buf, uint32_t msg_size,
     std::vector<uint32_t> &tally, uint32_t tot_padding);
 
+// Obliviously determine the number of messages exceeding the maximum that can
+// be sent to a storage server.
+std::vector<uint32_t> obliv_excess_stg(std::vector<uint32_t> &tally,
+    nodenum_t num_storage_nodes, uint32_t msgs_per_stg);
+
+// Obliviously determine the number of messages under the maximum that can
+// be sent to a storage server.
+std::vector<uint32_t> obliv_padding_stg(std::vector<uint32_t> &tally,
+    nodenum_t num_storage_nodes, uint32_t msgs_per_stg);
+
+void obliv_excess_to_padding(uint8_t *buf, uint32_t msg_size, uint32_t num_msgs,
+    std::vector<uint32_t> &excess, std::vector<uint32_t> &padding,
+    std::vector<uint32_t> &tally, nodenum_t num_storage_nodes);
 #endif

Enclave/route.cpp

@@ -50,7 +50,12 @@ bool route_init()
     // send at most m_priv_out messages.
     uint32_t users_per_ing = CEILDIV(g_teems_config.user_count,
         g_teems_config.num_ingestion_nodes);
-    uint32_t tot_msg_per_ing = users_per_ing * g_teems_config.m_priv_out;
+    uint32_t tot_msg_per_ing;
+    if (g_teems_config.private_routing) {
+        tot_msg_per_ing = users_per_ing * g_teems_config.m_priv_out;
+    } else {
+        tot_msg_per_ing = users_per_ing * g_teems_config.m_pub_out;
+    }
 
     // Compute the maximum number of messages we could receive in round 1
     // Each ingestion node will send us an our_weight/tot_weight
@@ -70,8 +75,12 @@ bool route_init()
         g_teems_config.num_storage_nodes);
 
     // And so can receive at most this many messages
-    uint32_t tot_msg_per_stg = users_per_stg *
-        g_teems_config.m_priv_in;
+    uint32_t tot_msg_per_stg;
+    if (g_teems_config.private_routing) {
+        tot_msg_per_stg = users_per_stg * g_teems_config.m_priv_in;
+    } else {
+        tot_msg_per_stg = users_per_stg * g_teems_config.m_pub_in;
+    }
 
     // Which will be at most this many from us
     uint32_t max_msg_to_each_stg = CEILDIV(tot_msg_per_stg,
@@ -612,6 +621,25 @@ void ecall_routing_proceed(void *cbpointer)
             printf_with_rtclock_diff(start_tally, "end tally (%u)\n", inserted);
 #endif
 
+            // For public routing, remove excess messages, making them padding
+            if (!g_teems_config.private_routing) {
+                // How many excess messages to remove per storage server
+                std::vector<uint32_t> excess = obliv_excess_stg(tally,
+                    num_storage_nodes, msgs_per_stg);
+                // How many padding messages to add per storage server
+                std::vector<uint32_t> padding = obliv_padding_stg(tally,
+                    num_storage_nodes, msgs_per_stg);
+                // Sort received messages by increasing storage node and
+                // priority. Smaller priority number indicates higher priority.
+                // Sorted messages are put back into source buffer.
+                sort_mtobliv<NidPriorityKey>(g_teems_config.nthreads,
+                    round1.buf, g_teems_config.msg_size, round1.inserted,
+                    round1.bufsize);
+                // Convert excess messages into padding
+                obliv_excess_to_padding(round1.buf, msg_size, round1.inserted,
+                    excess, padding, num_storage_nodes);
+            }
+
             // Note: tally contains private values!  It's OK to
             // non-obliviously check for an error condition, though.
             // While we're at it, obliviously change the tally of