1
use std::io::IoSlice;
2

            
3
const RETRY_CNT: u32 = 5;
4

            
5
use crate::{ Error, Result };
6
use crate::util::{ SocketAddr, UdpSocket };
7

            
8
use super::{ Request, RequestError, Datagram, Oack, Xfer, SequenceId };
9

            
10
#[derive(Default, Debug)]
11
pub struct Stats {
12
    pub filesize:	u64,
13
    pub xmitsz:		u64,
14
    pub retries:	u32,
15
    pub wastedsz:	u64,
16
    pub num_timeouts:	u32,
17
    pub window_size:	u16,
18
    pub block_size:	u16,
19
}
20

            
21
impl Stats {
22
    pub fn has_errors(&self) -> bool {
23
	self.filesize != self.xmitsz ||
24
	    self.retries != 0 ||
25
	    self.wastedsz != 0 ||
26
	    self.num_timeouts != 0
27
    }
28

            
29
    pub fn speed_bit_per_s(&self, duration: std::time::Duration) -> Option<(f32, f32)> {
30
	if duration.is_zero() {
31
	    return None;
32
	}
33

            
34
	Some(((self.filesize as f64 / duration.as_secs_f64()) as f32,
35
	      (self.xmitsz as f64 / duration.as_secs_f64()) as f32))
36
    }
37
}
38

            
39
impl std::fmt::Display for Stats {
40
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41
	use num_format::{ ToFormattedString, SystemLocale };
42
	let locale = SystemLocale::default().unwrap();
43

            
44
	if ! self.has_errors() {
45
            write!(f, "filesize={}, windowsize={}, blocksize={}",
46
		   self.filesize.to_formatted_string(&locale), self.window_size, self.block_size)
47
	} else {
48
            write!(f, "filesize={}, sent={} ({} retries, {} wasted, {} timeouts), windowsize={}, blocksize={}",
49
		   self.filesize.to_formatted_string(&locale),
50
		   self.xmitsz.to_formatted_string(&locale),
51
		   self.retries, self.wastedsz.to_formatted_string(&locale),
52
		   self.num_timeouts,
53
		   self.window_size, self.block_size)
54
	}
55
    }
56
}
57

            
58
pub struct Session<'a> {
59
    remote:	SocketAddr,
60
    sock:	UdpSocket,
61
    env:	&'a crate::Environment,
62

            
63
    window_size:	u16,
64
    block_size:		u16,
65
    timeout:		std::time::Duration,
66
}
67

            
68
impl <'a> Session<'a> {
69
    pub async fn new(env: &'a crate::Environment,
70
		     remote: SocketAddr,
71
		     local: std::net::IpAddr) -> Result<Session<'a>> {
72
	let local_addr = SocketAddr::new(local, 0);
73

            
74
	let sock = UdpSocket::bind(local_addr)?;
75

            
76
	tracing::Span::current().record("remote", &remote.to_string());
77
	tracing::Span::current().record("local",  &sock.local_addr().unwrap().to_string());
78

            
79
	Ok(Self {
80
	    remote:		remote,
81
	    sock:		sock,
82
	    env:		env,
83

            
84
	    window_size:	1,
85
	    block_size:		512,
86
	    timeout:		env.timeout,
87
	})
88
    }
89

            
90
    async fn send(&self, msg: &[u8]) -> Result<()>
91
    {
92
	self.sock.sendto(msg, self.remote).await?;
93
	Ok(())
94
    }
95

            
96
    async fn send_slice(&self, data: &[IoSlice<'_>]) -> Result<()>
97
    {
98
	self.sock.sendmsg(data, self.remote).await
99
    }
100

            
101
    async fn send_datagram(&self, msg: Datagram<'_>) -> Result<()>
102
    {
103
	match msg {
104
	    Datagram::Data(seq, data)	=> {
105
		let seq = seq.as_slice();
106
		let data = &[
107
		    IoSlice::new(&[0, 3]),
108
		    IoSlice::new(&seq),
109
		    IoSlice::new(data),
110
		];
111

            
112
		self.send_slice(data).await
113
	    },
114

            
115
	    _	=> Err(Error::Internal("send_datagram not implemented for this message"))
116
	}
117
    }
118

            
119
    async fn send_err(self, e: Error) -> Result<()>
120
    {
121
	let mut msg = Vec::<u8>::with_capacity(1500);
122

            
123
	warn!("error: {}", e);
124

            
125
	msg.extend([0, 5]);
126

            
127
	match e {
128
	    Error::RequestError(d)	=> {
129
		msg.extend([0, 4]);
130
		msg.extend(d.to_string().as_bytes());
131
		msg.push(0);
132
	    },
133

            
134
	    Error::FileMissing		=> {
135
		msg.extend([0, 1]);
136
	    },
137

            
138
	    Error::TooMuchClients	=> {
139
		msg.extend([0, 4]);
140
		msg.extend(b"too much clients");
141
		msg.push(0);
142
	    },
143
	    _				=> {
144
		msg.extend([0, 0]);
145
	    },
146
	};
147

            
148
	msg.push(0);
149

            
150
	self.send(&msg).await
151
    }
152

            
153
    async fn send_oack(&self, oack: Oack) -> Result<()>
154
    {
155
	let mut msg = Vec::<u8>::with_capacity(1500);
156

            
157
	oack.fill_buf(&mut msg);
158

            
159
	self.send(&msg).await
160
    }
161

            
162
    async fn run_wrq(self, _req: Request<'_>) -> Result<Stats>
163
    {
164
	self.send_err(RequestError::WriteUnsupported.into()).await?;
165

            
166
	Err(Error::NotImplemented)
167
    }
168

            
169
    async fn run_oack(&mut self, mut oack: Oack, file_size: Option<u64>) -> Result<()>
170
    {
171
	oack.update_tsize(file_size);
172
	oack.update_block_size(self.env.max_block_size,   |v| self.block_size = v);
173
	oack.update_window_size(self.env.max_window_size, |v| self.window_size = v);
174
	oack.update_timeout(|v| self.timeout = v);
175

            
176
	let mut buf = vec![0u8; 1500];
177

            
178
	self.send_oack(oack).await?;
179

            
180
	let resp = Datagram::recv(&self.sock, &mut buf, &self.remote, self.timeout).await?;
181

            
182
	match resp {
183
	    Datagram::Ack(id) if id.as_u16() == 0	=> {},
184
	    Datagram::Ack(id)	=> {
185
		warn!("ACK of OACK with invalid id {}", id.as_u16());
186
		return Err(Error::BadAck);
187
	    }
188
	    r			=> {
189
		warn!("bad response to OACK: {:?}", r);
190
		return Err(Error::Protocol("bad response to OACK"));
191
	    },
192
	};
193

            
194
	Ok(())
195
    }
196

            
197
    fn log_request(&self, req: &Request<'_>, op: &'static str)
198
    {
199
	tracing::Span::current().record("op", &op.to_string());
200
	tracing::Span::current().record("filename", &req.get_filename().to_string_lossy().into_owned());
201

            
202
	debug!("request={:?}", req);
203
    }
204

            
205
    async fn run_rrq(mut self, req: Request<'_>) -> Result<Stats>
206
    {
207
	use crate::fetcher::Builder;
208

            
209
	let mut stats = Stats::default();
210

            
211
	self.log_request(&req, "read");
212

            
213
	let mut fetcher = Builder::new(self.env).instanciate(&req.get_filename())?;
214

            
215
	if let Err(e) = fetcher.open() {
216
	    self.send_err(e.clone()).await?;
217
	    return Err(e);
218
	}
219

            
220
	let fsize = fetcher.get_size();
221

            
222
	if let Some(sz) = fsize {
223
	    stats.filesize = sz;
224
	    tracing::Span::current().record("filesize", &sz);
225
	}
226

            
227
	let mut seq = match req.has_options() {
228
	    false	=> SequenceId::new(0),
229
	    true	=> {
230
		self.run_oack(Oack::from_request(&req), fsize).await?;
231
		SequenceId::new(1)
232
	    }
233
	};
234

            
235
	stats.window_size = self.window_size;
236
	stats.block_size  = self.block_size;
237

            
238
	let mut xfer = Xfer::new(&fetcher, self.block_size, self.window_size);
239
	let mut retry = RETRY_CNT;
240
	let mut is_startup = true;
241

            
242
	loop {
243
	    match xfer.fill_window(seq, &mut fetcher).await? {
244
		0	=> {},
245
		v	=> {
246
		    debug!("retransmitting {:?}+", seq);
247

            
248
		    stats.retries += 1;
249
		    stats.wastedsz += v as u64;
250
		}
251
	    }
252

            
253

            
254
	    if xfer.is_eof() {
255
		break;
256
	    }
257

            
258
	    for d in xfer.iter() {
259
		stats.xmitsz += d.get_data_len() as u64;
260
		self.send_datagram(d).await?;
261
	    }
262

            
263

            
264
	    let mut buf = vec![0u8; 1500];
265
	    let resp = Datagram::recv(&self.sock, &mut buf, &self.remote, self.timeout).await;
266

            
267
	    match resp {
268
		Err(Error::Timeout) if retry > 0    => {
269
		    retry -= 1;
270
		    stats.num_timeouts += 1;
271
		    debug!("timeout; resending seq {}", seq.as_u16());
272
		},
273
		Ok(Datagram::Ack(id))	=> {
274
		    debug!("got ACK #{}", id.as_u16());
275
		    is_startup = false;
276
		    retry = RETRY_CNT;
277
		    seq = id + 1
278
		},
279

            
280
		Ok(Datagram::Error(code, info))	if is_startup => {
281
		    debug!("remote site sent error #{} ({}) on startup; probably just testing for existence",
282
			   code, String::from_utf8_lossy(info));
283
		    break;
284
		}
285

            
286
		Ok(Datagram::Error(code, info)) => {
287
		    info!("remote site sent error #{} ({})", code, String::from_utf8_lossy(info));
288
		    break;
289
		}
290

            
291
		Err(Error::Timeout)	=> {
292
		    warn!("timeout while waiting for ACK");
293
		    return Err(Error::Timeout);
294
		},
295
		r			=> {
296
		    warn!("bad response to DATA: {:?}", r);
297
		    return Err(Error::Protocol("bad response to DATA"));
298
		},
299
	    }
300
	}
301

            
302
	debug!("stats: {:?}", stats);
303

            
304
	Ok(stats)
305
    }
306

            
307
    pub async fn run(self, req: Vec<u8>) -> Result<Stats>
308
    {
309
	let op = Datagram::try_from(req.as_slice());
310

            
311
	match op {
312
	    Ok(Datagram::Write(r))	=> self.run_wrq(r).await,
313
	    Ok(Datagram::Read(r))	=> self.run_rrq(r).await,
314
	    Ok(_)	=> {
315
		self.send_err(RequestError::OperationUnsupported.into()).await?;
316
		Err(RequestError::OperationUnsupported.into())
317
	    },
318
	    Err(e)	=> {
319
		self.send_err(e.clone()).await?;
320
		Err(e)
321
	    }
322
	}
323
    }
324

            
325
    pub async fn do_reject(self) -> Result<Stats>
326
    {
327
	self.send_err(Error::TooMuchClients).await?;
328
	Err(Error::TooMuchClients)
329
    }
330
}