Support interception of mpc messages in sharded environment

[andyleiserson]

This will be used by the tests for malicious sharded shuffle.

[codecov]

Codecov Report

Attention: Patch coverage is 93.66197% with 9 lines in your changes missing coverage. Please review.

Please upload report for BASE (main@6d29275). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
ipa-core/src/helpers/transport/mod.rs	78.78%	7 Missing ⚠️
ipa-core/src/helpers/transport/in_memory/config.rs	92.30%	1 Missing ⚠️
ipa-core/src/test_fixture/world.rs	98.71%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1390   +/-   ##
=======================================
  Coverage        ?   93.67%           
=======================================
  Files           ?      223           
  Lines           ?    37737           
  Branches        ?        0           
=======================================
  Hits            ?    35351           
  Misses          ?     2386           
  Partials        ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[akoshelev]

I guess we don't need to intercept cross-shard messages unless we want to emulate network failures? Is there any other reason you want to corrupt them?

[andyleiserson]

I am considering three categories of messages:

1. Helper-to-helper messages in a non-sharded environment
2. Helper-to-helper messages in a sharded environment. These are similar to (1), but in order to uniquely identify a channel, the shard also needs to be specified. Step Foo, from shard 0 of H1 to shard 0 of H2, is a different channel from step Foo, from shard 1 of H1 to shard 1 of H2. This is the kind of message that I am adding support for in this PR.
3. Shard-to-shard messages in a sharded environment. i.e. step Bar from shard 0 of H1 to shard 1 of H1. I'm not intentionally targeting this case with these changes, although the situation may improve slightly. These messages aren't interesting to corrupt for the purpose of testing malicious security, because the shards of a helper fall within the same trust domain. Like you say, they might be interesting to corrupt to test a network failure, but that seems less important than testing the properties of malicious security protocols.

InspectContext::MpcMessage covers cases (1) and (2), depending on whether shard is None or Some(shard_index).

The reason I added InspectContext::ShardMessage, which is for case (3), even though it's not fully supported/tested, is to maintain roughly the same functionality of the code in <Weak<InMemoryTransport<I>> as Transport>::send. That code is used for all three categories of messages above. I also wanted to improve type safety for the dest member (previously Cow<'static, str>), to avoid potential confusion whether it should be interpreted as a helper ID or a shard index.

The MaliciousHelper utility is still restricted to cases (1) and (2). Since the goal of MaliciousHelper is to simplify writing tests that use the stream interceptor, it didn't seem appropriate to complicate the MaliciousHelper interface to try and support all three cases simultaneously. If we decide to write tests that fall under case (3), we can either code against the interceptor interface directly, or make a new MaliciousShard utility.

[cberkhoff]

Besides the shard index, the shard count is useful for some of the functionality I'm working on.

One option would be for me to expand ShardContext to also have the shard count but it will look very similar to Sharded. Is there any special additional meaning that ShardContext has that Sharded doesn't?

If they are interchangeable I could modify the code as follows:
fn context(&self) -> Sharded;
then that makes me wonder why have the context function to begin with?

If we don't need it why not just keep ShardBinding as it was, without the context function and have the following code?
impl ShardBinding for Sharded {}

[andyleiserson]

ShardContext is Option<ShardIndex>, it can represent either NotSharded or Sharded. Sharded is definitely sharded and contains a non-optional shard index (and the shard count).

What are you trying to do with ShardContext?

[cberkhoff]

Here's what I want to do:

I need shard transports (in particular the in-memory version) to know how many shards are in total. As an example, I use this functionality in the Prepare Query API (query processor) so that the leader shard can send messages to all other shards. Note there's no Context at this stage, these are the preliminary steps of the query.

My plan was to modify TransportConfig, part of InMemoryTransport. This config is now created by InMemoryMpcNetwork::with_stream_interceptor using the ShardContext.

[andyleiserson]

If you haven't already, it might be worth some time thinking about it from the other end -- how does trait Transport support determining the shard count? It seems related to fn identity, but doesn't directly fit into that.

I wonder if you can use HashMap::len on InMemoryTransport::connections to determine the shard count? (Or it might be something more like peer_count, since you could call it for an MPC transport as well, although there is little point in doing so, because it's always the same.)

I think it would be reasonable to change ShardContext to something like this:

enum TransportContext {
    // possibly shard count here, although I'm not sure you need it
    MpcNetwork { shard_index: Option<ShardIndex> }, 

    // possibly helper identity here, by analogy to MpcNetwork
    ShardNetwork { shard_count: ShardIndex },
}

I'm not sure how TransportContext relates to the existing identity and shard_index fields. Currently those fields are overloaded in that they have different meanings depending on whether the transport being described is a shard network or an MPC network. Maybe add a local_identity field to each of the TransportContext variants.

Another possibility here is to add a ShardBinding-like type parameter to InMemoryTransport. I don't think Sharded and NotSharded are going to work as-is though.

After this PR, there is a set of three functions on TransportConfigBuilder: with_sharding, bind_to_shard, and not_sharded. They should correspond in some manner to the enum variants above -- may make sense to rename these functions as well.