From 0a7da6a9b1b878cbfabf437759031f816161fab7 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Thu, 4 Jun 2015 23:57:23 +1200
Subject: [PATCH] Rudimentary support for reflected websocket frames.

---
 libpathod/language/__init__.py   | 10 ++++++++
 libpathod/language/websockets.py | 12 ++++++---
 libpathod/pathod.py              | 42 +++++++++++++++++++++++---------
 test/test_pathod.py              |  9 +++++--
 4 files changed, 55 insertions(+), 18 deletions(-)

diff --git a/libpathod/language/__init__.py b/libpathod/language/__init__.py
index 123786078..52ccf61c9 100644
--- a/libpathod/language/__init__.py
+++ b/libpathod/language/__init__.py
@@ -57,6 +57,16 @@ def parse_pathoc(s):
     return expanded
 
 
+def parse_websocket_frame(s):
+    try:
+        return websockets.WebsocketFrame.expr().parseString(
+            s,
+            parseAll = True
+        )[0]
+    except pp.ParseException as v:
+        raise exceptions.ParseException(v.msg, v.line, v.col)
+
+
 def serve(msg, fp, settings):
     """
         fp: The file pointer to write to.
diff --git a/libpathod/language/websockets.py b/libpathod/language/websockets.py
index 44c12f645..96dbce530 100644
--- a/libpathod/language/websockets.py
+++ b/libpathod/language/websockets.py
@@ -3,6 +3,7 @@ import netlib.websockets
 import pyparsing as pp
 from . import base, generators, actions, message
 
+NESTED_LEADER = "pathod!"
 
 class WF(base.CaselessLiteral):
     TOK = "wf"
@@ -160,6 +161,10 @@ class WebsocketFrame(message.Message):
         resp = resp.setParseAction(klass)
         return resp
 
+    @property
+    def nested_frame(self):
+        return self.tok(NestedFrame)
+
     def resolve(self, settings, msg=None):
         tokens = self.tokens[:]
         if not self.mask and settings.is_client:
@@ -181,6 +186,9 @@ class WebsocketFrame(message.Message):
         elif self.rawbody:
             bodygen = self.rawbody.value.get_generator(settings)
             length = len(self.rawbody.value.get_generator(settings))
+        elif self.nested_frame:
+            bodygen = NESTED_LEADER + self.nested_frame.parsed.spec()
+            length = len(bodygen)
         else:
             bodygen = None
             length = 0
@@ -228,7 +236,3 @@ class WebsocketClientFrame(WebsocketFrame):
     components = COMPONENTS + (
         NestedFrame,
     )
-
-    @property
-    def nested_frame(self):
-        return self.tok(NestedFrame)
diff --git a/libpathod/pathod.py b/libpathod/pathod.py
index 6b385a47a..abce15fe7 100644
--- a/libpathod/pathod.py
+++ b/libpathod/pathod.py
@@ -11,6 +11,8 @@ from netlib import tcp, http, wsgi, certutils, websockets
 from . import version, app, language, utils, log
 import language.http
 import language.actions
+import language.exceptions
+import language.websockets
 
 
 DEFAULT_CERT_DOMAIN = "pathod.net"
@@ -102,21 +104,37 @@ class PathodHandler(tcp.BaseHandler):
     def handle_websocket(self):
         lr = self.rfile if self.server.logreq else None
         lw = self.wfile if self.server.logresp else None
-        with log.Log(self.logfp, self.server.hexdump, lr, lw) as lg:
-            while True:
+        while True:
+            with log.Log(self.logfp, self.server.hexdump, lr, lw) as lg:
                 try:
                     frm = websockets.Frame.from_file(self.rfile)
-                    retlog = dict(
-                        type="wsframe",
-                        frame=dict(
-                        ),
-                        cipher=None,
-                    )
-                    self.addlog(retlog)
+                except tcp.NetLibIncomplete, e:
+                    lg("Error reading websocket frame: %s"%e)
                     break
-                except tcp.NetLibTimeout: # pragma: no cover
-                    pass
-            lg(frm.human_readable())
+                lg(frm.human_readable())
+                retlog = dict(
+                    type="wsframe",
+                    frame=dict(
+                    ),
+                    cipher=None,
+                )
+                ld = language.websockets.NESTED_LEADER
+                if frm.payload.startswith(ld):
+                    nest = frm.payload[len(ld):]
+                    try:
+                        wf = language.parse_websocket_frame(nest)
+                    except language.exceptions.ParseException, v:
+                        lg(
+                            "Parse error in reflected frame specifcation:"
+                            " %s" % v.msg
+                        )
+                        break
+                    frame_log = language.serve(
+                        wf,
+                        self.wfile,
+                        self.settings
+                    )
+                self.addlog(retlog)
         return self.handle_websocket, None
 
     def handle_http_connect(self, connect, lg):
diff --git a/test/test_pathod.py b/test/test_pathod.py
index 43fb8368e..c7ea47e3c 100644
--- a/test/test_pathod.py
+++ b/test/test_pathod.py
@@ -2,6 +2,7 @@ import cStringIO
 from libpathod import pathod, version
 from netlib import tcp, http
 import time
+import sys
 
 import tutils
 
@@ -211,10 +212,14 @@ class CommonTests(tutils.DaemonTests):
 
     def test_websocket_frame(self):
         r = self.pathoc(["ws:/p/", "wf:b@10"], ws_read_limit=0)
-        print r
-        print self.d.log()
+        print >> sys.stderr, r
+        print >> sys.stderr, self.d.log()
         assert self.d.last_log()["type"] == "wsframe"
 
+    def test_websocket_reflected_frame(self):
+        r = self.pathoc(["ws:/p/", "wf:f'wf'"], ws_read_limit=0)
+        assert r
+
 
 class TestDaemon(CommonTests):
     ssl = False