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
149
150
151
152
153

use super::{
    error::Error,
    protocol::{Link, Protocol, ProtocolV1},
    AvailableDeviceTransport, ProtoMessage, Transport,
};
use crate::{AvailableDevice, Model};
use std::{fmt, net::UdpSocket, result::Result, time::Duration};

// A collection of constants related to the Emulator Ports.
mod constants {
    pub(crate) const DEFAULT_HOST: &str = "127.0.0.1";
    pub(crate) const DEFAULT_PORT: &str = "21324";
    pub(crate) const DEFAULT_DEBUG_PORT: &str = "21325";
    pub(crate) const LOCAL_LISTENER: &str = "127.0.0.1:0";
}

use constants::{DEFAULT_DEBUG_PORT, DEFAULT_HOST, DEFAULT_PORT, LOCAL_LISTENER};

/// The chunk size for the serial protocol.
const CHUNK_SIZE: usize = 64;

const READ_TIMEOUT_MS: u64 = 100000;
const WRITE_TIMEOUT_MS: u64 = 100000;

/// An available transport for connecting with a device.
#[derive(Debug)]
pub struct AvailableUdpTransport {
    pub host: String,
    pub port: String,
}

impl fmt::Display for AvailableUdpTransport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "udp:{}:{}", self.host, self.port)
    }
}

/// An actual serial HID USB link to a device over which bytes can be sent.
struct UdpLink {
    socket: UdpSocket,
    device: (String, String),
}
// No need to implement drop as every member is owned

impl Link for UdpLink {
    fn write_chunk(&mut self, chunk: Vec<u8>) -> Result<(), Error> {
        debug_assert_eq!(CHUNK_SIZE, chunk.len());
        let timeout = Duration::from_millis(WRITE_TIMEOUT_MS);
        self.socket.set_write_timeout(Some(timeout))?;
        if let Err(e) = self.socket.send(&chunk) {
            return Err(e.into())
        }
        Ok(())
    }

    fn read_chunk(&mut self) -> Result<Vec<u8>, Error> {
        let mut chunk = vec![0; CHUNK_SIZE];
        let timeout = Duration::from_millis(READ_TIMEOUT_MS);
        self.socket.set_read_timeout(Some(timeout))?;

        let n = self.socket.recv(&mut chunk)?;
        if n == CHUNK_SIZE {
            Ok(chunk)
        } else {
            Err(Error::DeviceReadTimeout)
        }
    }
}

impl UdpLink {
    fn open(path: &str) -> Result<UdpLink, Error> {
        let mut parts = path.split(':');
        let link = Self {
            socket: UdpSocket::bind(LOCAL_LISTENER)?,
            device: (
                parts.next().expect("Incorrect Path").to_owned(),
                parts.next().expect("Incorrect Path").to_owned(),
            ),
        };
        link.socket.connect(path)?;
        Ok(link)
    }

    // Ping the port and compare against expected response
    fn ping(&self) -> Result<bool, Error> {
        let mut resp = [0; CHUNK_SIZE];
        self.socket.send("PINGPING".as_bytes())?;
        let size = self.socket.recv(&mut resp)?;
        Ok(&resp[..size] == "PONGPONG".as_bytes())
    }
}

/// An implementation of the Transport interface for UDP devices.
pub struct UdpTransport {
    protocol: ProtocolV1<UdpLink>,
}

impl UdpTransport {
    pub fn find_devices(debug: bool, path: Option<&str>) -> Result<Vec<AvailableDevice>, Error> {
        let mut devices = Vec::new();
        let mut dest = String::new();
        match path {
            Some(p) => dest = p.to_owned(),
            None => {
                dest.push_str(DEFAULT_HOST);
                dest.push(':');
                dest.push_str(if debug { DEFAULT_DEBUG_PORT } else { DEFAULT_PORT });
            }
        };
        let link = UdpLink::open(&dest)?;
        if link.ping()? {
            devices.push(AvailableDevice {
                model: Model::TrezorEmulator,
                debug,
                transport: AvailableDeviceTransport::Udp(AvailableUdpTransport {
                    host: link.device.0,
                    port: link.device.1,
                }),
            });
        }
        Ok(devices)
    }

    /// Connect to a device over the UDP transport.
    pub fn connect(device: &AvailableDevice) -> Result<Box<dyn Transport>, Error> {
        let transport = match device.transport {
            AvailableDeviceTransport::Udp(ref t) => t,
            _ => panic!("passed wrong AvailableDevice in UdpTransport::connect"),
        };
        let mut path = String::new();
        path.push_str(&transport.host);
        path.push(':');
        path.push_str(&transport.port);
        let link = UdpLink::open(&path)?;
        Ok(Box::new(UdpTransport { protocol: ProtocolV1 { link } }))
    }
}

impl super::Transport for UdpTransport {
    fn session_begin(&mut self) -> Result<(), Error> {
        self.protocol.session_begin()
    }
    fn session_end(&mut self) -> Result<(), Error> {
        self.protocol.session_end()
    }

    fn write_message(&mut self, message: ProtoMessage) -> Result<(), Error> {
        self.protocol.write(message)
    }
    fn read_message(&mut self) -> Result<ProtoMessage, Error> {
        self.protocol.read()
    }
}