From d80fbc24107ac2166b80ccb5f4a240a7f480497b Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sun, 17 May 2015 15:38:13 +1200
Subject: [PATCH] websockets: raw body specification

---
 libpathod/language/websockets.py              | 28 +++++++++++++++----
 libpathod/templates/docs_lang_websockets.html |  8 ++++++
 test/test_language_websocket.py               |  8 ++++++
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/libpathod/language/websockets.py b/libpathod/language/websockets.py
index a006c50df..12b59bc67 100644
--- a/libpathod/language/websockets.py
+++ b/libpathod/language/websockets.py
@@ -8,8 +8,8 @@ from . import base, generators, actions, message
     wf:c15:r'foo'
     wf:fin:rsv1:rsv2:rsv3:mask
     wf:-fin:-rsv1:-rsv2:-rsv3:-mask
-
     wf:l234
+
     wf:mask:r"foo
 """
 
@@ -35,6 +35,11 @@ class Body(base.Value):
     preamble = "b"
 
 
+class RawBody(base.Value):
+    unique_name = "Body"
+    preamble = "r"
+
+
 class Fin(base.Boolean):
     name = "fin"
 
@@ -66,14 +71,16 @@ class KeyNone(base.CaselessLiteral):
 
 
 class Length(base.Integer):
-    bounds = (0, 1<<64)
+    bounds = (0, 1 << 64)
     preamble = "l"
 
 
+class Times(base.Integer):
+    preamble = "x"
+
+
 class WebsocketFrame(message.Message):
     comps = (
-        Body,
-
         OpCode,
         Length,
         # Bit flags
@@ -87,6 +94,10 @@ class WebsocketFrame(message.Message):
         actions.InjectAt,
         KeyNone,
         Key,
+        Times,
+
+        Body,
+        RawBody,
     )
     logattrs = ["body"]
     @property
@@ -97,6 +108,10 @@ class WebsocketFrame(message.Message):
     def body(self):
         return self.tok(Body)
 
+    @property
+    def rawbody(self):
+        return self.tok(RawBody)
+
     @property
     def opcode(self):
         return self.tok(OpCode)
@@ -165,6 +180,9 @@ class WebsocketFrame(message.Message):
         if self.body:
             bodygen = self.body.value.get_generator(settings)
             length = len(self.body.value.get_generator(settings))
+        elif self.rawbody:
+            bodygen = self.rawbody.value.get_generator(settings)
+            length = len(self.rawbody.value.get_generator(settings))
         else:
             bodygen = None
             length = 0
@@ -187,7 +205,7 @@ class WebsocketFrame(message.Message):
         frame = netlib.websockets.FrameHeader(**frameparts)
         vals = [frame.to_bytes()]
         if bodygen:
-            if frame.masking_key:
+            if frame.masking_key and not self.rawbody:
                 masker = netlib.websockets.Masker(frame.masking_key)
                 vals.append(
                     generators.TransformGenerator(
diff --git a/libpathod/templates/docs_lang_websockets.html b/libpathod/templates/docs_lang_websockets.html
index b21f84a1b..8cd176fc3 100644
--- a/libpathod/templates/docs_lang_websockets.html
+++ b/libpathod/templates/docs_lang_websockets.html
@@ -76,6 +76,14 @@
             </td>
         </tr>
 
+        <tr>
+            <td> r<a href="#valuespec">VALUE</a> </td>
+            <td>
+                Set the raw frame payload. This diesables masking, even if the
+                key is present.
+            </td>
+        </tr>
+
         <tr>
             <td> [-]rsv1 </td>
             <td>
diff --git a/test/test_language_websocket.py b/test/test_language_websocket.py
index 0228c350b..faf9299fb 100644
--- a/test/test_language_websocket.py
+++ b/test/test_language_websocket.py
@@ -15,6 +15,7 @@ class TestWebsocketFrame:
             "wf",
             "wf:dr",
             "wf:b'foo'",
+            "wf:mask:r'foo'",
             "wf:l1024:b'foo'",
             "wf:cbinary",
             "wf:c1",
@@ -68,6 +69,13 @@ class TestWebsocketFrame:
         assert self.fr("wf:ctext").header.opcode ==\
             netlib.websockets.OPCODE.TEXT
 
+    def test_rawbody(self):
+        frm = self.fr("wf:mask:r'foo'")
+        assert len(frm.payload) == 3
+        assert frm.payload != "foo"
+
+        assert self.fr("wf:r'foo'").payload == "foo"
+
     def test_construction(self):
         # Simple server frame
         frm = self.fr("wf:b'foo'")