1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    protocols::{
        identity::exchange_handshake,
        wire::handshake::v1::{HandshakeMsg, MessagingProtocolVersion, SupportedProtocols},
    },
    testutils::fake_socket::ReadOnlyTestSocketVec,
};
use diem_config::network_id::NetworkId;
use diem_types::chain_id::ChainId;
use futures::executor::block_on;
use proptest::{collection::btree_map, prelude::*};

//
// Handshake Protocol Fuzzer
// =========================
//

/// Serializes a HandshakeMsg by simulating sending it over a socket
fn serialize_handshake_message(handshake_msg: &HandshakeMsg) -> Vec<u8> {
    // serialize with BCS
    let handshake_msg = bcs::to_bytes(handshake_msg).unwrap();
    // prepend a 2-byte prefix indicating the message length
    let mut serialized = (handshake_msg.len() as u16).to_be_bytes().to_vec();
    serialized.extend_from_slice(&handshake_msg);
    serialized
}

/// Fuzzing the handshake protocol, which negotiates protocols supported by both
/// the client and the server.
/// At the moment, fuzzing the client or the server leads to the same logic.
pub fn fuzz_network_handshake_protocol_exchange(self_handshake: &HandshakeMsg, data: Vec<u8>) {
    // fake socket to read the other peer's serialized HandshakeMsg from
    let mut fake_socket = ReadOnlyTestSocketVec::new(data);
    fake_socket.set_trailing();

    // fuzz the network exchange of HandshakeMsg first
    block_on(async move {
        if let Ok(remote_handshake_msg) = exchange_handshake(self_handshake, &mut fake_socket).await
        {
            // then perform the negotiation
            let _ = self_handshake.perform_handshake(&remote_handshake_msg);
        }
    });
}

/// Same function as fuzz_network_handshake_protocol_exchange except that the network exchange is skipped,
/// letting us skip BCS deserialization (and potentially other logic) and fuzz the negotiation of protocols directly.
pub fn fuzz_network_handshake_protocol_negotiation(
    self_handshake: &HandshakeMsg,
    remote_handshake: &HandshakeMsg,
) {
    let _ = self_handshake.perform_handshake(remote_handshake);
}

prop_compose! {
  /// Builds an arbitrary HandshakeMsg
  fn build_handshake_msg()(
    supported_protocols in btree_map(
      any::<MessagingProtocolVersion>(),
      any::<SupportedProtocols>(),
      0..5
    ),
  ) -> HandshakeMsg {
    HandshakeMsg {
      supported_protocols,
      chain_id: ChainId::new(1), // doesn't matter for handshake protocol
      network_id: NetworkId::Validator, // doesn't matter for handshake protocol
    }
  }
}

prop_compose! {
  /// Builds two HandshakeMsg and serializes the second one.
  /// It is the input expected for the fuzzer.
  pub fn exchange_handshake_input()(
    self_handshake in build_handshake_msg(),
    remote_handshake in build_handshake_msg(),
  ) -> (HandshakeMsg, Vec<u8>) {
    (self_handshake, serialize_handshake_message(&remote_handshake))
  }
}

prop_compose! {
  /// Builds two HandshakeMsg and serializes the second one.
  /// It is the input expected for the fuzzer.
  pub fn perform_handshake_input()(
    self_handshake in build_handshake_msg(),
    remote_handshake in build_handshake_msg(),
  ) -> (HandshakeMsg, HandshakeMsg) {
    (self_handshake, remote_handshake)
  }
}

proptest! {
  #[test]
  fn test_handshake_exchange_fuzzer((self_handshake, remote_handshake) in exchange_handshake_input()) {
    fuzz_network_handshake_protocol_exchange(&self_handshake, remote_handshake);
  }

  #[test]
  fn test_handshake_negotiation_fuzzer((self_handshake, remote_handshake) in perform_handshake_input()) {
    fuzz_network_handshake_protocol_negotiation(&self_handshake, &remote_handshake);
  }
}