Skip to content

Commit 1f73393

Browse files
palavivpalaviv
committed
Move more socket methods to new arg style
1 parent 9954361 commit 1f73393

File tree

1 file changed

+85
-116
lines changed

1 file changed

+85
-116
lines changed

‎vm/src/stdlib/socket.rs‎

Lines changed: 85 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,8 @@ use std::io;
33
use std::io::Read;
44
use std::io::Write;
55
use std::net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs, UdpSocket};
6-
use std::ops::Deref;
76

87
use crate::function::PyFuncArgs;
9-
use crate::obj::objbytes;
108
use crate::obj::objbytes::PyBytesRef;
119
use crate::obj::objint;
1210
use crate::obj::objint::PyIntRef;
@@ -163,10 +161,6 @@ impl Socket {
163161
}
164162
}
165163

166-
fn get_socket<'a>(obj: &'a PyObjectRef) -> impl Deref<Target = Socket> + 'a {
167-
obj.payload::<Socket>().unwrap()
168-
}
169-
170164
type SocketRef = PyRef<Socket>;
171165

172166
impl SocketRef {
@@ -226,33 +220,6 @@ impl SocketRef {
226220
}
227221
}
228222

229-
fn sendto(self, bytes: PyBytesRef, address: PyTupleRef, vm: &VirtualMachine) -> PyResult {
230-
let address_string = get_address_string(vm, address)?;
231-
232-
match self.socket_kind {
233-
SocketKind::Dgram => {
234-
if let Some(v) = self.con.borrow().as_ref() {
235-
return match v.send_to(&bytes, address_string) {
236-
Ok(_) => Ok(vm.get_none()),
237-
Err(s) => Err(vm.new_os_error(s.to_string())),
238-
};
239-
}
240-
// Doing implicit bind
241-
match UdpSocket::bind("0.0.0.0:0") {
242-
Ok(dgram) => match dgram.send_to(&bytes, address_string) {
243-
Ok(_) => {
244-
self.con.borrow_mut().replace(Connection::UdpSocket(dgram));
245-
Ok(vm.get_none())
246-
}
247-
Err(s) => Err(vm.new_os_error(s.to_string())),
248-
},
249-
Err(s) => Err(vm.new_os_error(s.to_string())),
250-
}
251-
}
252-
_ => Err(vm.new_not_implemented_error("".to_string())),
253-
}
254-
}
255-
256223
fn listen(self, _num: PyIntRef, _vm: &VirtualMachine) -> () {}
257224

258225
fn accept(self, vm: &VirtualMachine) -> PyResult {
@@ -289,6 +256,86 @@ impl SocketRef {
289256
};
290257
Ok(vm.ctx.new_bytes(buffer))
291258
}
259+
260+
fn recvfrom(self, bufsize: PyIntRef, vm: &VirtualMachine) -> PyResult {
261+
let mut buffer = vec![0u8; bufsize.as_bigint().to_usize().unwrap()];
262+
let ret = match self.con.borrow().as_ref() {
263+
Some(v) => v.recv_from(&mut buffer),
264+
None => return Err(vm.new_type_error("".to_string())),
265+
};
266+
267+
let addr = match ret {
268+
Ok((_size, addr)) => addr,
269+
Err(s) => return Err(vm.new_os_error(s.to_string())),
270+
};
271+
272+
let addr_tuple = get_addr_tuple(vm, addr)?;
273+
274+
Ok(vm.ctx.new_tuple(vec![vm.ctx.new_bytes(buffer), addr_tuple]))
275+
}
276+
277+
fn send(self, bytes: PyBytesRef, vm: &VirtualMachine) -> PyResult {
278+
match self.con.borrow_mut().as_mut() {
279+
Some(v) => match v.write(&bytes) {
280+
Ok(_) => (),
281+
Err(s) => return Err(vm.new_os_error(s.to_string())),
282+
},
283+
None => return Err(vm.new_type_error("".to_string())),
284+
};
285+
Ok(vm.get_none())
286+
}
287+
288+
fn sendto(self, bytes: PyBytesRef, address: PyTupleRef, vm: &VirtualMachine) -> PyResult {
289+
let address_string = get_address_string(vm, address)?;
290+
291+
match self.socket_kind {
292+
SocketKind::Dgram => {
293+
if let Some(v) = self.con.borrow().as_ref() {
294+
return match v.send_to(&bytes, address_string) {
295+
Ok(_) => Ok(vm.get_none()),
296+
Err(s) => Err(vm.new_os_error(s.to_string())),
297+
};
298+
}
299+
// Doing implicit bind
300+
match UdpSocket::bind("0.0.0.0:0") {
301+
Ok(dgram) => match dgram.send_to(&bytes, address_string) {
302+
Ok(_) => {
303+
self.con.borrow_mut().replace(Connection::UdpSocket(dgram));
304+
Ok(vm.get_none())
305+
}
306+
Err(s) => Err(vm.new_os_error(s.to_string())),
307+
},
308+
Err(s) => Err(vm.new_os_error(s.to_string())),
309+
}
310+
}
311+
_ => Err(vm.new_not_implemented_error("".to_string())),
312+
}
313+
}
314+
315+
fn close(self, vm: &VirtualMachine) -> PyResult {
316+
self.con.borrow_mut().take();
317+
Ok(vm.get_none())
318+
}
319+
320+
fn fileno(self, vm: &VirtualMachine) -> PyResult {
321+
let fileno = match self.con.borrow_mut().as_mut() {
322+
Some(v) => v.fileno(),
323+
None => return Err(vm.new_type_error("".to_string())),
324+
};
325+
Ok(vm.ctx.new_int(fileno))
326+
}
327+
328+
fn getsockname(self, vm: &VirtualMachine) -> PyResult {
329+
let addr = match self.con.borrow().as_ref() {
330+
Some(v) => v.local_addr(),
331+
None => return Err(vm.new_type_error("".to_string())),
332+
};
333+
334+
match addr {
335+
Ok(addr) => get_addr_tuple(vm, addr),
336+
Err(s) => Err(vm.new_os_error(s.to_string())),
337+
}
338+
}
292339
}
293340

294341
fn get_address_string(vm: &VirtualMachine, address: PyTupleRef) -> Result<String, PyObjectRef> {
@@ -312,84 +359,6 @@ fn get_address_string(vm: &VirtualMachine, address: PyTupleRef) -> Result<String
312359
))
313360
}
314361

315-
fn socket_recvfrom(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
316-
arg_check!(
317-
vm,
318-
args,
319-
required = [(zelf, None), (bufsize, Some(vm.ctx.int_type()))]
320-
);
321-
322-
let socket = get_socket(zelf);
323-
324-
let mut buffer = vec![0u8; objint::get_value(bufsize).to_usize().unwrap()];
325-
let ret = match socket.con.borrow().as_ref() {
326-
Some(v) => v.recv_from(&mut buffer),
327-
None => return Err(vm.new_type_error("".to_string())),
328-
};
329-
330-
let addr = match ret {
331-
Ok((_size, addr)) => addr,
332-
Err(s) => return Err(vm.new_os_error(s.to_string())),
333-
};
334-
335-
let addr_tuple = get_addr_tuple(vm, addr)?;
336-
337-
Ok(vm.ctx.new_tuple(vec![vm.ctx.new_bytes(buffer), addr_tuple]))
338-
}
339-
340-
fn socket_send(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
341-
arg_check!(
342-
vm,
343-
args,
344-
required = [(zelf, None), (bytes, Some(vm.ctx.bytes_type()))]
345-
);
346-
let socket = get_socket(zelf);
347-
348-
match socket.con.borrow_mut().as_mut() {
349-
Some(v) => match v.write(&objbytes::get_value(&bytes)) {
350-
Ok(_) => (),
351-
Err(s) => return Err(vm.new_os_error(s.to_string())),
352-
},
353-
None => return Err(vm.new_type_error("".to_string())),
354-
};
355-
Ok(vm.get_none())
356-
}
357-
358-
fn socket_close(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
359-
arg_check!(vm, args, required = [(zelf, None)]);
360-
361-
let socket = get_socket(zelf);
362-
socket.con.borrow_mut().take();
363-
Ok(vm.get_none())
364-
}
365-
366-
fn socket_fileno(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
367-
arg_check!(vm, args, required = [(zelf, None)]);
368-
369-
let socket = get_socket(zelf);
370-
371-
let fileno = match socket.con.borrow_mut().as_mut() {
372-
Some(v) => v.fileno(),
373-
None => return Err(vm.new_type_error("".to_string())),
374-
};
375-
Ok(vm.ctx.new_int(fileno))
376-
}
377-
378-
fn socket_getsockname(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
379-
arg_check!(vm, args, required = [(zelf, None)]);
380-
let socket = get_socket(zelf);
381-
382-
let addr = match socket.con.borrow().as_ref() {
383-
Some(v) => v.local_addr(),
384-
None => return Err(vm.new_type_error("".to_string())),
385-
};
386-
387-
match addr {
388-
Ok(addr) => get_addr_tuple(vm, addr),
389-
Err(s) => Err(vm.new_os_error(s.to_string())),
390-
}
391-
}
392-
393362
fn get_addr_tuple(vm: &VirtualMachine, addr: SocketAddr) -> PyResult {
394363
let port = vm.ctx.new_int(addr.port());
395364
let ip = vm.ctx.new_str(addr.ip().to_string());
@@ -404,15 +373,15 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
404373
"__new__" => ctx.new_rustfunc(SocketRef::new),
405374
"connect" => ctx.new_rustfunc(SocketRef::connect),
406375
"recv" => ctx.new_rustfunc(SocketRef::recv),
407-
"send" => ctx.new_rustfunc(socket_send),
376+
"send" => ctx.new_rustfunc(SocketRef::send),
408377
"bind" => ctx.new_rustfunc(SocketRef::bind),
409378
"accept" => ctx.new_rustfunc(SocketRef::accept),
410379
"listen" => ctx.new_rustfunc(SocketRef::listen),
411-
"close" => ctx.new_rustfunc(socket_close),
412-
"getsockname" => ctx.new_rustfunc(socket_getsockname),
380+
"close" => ctx.new_rustfunc(SocketRef::close),
381+
"getsockname" => ctx.new_rustfunc(SocketRef::getsockname),
413382
"sendto" => ctx.new_rustfunc(SocketRef::sendto),
414-
"recvfrom" => ctx.new_rustfunc(socket_recvfrom),
415-
"fileno" => ctx.new_rustfunc(socket_fileno),
383+
"recvfrom" => ctx.new_rustfunc(SocketRef::recvfrom),
384+
"fileno" => ctx.new_rustfunc(SocketRef::fileno),
416385
});
417386

418387
py_module!(vm, "socket", {

0 commit comments

Comments
 (0)