1 files changed, 26 insertions, 2 deletions
diff --git a/test/test_socks.py b/test/test_socks.py
index 964678de..740fdb9c 100644
--- a/test/test_socks.py
+++ b/test/test_socks.py
@@ -1,7 +1,7 @@
 from cStringIO import StringIO
 import socket
 from nose.plugins.skip import SkipTest
-from netlib import socks, utils
+from netlib import socks, tcp
 import tutils
 
 
@@ -42,10 +42,16 @@ def test_message():
     assert msg.atyp == 0x03
     assert msg.addr == ("example.com", 0xDEAD)
 
+
+def test_message_ipv4():
     # Test ATYP=0x01 (IPV4)
     raw = StringIO("\x05\x01\x00\x01\x7f\x00\x00\x01\xDE\xAD\xBE\xEF")
+    out = StringIO()
     msg = socks.Message.from_file(raw)
     assert raw.read(2) == "\xBE\xEF"
+    msg.to_file(out)
+
+    assert out.getvalue() == raw.getvalue()[:-2]
     assert msg.addr == ("127.0.0.1", 0xDEAD)
 
 
@@ -54,7 +60,25 @@ def test_message_ipv6():
         raise SkipTest("Skipped because inet_ntop is not available")
     # Test ATYP=0x04 (IPV6)
     ipv6_addr = "2001:db8:85a3:8d3:1319:8a2e:370:7344"
+
     raw = StringIO("\x05\x01\x00\x04" + socket.inet_pton(socket.AF_INET6, ipv6_addr) + "\xDE\xAD\xBE\xEF")
+    out = StringIO()
     msg = socks.Message.from_file(raw)
     assert raw.read(2) == "\xBE\xEF"
-    assert msg.addr.host == ipv6_addr
-\ No newline at end of file
+    msg.to_file(out)
+
+    assert out.getvalue() == raw.getvalue()[:-2]
+    assert msg.addr.host == ipv6_addr
+
+
+def test_message_invalid_rsv():
+    raw = StringIO("\x05\x01\xFF\x01\x7f\x00\x00\x01\xDE\xAD\xBE\xEF")
+    tutils.raises(socks.SocksError, socks.Message.from_file, raw)
+
+
+def test_message_unknown_atyp():
+    raw = StringIO("\x05\x02\x00\x02\x7f\x00\x00\x01\xDE\xAD\xBE\xEF")
+    tutils.raises(socks.SocksError, socks.Message.from_file, raw)
+
+    m = socks.Message(5, 1, 0x02, tcp.Address(("example.com", 5050)))
+    tutils.raises(socks.SocksError, m.to_file, StringIO())
+\ No newline at end of file