From 18489de6d46d23332dea5613bc2a44a8f09d30c2 Mon Sep 17 00:00:00 2001
From: Pedro Ferreira <pedro@dete.st>
Date: Mon, 5 Jun 2017 22:49:27 +0200
Subject: [PATCH] Start using readexaclty()

Since read(N) is not guaranteed to read exactly N bytes
---
 secret_handshake/boxstream.py      | 10 ++++++----
 secret_handshake/network.py        |  8 ++++----
 secret_handshake/test_boxstream.py |  1 +
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/secret_handshake/boxstream.py b/secret_handshake/boxstream.py
index f3bb575..b4739b3 100644
--- a/secret_handshake/boxstream.py
+++ b/secret_handshake/boxstream.py
@@ -1,4 +1,6 @@
 import struct
+from asyncio import IncompleteReadError
+
 from nacl.secret import SecretBox
 
 from .util import split_chunks, inc_nonce
@@ -32,9 +34,9 @@ class UnboxStream(object):
         self.closed = False
 
     async def read(self):
-        data = await self.reader.read(HEADER_LENGTH)
-
-        if not data:
+        try:
+            data = await self.reader.readexactly(HEADER_LENGTH)
+        except IncompleteReadError:
             self.closed = True
             return None
 
@@ -49,7 +51,7 @@ class UnboxStream(object):
         length = struct.unpack('>H', header[:2])[0]
         mac = header[2:]
 
-        data = await self.reader.read(length)
+        data = await self.reader.readexactly(length)
 
         body = box.decrypt(mac + data, inc_nonce(self.nonce))
 
diff --git a/secret_handshake/network.py b/secret_handshake/network.py
index 3a06711..21495f9 100644
--- a/secret_handshake/network.py
+++ b/secret_handshake/network.py
@@ -59,13 +59,13 @@ class SHSServer(SHSSocket):
         self.crypto = SHSServerCrypto(server_kp, application_key=application_key)
 
     async def _handshake(self, reader, writer):
-        data = await reader.read(64)
+        data = await reader.readexactly(64)
         if not self.crypto.verify_challenge(data):
             raise SHSClientException('Client challenge is not valid')
 
         writer.write(self.crypto.generate_challenge())
 
-        data = await reader.read(112)
+        data = await reader.readexactly(112)
         if not self.crypto.verify_client_auth(data):
             raise SHSClientException('Client auth is not valid')
 
@@ -99,13 +99,13 @@ class SHSClient(SHSSocket):
     async def _handshake(self, reader, writer):
         writer.write(self.crypto.generate_challenge())
 
-        data = await reader.read(64)
+        data = await reader.readexactly(64)
         if not self.crypto.verify_server_challenge(data):
             raise SHSClientException('Server challenge is not valid')
 
         writer.write(self.crypto.generate_client_auth())
 
-        data = await reader.read(80)
+        data = await reader.readexactly(80)
         if not self.crypto.verify_server_accept(data):
             raise SHSClientException('Server accept is not valid')
 
diff --git a/secret_handshake/test_boxstream.py b/secret_handshake/test_boxstream.py
index 98f3a66..fe2db6b 100644
--- a/secret_handshake/test_boxstream.py
+++ b/secret_handshake/test_boxstream.py
@@ -37,6 +37,7 @@ class AsyncBuffer(BytesIO):
     """Just a BytesIO with an async read method."""
     async def read(self, n=None):
         return super(AsyncBuffer, self).read(n)
+    readexactly = read
 
 
 @pytest.mark.asyncio