Skip to content

Commit 740f00d

Browse files
jkcoxsonjkcoxson
committed
Refactor out dangerous unwraps
1 parent 067f688 commit 740f00d

File tree

4 files changed

+155
-60
lines changed

4 files changed

+155
-60
lines changed

‎src/devices.rs

Lines changed: 76 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use std::{collections::HashMap, io::Read, net::IpAddr, path::PathBuf, sync::Arc}
55
use log::{debug, info, trace, warn};
66
use tokio::{
77
io::AsyncReadExt,
8-
sync::{mpsc::UnboundedSender, Mutex},
8+
sync::{oneshot::Sender, Mutex},
99
};
1010

1111
use crate::heartbeat;
@@ -29,7 +29,7 @@ pub struct MuxerDevice {
2929

3030
// Network types
3131
pub network_address: Option<IpAddr>,
32-
pub heartbeat_handle: Option<UnboundedSender<()>>,
32+
pub heartbeat_handle: Option<Sender<()>>,
3333
pub service_name: Option<String>,
3434

3535
// USB types
@@ -125,16 +125,21 @@ impl SharedDevices {
125125
return;
126126
}
127127
info!("Removing device: {:?}", udid);
128-
let _ = &self
129-
.devices
130-
.get(udid)
131-
.unwrap()
132-
.heartbeat_handle
133-
.as_ref()
134-
.unwrap()
135-
.send(())
136-
.unwrap();
137-
self.devices.remove(udid);
128+
match self.devices.remove(udid) {
129+
Some(d) => match d.heartbeat_handle {
130+
Some(h) => {
131+
if let Err(e) = h.send(()) {
132+
warn!("Couldn't send kill signal to heartbeat thread: {e:?}");
133+
}
134+
}
135+
None => {
136+
warn!("Heartbeat handle option has none, can't kill");
137+
}
138+
},
139+
None => {
140+
warn!("No heartbeat handle found for device");
141+
}
142+
};
138143
}
139144
pub async fn get_pairing_record(&self, udid: String) -> Result<Vec<u8>, std::io::Error> {
140145
let path = PathBuf::from(self.plist_storage.clone()).join(format!("{}.plist", udid));
@@ -166,13 +171,22 @@ impl SharedDevices {
166171
}
167172
// Read the file to a string
168173
debug!("Reading SystemConfiguration.plist");
169-
let mut file = std::fs::File::open(path).unwrap();
174+
let mut file = std::fs::File::open(path)?;
170175
let mut contents = Vec::new();
171-
file.read_to_end(&mut contents).unwrap();
176+
file.read_to_end(&mut contents)?;
172177

173178
// Parse the string into a plist
174179
debug!("Parsing SystemConfiguration.plist");
175-
let plist = plist::from_bytes::<plist::Dictionary>(&contents).unwrap();
180+
let plist = match plist::from_bytes::<plist::Dictionary>(&contents) {
181+
Ok(p) => p,
182+
Err(e) => {
183+
log::error!("Failed to parse plist: {e:?}");
184+
return Err(std::io::Error::new(
185+
std::io::ErrorKind::InvalidInput,
186+
"unable to parse plist",
187+
));
188+
}
189+
};
176190
match plist.get("SystemBUID") {
177191
Some(plist::Value::String(b)) => Ok(b.to_owned()),
178192
_ => Err(std::io::Error::new(
@@ -187,14 +201,32 @@ impl SharedDevices {
187201
trace!("Updating plist cache");
188202
let path = PathBuf::from(self.plist_storage.clone());
189203
for entry in std::fs::read_dir(path).expect("Plist storage is unreadable!!") {
190-
let entry = entry.unwrap();
204+
let entry = match entry {
205+
Ok(e) => e,
206+
Err(e) => {
207+
warn!("Unable to read entry in plist storage: {e:?}");
208+
continue;
209+
}
210+
};
191211
let path = entry.path();
192212
trace!("Attempting to read {:?}", path);
193213
if path.is_file() {
194-
let mut file = tokio::fs::File::open(&path).await.unwrap();
214+
let mut file = match tokio::fs::File::open(&path).await {
215+
Ok(f) => f,
216+
Err(e) => {
217+
warn!("Unable to read plist storage entry to memory: {e:?}");
218+
continue;
219+
}
220+
};
195221
let mut contents = Vec::new();
196222
let plist: plist::Dictionary = match file.read_to_end(&mut contents).await {
197-
Ok(_) => plist::from_bytes(&contents).unwrap(),
223+
Ok(_) => match plist::from_bytes(&contents) {
224+
Ok(p) => p,
225+
Err(e) => {
226+
warn!("Unable to parse entry file to plist: {e:?}");
227+
continue;
228+
}
229+
},
198230
Err(e) => {
199231
trace!("Could not read plist to memory: {e:?}");
200232
continue;
@@ -223,20 +255,30 @@ impl SharedDevices {
223255
// This is just used as a last resort, but might not be correct so we'll pass a warning
224256
warn!("Using the file name as the UDID");
225257
match path.file_name() {
226-
Some(f) => {
227-
f.to_str().unwrap().split('.').collect::<Vec<&str>>()[0].to_string()
228-
}
258+
Some(f) => match f.to_str() {
259+
Some(f) => f.split('.').collect::<Vec<&str>>()[0].to_string(),
260+
None => {
261+
warn!("Failed to get entry file name string");
262+
continue;
263+
}
264+
},
229265
None => {
230266
trace!("File had no name");
231267
continue;
232268
}
233269
}
234270
};
235271

236-
self.known_mac_addresses.insert(
237-
mac_addr.to_owned(),
238-
path.file_stem().unwrap().to_string_lossy().to_string(),
239-
);
272+
let stem = match path.file_stem() {
273+
Some(s) => s,
274+
None => {
275+
warn!("Failed to get file stem for entry");
276+
continue;
277+
}
278+
};
279+
280+
self.known_mac_addresses
281+
.insert(mac_addr.to_owned(), stem.to_string_lossy().to_string());
240282
if self.paired_udids.contains(&udid) {
241283
trace!("Cache already contained this UDID");
242284
continue;
@@ -277,15 +319,22 @@ impl From<&MuxerDevice> for plist::Dictionary {
277319
if device.connection_type == "Network" {
278320
p.insert(
279321
"EscapedFullServiceName".into(),
280-
device.service_name.clone().unwrap().into(),
322+
device
323+
.service_name
324+
.clone()
325+
.expect("Network device, but no service name")
326+
.into(),
281327
);
282328
}
283329
p.insert("InterfaceIndex".into(), device.interface_index.into());
284330

285331
// Reassemble the network address back into bytes
286332
if device.connection_type == "Network" {
287333
let mut data = [0u8; 152];
288-
match device.network_address.unwrap() {
334+
match device
335+
.network_address
336+
.expect("Network device, but no address")
337+
{
289338
IpAddr::V4(ip_addr) => {
290339
data[0] = 0x02;
291340
data[1] = 0x00;

‎src/heartbeat.rs

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
// jkcoxson
22

33
use idevice::{heartbeat::HeartbeatClient, lockdownd::LockdowndClient, Idevice};
4-
use log::info;
4+
use log::{info, warn};
55
use std::{
66
net::{IpAddr, SocketAddr},
7-
sync::{Arc, Mutex},
7+
sync::Arc,
88
};
9-
use tokio::sync::mpsc::UnboundedSender;
9+
use tokio::sync::oneshot::{error::TryRecvError, Sender};
1010

1111
use crate::devices::SharedDevices;
1212

@@ -15,10 +15,8 @@ pub async fn heartbeat(
1515
udid: String,
1616
pairing_file: idevice::pairing_file::PairingFile,
1717
data: Arc<tokio::sync::Mutex<SharedDevices>>,
18-
) -> Result<UnboundedSender<()>, Box<dyn std::error::Error>> {
19-
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
20-
let pls_stop = Arc::new(Mutex::new(false));
21-
let pls_stop_clone = pls_stop.clone();
18+
) -> Result<Sender<()>, Box<dyn std::error::Error>> {
19+
let (tx, mut rx) = tokio::sync::oneshot::channel();
2220

2321
let socket = SocketAddr::new(ip_addr, LockdowndClient::LOCKDOWND_PORT);
2422

@@ -31,8 +29,7 @@ pub async fn heartbeat(
3129

3230
let (port, _) = lockdown_client
3331
.start_service("com.apple.mobile.heartbeat")
34-
.await
35-
.unwrap();
32+
.await?;
3633

3734
let socket = SocketAddr::new(ip_addr, port);
3835
let socket = tokio::net::TcpStream::connect(socket).await?;
@@ -41,10 +38,9 @@ pub async fn heartbeat(
4138

4239
idevice.start_session(&pairing_file).await?;
4340

44-
tokio::spawn(async {
41+
tokio::spawn(async move {
4542
let mut interval = 10;
4643
let mut heartbeat_client = HeartbeatClient { idevice };
47-
let pls_stop = pls_stop;
4844
loop {
4945
match heartbeat_client.get_marco(interval + 5).await {
5046
Ok(i) => {
@@ -58,8 +54,15 @@ pub async fn heartbeat(
5854
break;
5955
}
6056
}
61-
if *pls_stop.lock().unwrap() {
62-
break;
57+
match rx.try_recv() {
58+
Ok(_) => {
59+
info!("Heartbeat instructed to die")
60+
}
61+
Err(TryRecvError::Closed) => {
62+
warn!("Heartbeat killer closed");
63+
break;
64+
}
65+
_ => {}
6366
}
6467
if let Err(e) = heartbeat_client.send_polo().await {
6568
info!("Heartbeat send failed: {:?}", e);
@@ -70,10 +73,6 @@ pub async fn heartbeat(
7073
}
7174
}
7275
});
73-
tokio::spawn(async move {
74-
rx.recv().await;
75-
*pls_stop_clone.lock().unwrap() = true;
76-
});
7776
Ok(tx)
7877
}
7978

‎src/main.rs

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -42,15 +42,28 @@ async fn main() {
4242
while i < std::env::args().len() {
4343
match std::env::args().nth(i).unwrap().as_str() {
4444
"-p" | "--port" => {
45-
port = std::env::args().nth(i + 1).unwrap().parse::<i32>().unwrap();
45+
port = std::env::args()
46+
.nth(i + 1)
47+
.expect("port flag passed without number")
48+
.parse::<i32>()
49+
.expect("port isn't a number");
4650
i += 2;
4751
}
4852
"--host" => {
49-
host = Some(std::env::args().nth(i + 1).unwrap().to_string());
53+
host = Some(
54+
std::env::args()
55+
.nth(i + 1)
56+
.expect("host flag passed without host")
57+
.to_string(),
58+
);
5059
i += 2;
5160
}
5261
"--plist-storage" => {
53-
plist_storage = Some(std::env::args().nth(i + 1).unwrap());
62+
plist_storage = Some(
63+
std::env::args()
64+
.nth(i + 1)
65+
.expect("flag passed without value"),
66+
);
5467
i += 1;
5568
}
5669
#[cfg(unix)]
@@ -110,7 +123,7 @@ async fn main() {
110123
// Create TcpListener
111124
let listener = tokio::net::TcpListener::bind(format!("{}:{}", host, port))
112125
.await
113-
.unwrap();
126+
.expect("Unable to bind to TCP listener");
114127

115128
println!("Listening on {}:{}", host, port);
116129
#[cfg(unix)]
@@ -137,10 +150,12 @@ async fn main() {
137150
std::fs::remove_file("/var/run/usbmuxd").unwrap_or_default();
138151
// Create UnixListener
139152
info!("Binding to new Unix socket");
140-
let listener = tokio::net::UnixListener::bind("/var/run/usbmuxd").unwrap();
153+
let listener = tokio::net::UnixListener::bind("/var/run/usbmuxd")
154+
.expect("Unable to bind to unix socket");
141155
// Change the permission of the socket
142156
info!("Changing permissions of socket");
143-
fs::set_permissions("/var/run/usbmuxd", fs::Permissions::from_mode(0o666)).unwrap();
157+
fs::set_permissions("/var/run/usbmuxd", fs::Permissions::from_mode(0o666))
158+
.expect("Unable to set socket file permissions");
144159

145160
println!("Listening on /var/run/usbmuxd");
146161

@@ -207,7 +222,14 @@ async fn handle_stream(
207222
info!("Only read the header, pulling more bytes");
208223
// Get the number of bytes to pull
209224
let packet_size = &buffer[0..4];
210-
let packet_size = u32::from_le_bytes(packet_size.try_into().unwrap());
225+
let packet_size = match packet_size.try_into() {
226+
Ok(p) => p,
227+
Err(e) => {
228+
warn!("Failed to read packet size: {e:?}");
229+
return;
230+
}
231+
};
232+
let packet_size = u32::from_le_bytes(packet_size);
211233
info!("Packet size: {}", packet_size);
212234
// Pull the rest of the packet
213235
let mut packet = vec![0; packet_size as usize];
@@ -232,7 +254,7 @@ async fn handle_stream(
232254
return;
233255
}
234256
};
235-
trace!("Recv'd plist: {parsed:?}");
257+
trace!("Recv'd plist: {parsed:#?}");
236258

237259
match current_directions {
238260
Directions::None => {
@@ -397,7 +419,13 @@ async fn handle_stream(
397419
}
398420
"ReadBUID" => {
399421
let lock = data.lock().await;
400-
let buid = lock.get_buid().await.unwrap();
422+
let buid = match lock.get_buid().await {
423+
Ok(b) => b,
424+
Err(e) => {
425+
log::error!("Failed to get buid: {e:?}");
426+
return;
427+
}
428+
};
401429

402430
let mut p = plist::Dictionary::new();
403431
p.insert("BUID".into(), buid.into());

0 commit comments

Comments
 (0)