Allow limits for maximum payload

src/Messaging/Message.php

@@ -7,8 +7,14 @@ class Message implements \IteratorAggregate, MessageInterface {
      */
     private $_frames;
 
+    /**
+     * @var int
+     */
+    private $len;
+
     public function __construct() {
         $this->_frames = new \SplDoublyLinkedList;
+        $this->len = 0;
     }
 
     public function getIterator() {
@@ -39,6 +45,7 @@ class Message implements \IteratorAggregate, MessageInterface {
      * {@inheritdoc}
      */
     public function addFrame(FrameInterface $fragment) {
+        $this->len += $fragment->getPayloadLength();
         $this->_frames->push($fragment);
 
         return $this;
@@ -59,17 +66,7 @@ class Message implements \IteratorAggregate, MessageInterface {
      * {@inheritdoc}
      */
     public function getPayloadLength() {
-        $len = 0;
+        return $this->len;
-
-        foreach ($this->_frames as $frame) {
-            try {
-                $len += $frame->getPayloadLength();
-            } catch (\UnderflowException $e) {
-                // Not an error, want the current amount buffered
-            }
-        }
-
-        return $len;
     }
 
     /**

src/Messaging/MessageBuffer.php

@@ -42,12 +42,24 @@ class MessageBuffer {
      */
     private $leftovers;
 
+    /**
+     * @var int
+     */
+    private $maxMessagePayloadSize;
+
+    /**
+     * @var int
+     */
+    private $maxFramePayloadSize;
+
     function __construct(
         CloseFrameChecker $frameChecker,
         callable $onMessage,
         callable $onControl = null,
         $expectMask = true,
-        $exceptionFactory = null
+        $exceptionFactory = null,
+        $maxMessagePayloadSize = null, // null for default - zero for no limit
+        $maxFramePayloadSize = null    // null for default - zero for no limit
     ) {
         $this->closeFrameChecker = $frameChecker;
         $this->checkForMask = (bool)$expectMask;
@@ -60,6 +72,31 @@ class MessageBuffer {
         $this->onControl = $onControl ?: function() {};
 
         $this->leftovers = '';
+
+        $memory_limit = \trim(\ini_get('memory_limit'));
+        $memory_limit_bytes = 0;
+        if ($memory_limit !== '') {
+            $shifty = ['k' => 0, 'm' => 10, 'g' => 20];
+            $multiplier = strtolower($memory_limit)[-1];
+            $memory_limit = (int)$memory_limit;
+            $memory_limit_bytes = in_array($multiplier, array_keys($shifty), true) ? $memory_limit * 1024 << $shifty[$multiplier] : $memory_limit;
+        }
+        if ($maxMessagePayloadSize === null) {
+            $maxMessagePayloadSize = $memory_limit_bytes / 4;
+        }
+        if ($maxFramePayloadSize === null) {
+            $maxFramePayloadSize = $memory_limit_bytes / 4;
+        }
+
+        if (!is_int($maxFramePayloadSize) || $maxFramePayloadSize > 0x7FFFFFFFFFFFFFFF || $maxFramePayloadSize < 0) { // this should be interesting on non-64 bit systems
+            throw New \InvalidArgumentException('maxFramePayloadSize is not valid');
+        }
+        $this->maxFramePayloadSize = $maxFramePayloadSize;
+
+        if (!is_int($maxMessagePayloadSize) || $maxMessagePayloadSize > 0x7FFFFFFFFFFFFFFF || $maxMessagePayloadSize < 0) {
+            throw New \InvalidArgumentException('maxMessagePayloadSize is not valid');
+        }
+        $this->maxMessagePayloadSize = $maxMessagePayloadSize;
     }
 
     public function onData($data) {
@@ -90,6 +127,29 @@ class MessageBuffer {
                     : unpack('J', $bytesToUpack)[1];
             }
 
+            $closeFrame = null;
+
+            if ($payload_length < 0) {
+                // this can happen when unpacking in php
+                $closeFrame = $this->newCloseFrame(Frame::CLOSE_PROTOCOL, 'Invalid frame length');
+            }
+
+            if (!$closeFrame && $this->maxFramePayloadSize > 1 && $payload_length > $this->maxFramePayloadSize) {
+                $closeFrame = $this->newCloseFrame(Frame::CLOSE_TOO_BIG, 'Maximum frame size exceeded');
+            }
+
+            if (!$closeFrame && $this->maxMessagePayloadSize > 0
+                && $payload_length + ($this->messageBuffer ? $this->messageBuffer->getPayloadLength() : 0) > $this->maxMessagePayloadSize) {
+                $closeFrame = $this->newCloseFrame(Frame::CLOSE_TOO_BIG, 'Maximum message size exceeded');
+            }
+
+            if ($closeFrame !== null) {
+                $onControl = $this->onControl;
+                $onControl($closeFrame);
+                $this->leftovers = '';
+                return;
+            }
+
             $isCoalesced = $dataLen - $frameStart >= $payload_length + $headerSize;
             if (!$isCoalesced) {
                 break;

tests/unit/Messaging/MessageBufferTest.php

@@ -69,4 +69,158 @@ class MessageBufferTest extends \PHPUnit_Framework_TestCase
 
         $this->assertTrue($bReceived);
     }
+
+    public function testInvalidFrameLength() {
+        $frame = new Frame(str_repeat('a', 200), true, Frame::OP_TEXT);
+
+        $frameRaw = $frame->getContents();
+
+        $frameRaw[1] = "\x7f"; // 127 in the first spot
+
+        $frameRaw[2] = "\xff"; // this will unpack to -1
+        $frameRaw[3] = "\xff";
+        $frameRaw[4] = "\xff";
+        $frameRaw[5] = "\xff";
+        $frameRaw[6] = "\xff";
+        $frameRaw[7] = "\xff";
+        $frameRaw[8] = "\xff";
+        $frameRaw[9] = "\xff";
+
+        /** @var Frame $controlFrame */
+        $controlFrame = null;
+        $messageCount = 0;
+
+        $messageBuffer = new MessageBuffer(
+            new CloseFrameChecker(),
+            function (Message $message) use (&$messageCount) {
+                $messageCount++;
+            },
+            function (Frame $frame) use (&$controlFrame) {
+                $this->assertNull($controlFrame);
+                $controlFrame = $frame;
+            },
+            false,
+            null,
+            0,
+            10
+        );
+
+        $messageBuffer->onData($frameRaw);
+
+        $this->assertEquals(0, $messageCount);
+        $this->assertTrue($controlFrame instanceof Frame);
+        $this->assertEquals(Frame::OP_CLOSE, $controlFrame->getOpcode());
+        $this->assertEquals([Frame::CLOSE_PROTOCOL], array_merge(unpack('n*', substr($controlFrame->getPayload(), 0, 2))));
+
+    }
+
+    public function testFrameLengthTooBig() {
+        $frame = new Frame(str_repeat('a', 200), true, Frame::OP_TEXT);
+
+        $frameRaw = $frame->getContents();
+
+        $frameRaw[1] = "\x7f"; // 127 in the first spot
+
+        $frameRaw[2] = "\x7f"; // this will unpack to -1
+        $frameRaw[3] = "\xff";
+        $frameRaw[4] = "\xff";
+        $frameRaw[5] = "\xff";
+        $frameRaw[6] = "\xff";
+        $frameRaw[7] = "\xff";
+        $frameRaw[8] = "\xff";
+        $frameRaw[9] = "\xff";
+
+        /** @var Frame $controlFrame */
+        $controlFrame = null;
+        $messageCount = 0;
+
+        $messageBuffer = new MessageBuffer(
+            new CloseFrameChecker(),
+            function (Message $message) use (&$messageCount) {
+                $messageCount++;
+            },
+            function (Frame $frame) use (&$controlFrame) {
+                $this->assertNull($controlFrame);
+                $controlFrame = $frame;
+            },
+            false,
+            null,
+            0,
+            10
+        );
+
+        $messageBuffer->onData($frameRaw);
+
+        $this->assertEquals(0, $messageCount);
+        $this->assertTrue($controlFrame instanceof Frame);
+        $this->assertEquals(Frame::OP_CLOSE, $controlFrame->getOpcode());
+        $this->assertEquals([Frame::CLOSE_TOO_BIG], array_merge(unpack('n*', substr($controlFrame->getPayload(), 0, 2))));
+    }
+
+    public function testFrameLengthBiggerThanMaxMessagePayload() {
+        $frame = new Frame(str_repeat('a', 200), true, Frame::OP_TEXT);
+
+        $frameRaw = $frame->getContents();
+
+        /** @var Frame $controlFrame */
+        $controlFrame = null;
+        $messageCount = 0;
+
+        $messageBuffer = new MessageBuffer(
+            new CloseFrameChecker(),
+            function (Message $message) use (&$messageCount) {
+                $messageCount++;
+            },
+            function (Frame $frame) use (&$controlFrame) {
+                $this->assertNull($controlFrame);
+                $controlFrame = $frame;
+            },
+            false,
+            null,
+            100,
+            0
+        );
+
+        $messageBuffer->onData($frameRaw);
+
+        $this->assertEquals(0, $messageCount);
+        $this->assertTrue($controlFrame instanceof Frame);
+        $this->assertEquals(Frame::OP_CLOSE, $controlFrame->getOpcode());
+        $this->assertEquals([Frame::CLOSE_TOO_BIG], array_merge(unpack('n*', substr($controlFrame->getPayload(), 0, 2))));
+    }
+
+    public function testSecondFrameLengthPushesPastMaxMessagePayload() {
+        $frame = new Frame(str_repeat('a', 200), false, Frame::OP_TEXT);
+        $firstFrameRaw = $frame->getContents();
+        $frame = new Frame(str_repeat('b', 200), true, Frame::OP_TEXT);
+        $secondFrameRaw = $frame->getContents();
+
+        /** @var Frame $controlFrame */
+        $controlFrame = null;
+        $messageCount = 0;
+
+        $messageBuffer = new MessageBuffer(
+            new CloseFrameChecker(),
+            function (Message $message) use (&$messageCount) {
+                $messageCount++;
+            },
+            function (Frame $frame) use (&$controlFrame) {
+                $this->assertNull($controlFrame);
+                $controlFrame = $frame;
+            },
+            false,
+            null,
+            300,
+            0
+        );
+
+        $messageBuffer->onData($firstFrameRaw);
+        // only put part of the second frame in to watch it fail fast
+        $messageBuffer->onData(substr($secondFrameRaw, 0, 150));
+
+        $this->assertEquals(0, $messageCount);
+        $this->assertTrue($controlFrame instanceof Frame);
+        $this->assertEquals(Frame::OP_CLOSE, $controlFrame->getOpcode());
+        $this->assertEquals([Frame::CLOSE_TOO_BIG], array_merge(unpack('n*', substr($controlFrame->getPayload(), 0, 2))));
+    }
 }