websockets: replace masker with more optimized one

1 files changed, 10 insertions, 6 deletions

diff --git a/mitmproxy/net/websockets/masker.py b/mitmproxy/net/websockets/masker.py
index 47b1a688..6134e09e 100644
--- a/mitmproxy/net/websockets/masker.py
+++ b/mitmproxy/net/websockets/masker.py
@@ -1,3 +1,6 @@
+import sys
+
+
 class Masker:
     """
     Data sent from the server must be masked to prevent malicious clients
@@ -12,12 +15,13 @@ class Masker:
         self.offset = 0
 
     def mask(self, offset, data):
-        result = bytearray(data)
-        for i in range(len(data)):
-            result[i] ^= self.key[offset % 4]
-            offset += 1
-        result = bytes(result)
-        return result
+        datalen = len(data)
+        offset_mod = offset % 4
+        data = int.from_bytes(data, sys.byteorder)
+        num_keys = (datalen + offset_mod + 3) // 4
+        mask = int.from_bytes((self.key * num_keys)[offset_mod:datalen +
+                                                    offset_mod], sys.byteorder)
+        return (data ^ mask).to_bytes(datalen, sys.byteorder)
 
     def __call__(self, data):
         ret = self.mask(self.offset, data)