Add validation to connect and bind address

palaviv · palaviv · commit 716eb4afca12 · 2019-02-27T19:54:39.000+02:00
diff --git a/tests/snippets/stdlib_socket.py b/tests/snippets/stdlib_socket.py
@@ -1,4 +1,5 @@
 import socket
+from testutils import assertRaises
 
 listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 listener.bind(("127.0.0.1", 8080))
@@ -22,3 +23,14 @@
 connector.close()
 listener.close()
 
+s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+with assertRaises(TypeError):
+	s.connect(("127.0.0.1", 8888, 8888))
+
+with assertRaises(TypeError):
+	s.bind(("127.0.0.1", 8888, 8888))
+
+with assertRaises(TypeError):
+	s.bind((888, 8888))
+
+s.close()
diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs
@@ -130,11 +130,7 @@ fn socket_connect(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(zelf, None), (address, Some(vm.ctx.tuple_type()))]
     );
 
-    let elements = get_elements(address);
-    let host = objstr::get_value(&elements[0]);
-    let port = objint::get_value(&elements[1]);
-
-    let address_string = format!("{}:{}", host, port.to_string());
+    let address_string = get_address_string(vm, address)?;
 
     match zelf.payload {
         PyObjectPayload::Socket { ref socket } => {
@@ -157,11 +153,7 @@ fn socket_bind(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(zelf, None), (address, Some(vm.ctx.tuple_type()))]
     );
 
-    let elements = get_elements(address);
-    let host = objstr::get_value(&elements[0]);
-    let port = objint::get_value(&elements[1]);
-
-    let address_string = format!("{}:{}", host, port.to_string());
+    let address_string = get_address_string(vm, address)?;
 
     match zelf.payload {
         PyObjectPayload::Socket { ref socket } => {
@@ -177,6 +169,30 @@ fn socket_bind(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
+fn get_address_string(
+    vm: &mut VirtualMachine,
+    address: &PyObjectRef,
+) -> Result<String, PyObjectRef> {
+    let args = PyFuncArgs {
+        args: get_elements(address).to_vec(),
+        kwargs: vec![],
+    };
+    arg_check!(
+        vm,
+        args,
+        required = [
+            (host, Some(vm.ctx.str_type())),
+            (port, Some(vm.ctx.int_type()))
+        ]
+    );
+
+    Ok(format!(
+        "{}:{}",
+        objstr::get_value(host),
+        objint::get_value(port).to_string()
+    ))
+}
+
 fn socket_listen(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.get_none())
 }
