Fork of the espurna firmware for `mhsw` switches
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

83 lines
2.9 KiB

  1. /*
  2. WebSocketIncommingBuffer
  3. Code by Hermann Kraus (https://bitbucket.org/hermr2d2/)
  4. and slightly modified.
  5. https://bitbucket.org/xoseperez/espurna/pull-requests/30/safer-buffer-handling-for-websocket-data
  6. */
  7. #pragma once
  8. #define MAX_WS_MSG_SIZE 4000
  9. typedef std::function<void(AsyncWebSocketClient *client, uint8_t *data, size_t len)> AwsMessageHandler;
  10. class WebSocketIncommingBuffer {
  11. public:
  12. WebSocketIncommingBuffer(AwsMessageHandler cb, bool terminate_string = true, bool cb_on_fragments = false) :
  13. _cb(cb),
  14. _cb_on_fragments(cb_on_fragments),
  15. _terminate_string(terminate_string),
  16. _buffer(0)
  17. {}
  18. ~WebSocketIncommingBuffer() {
  19. if (_buffer) delete _buffer;
  20. }
  21. void data_event(AsyncWebSocketClient *client, AwsFrameInfo *info, uint8_t *data, size_t len) {
  22. if((info->final || _cb_on_fragments) &&
  23. !_terminate_string && info->index == 0 && info->len == len) {
  24. /* The whole message is in a single frame and we got all of it's
  25. data therefore we can parse it without copying the data first.*/
  26. _cb(client, data, len);
  27. } else {
  28. if (info->len > MAX_WS_MSG_SIZE) return;
  29. /* Check if previous fragment was discarded because it was too long. */
  30. if (!_cb_on_fragments && info->num > 0 && !_buffer) return;
  31. if (!_buffer) {
  32. _buffer = new std::vector<uint8_t>();
  33. }
  34. if (info->index == 0) {
  35. //New frame => preallocate memory
  36. if (_cb_on_fragments) {
  37. _buffer->reserve(info->len + 1);
  38. } else {
  39. /* The current fragment would lead to a message which is
  40. too long. So discard everything received so far. */
  41. if (info->len + _buffer->size() > MAX_WS_MSG_SIZE) {
  42. delete _buffer;
  43. _buffer = 0;
  44. return;
  45. } else {
  46. _buffer->reserve(info->len + _buffer->size() + 1);
  47. }
  48. }
  49. }
  50. //assert(_buffer->size() == info->index);
  51. _buffer->insert(_buffer->end(), data, data+len);
  52. if (info->index + len == info->len &&
  53. (info->final || _cb_on_fragments)) {
  54. // Frame/message complete
  55. if (_terminate_string) {
  56. _buffer->push_back(0);
  57. }
  58. _cb(client, _buffer->data(), _buffer->size());
  59. _buffer->clear();
  60. }
  61. }
  62. }
  63. private:
  64. AwsMessageHandler _cb;
  65. bool _cb_on_fragments;
  66. bool _terminate_string;
  67. std::vector<uint8_t> *_buffer;
  68. };