Merge branch 'maxpayload'

ratchetphp · Dec 15, 2019 · 879e48c · 879e48c
2 parents c033a39 + e861242
commit 879e48c
Show file tree

Hide file tree

Showing 3 changed files with 381 additions and 13 deletions.
diff --git a/src/Messaging/Message.php b/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 @@ public function isCoalesced() {
      * {@inheritdoc}
      */
     public function addFrame(FrameInterface $fragment) {
+        $this->len += $fragment->getPayloadLength();
         $this->_frames->push($fragment);
 
         return $this;
@@ -59,17 +66,7 @@ public function getOpcode() {
      * {@inheritdoc}
      */
     public function getPayloadLength() {
-        $len = 0;
-
-        foreach ($this->_frames as $frame) {
-            try {
-                $len += $frame->getPayloadLength();
-            } catch (\UnderflowException $e) {
-                // Not an error, want the current amount buffered
-            }
-        }
-
-        return $len;
+        return $this->len;
     }
 
     /**

diff --git a/src/Messaging/MessageBuffer.php b/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,25 @@ function __construct(
         $this->onControl = $onControl ?: function() {};
 
         $this->leftovers = '';
+
+        $memory_limit_bytes = static::getMemoryLimit();
+
+        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 a valid maxFramePayloadSize');
+        }
+        $this->maxFramePayloadSize = $maxFramePayloadSize;
+
+        if (!is_int($maxMessagePayloadSize) || $maxMessagePayloadSize > 0x7FFFFFFFFFFFFFFF || $maxMessagePayloadSize < 0) {
+            throw new \InvalidArgumentException($maxMessagePayloadSize . 'is not a valid maxMessagePayloadSize');
+        }
+        $this->maxMessagePayloadSize = $maxMessagePayloadSize;
     }
 
     public function onData($data) {
@@ -68,6 +99,7 @@ public function onData($data) {
 
         if ($dataLen < 2) {
             $this->leftovers = $data;
+
             return;
         }
 
@@ -90,6 +122,30 @@ public function onData($data) {
                     : 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;
@@ -264,4 +320,24 @@ public function newFrame($payload = null, $final = null, $opcode = null) {
     public function newCloseFrame($code, $reason = '') {
         return $this->newFrame(pack('n', $code) . $reason, true, Frame::OP_CLOSE);
     }
-}
+
+    /**
+     * This is a separate function for testing purposes
+     * $memory_limit is only used for testing
+     *
+     * @param null|string $memory_limit
+     * @return int
+     */
+    private static function getMemoryLimit($memory_limit = null) {
+        $memory_limit = $memory_limit === null ? \trim(\ini_get('memory_limit')) : $memory_limit;
+        $memory_limit_bytes = 0;
+        if ($memory_limit !== '') {
+            $shifty = ['k' => 0, 'm' => 10, 'g' => 20];
+            $multiplier = strlen($memory_limit) > 1 ? substr(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;
+        }
+
+        return $memory_limit_bytes < 0 ? 0 : $memory_limit_bytes;
+    }
+}