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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

use bytes::Bytes;
use chain::{
    sync_packet::{PacketInfo, SyncPacket},
    ChainSync, PeerInfo,
};
use network::PeerId;
use rlp::{DecoderError, Rlp, RlpStream};

pub type RequestId = u64;

// Separate the eth/66 request id from a packet, if it exists.
pub fn strip_request_id<'a>(
    data: &'a [u8],
    sync: &ChainSync,
    peer: &PeerId,
    packet_id: &SyncPacket,
) -> Result<(Rlp<'a>, Option<RequestId>), DecoderError> {
    let protocol_version = if let Some(peer_info) = sync.peers.get(peer) {
        peer_info.protocol_version
    } else {
        trace!(
            "Peer info missing for peer {}, assuming protocol version 66",
            peer
        );
        66
    };

    let has_request_id = protocol_version >= 66 && packet_id.has_request_id_in_eth_66();

    do_strip_request_id(data, has_request_id)
}

fn do_strip_request_id<'a>(
    data: &'a [u8],
    has_request_id: bool,
) -> Result<(Rlp<'a>, Option<RequestId>), DecoderError> {
    let rlp = Rlp::new(data);

    if has_request_id {
        let request_id: RequestId = rlp.val_at(0)?;
        let stripped_rlp = rlp.at(1)?;
        Ok((stripped_rlp, Some(request_id)))
    } else {
        Ok((rlp, None))
    }
}

// Add a given eth/66 request id to a packet being built.
pub fn prepend_request_id(rlp: RlpStream, request_id: Option<RequestId>) -> RlpStream {
    match request_id {
        Some(ref id) => {
            let mut stream = RlpStream::new_list(2);
            stream.append(id);
            stream.append_raw(&rlp.out(), 1);
            stream
        }
        None => rlp,
    }
}

/// Prepend a new eth/66 request id to the packet if appropriate.
pub fn generate_request_id(
    packet: Bytes,
    peer: &PeerInfo,
    packet_id: SyncPacket,
) -> (Bytes, Option<RequestId>) {
    if peer.protocol_version >= 66 && packet_id.has_request_id_in_eth_66() {
        do_generate_request_id(&packet)
    } else {
        (packet, None)
    }
}

fn do_generate_request_id(packet: &Bytes) -> (Bytes, Option<RequestId>) {
    let request_id: RequestId = rand::random();

    let mut rlp = RlpStream::new_list(2);
    rlp.append(&request_id);
    rlp.append_raw(packet, 1);

    (rlp.out(), Some(request_id))
}

#[cfg(test)]
mod tests {
    use super::*;
    use ethereum_types::H256;

    #[test]
    fn test_prepend_request_id() {
        let mut request = RlpStream::new_list(2);
        request.append(&H256::from_low_u64_be(1));
        request.append(&H256::from_low_u64_be(2));

        let with_id = prepend_request_id(request, Some(10));
        let rlp = Rlp::new(with_id.as_raw());
        let recovered_id: RequestId = rlp.val_at(0).unwrap();
        let recovered_request: Vec<H256> = rlp.at(1).unwrap().as_list().unwrap();

        assert_eq!(recovered_id, 10);
        assert_eq!(
            recovered_request,
            [H256::from_low_u64_be(1), H256::from_low_u64_be(2)]
        );
    }

    #[test]
    fn test_strip_request_id() {
        let request = vec![
            H256::from_low_u64_be(1),
            H256::from_low_u64_be(2),
            H256::from_low_u64_be(3),
        ];

        let mut request_with_id = RlpStream::new_list(2);
        request_with_id.append(&20u64);
        request_with_id.append_list(&request);
        let data = request_with_id.out();

        let (rlp, id) = do_strip_request_id(&data, true).unwrap();

        assert_eq!(id, Some(20));
        assert_eq!(rlp.as_list::<H256>().unwrap(), request);
    }

    #[test]
    fn test_generate_request_id() {
        let request = vec![
            H256::from_low_u64_be(1),
            H256::from_low_u64_be(2),
            H256::from_low_u64_be(3),
        ];

        let mut stream = RlpStream::new_list(3);
        for hash in &request {
            stream.append(hash);
        }
        let data = stream.out();

        let (new_data, id) = do_generate_request_id(&data);

        let recovered = Rlp::new(&new_data);
        let recovered_id: RequestId = recovered.val_at(0).unwrap();
        let recovered_request: Vec<H256> = recovered.at(1).unwrap().as_list().unwrap();
        assert_eq!(recovered_id, id.unwrap());
        assert_eq!(recovered_request, request);
    }
}