1
0

WebSocketChannel.h 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. #ifndef HV_WEBSOCKET_CHANNEL_H_
  2. #define HV_WEBSOCKET_CHANNEL_H_
  3. #include <mutex>
  4. #include "Channel.h"
  5. #include "wsdef.h"
  6. #include "hmath.h"
  7. namespace hv {
  8. class WebSocketChannel : public SocketChannel {
  9. public:
  10. ws_session_type type;
  11. WebSocketChannel(hio_t* io, ws_session_type type = WS_CLIENT)
  12. : SocketChannel(io)
  13. , type(type)
  14. {}
  15. ~WebSocketChannel() {}
  16. // isConnected, send, close
  17. int send(const std::string& msg, enum ws_opcode opcode = WS_OPCODE_TEXT, bool fin = true) {
  18. return send(msg.c_str(), msg.size(), opcode, fin);
  19. }
  20. int send(const char* buf, int len, enum ws_opcode opcode = WS_OPCODE_BINARY, bool fin = true) {
  21. bool has_mask = false;
  22. char mask[4] = {0};
  23. if (type == WS_CLIENT) {
  24. has_mask = true;
  25. *(int*)mask = rand();
  26. }
  27. int frame_size = ws_calc_frame_size(len, has_mask);
  28. std::lock_guard<std::mutex> locker(mutex_);
  29. if (sendbuf_.len < frame_size) {
  30. sendbuf_.resize(ceil2e(frame_size));
  31. }
  32. ws_build_frame(sendbuf_.base, buf, len, mask, has_mask, opcode, fin);
  33. return write(sendbuf_.base, frame_size);
  34. }
  35. // websocket fragment
  36. // send(p, fragment, opcode, false) ->
  37. // send(p, fragment, WS_OPCODE_CONTINUE, false) ->
  38. // ... ->
  39. // send(p, remain, WS_OPCODE_CONTINUE, true)
  40. int send(const char* buf, int len, int fragment, enum ws_opcode opcode = WS_OPCODE_BINARY) {
  41. if (len <= fragment) {
  42. return send(buf, len, opcode, true);
  43. }
  44. // first fragment
  45. int nsend = send(buf, fragment, opcode, false);
  46. if (nsend < 0) return nsend;
  47. const char* p = buf + fragment;
  48. int remain = len - fragment;
  49. while (remain > fragment) {
  50. nsend = send(p, fragment, WS_OPCODE_CONTINUE, false);
  51. if (nsend < 0) return nsend;
  52. p += fragment;
  53. remain -= fragment;
  54. }
  55. // last fragment
  56. nsend = send(p, remain, WS_OPCODE_CONTINUE, true);
  57. if (nsend < 0) return nsend;
  58. return len;
  59. }
  60. private:
  61. Buffer sendbuf_;
  62. std::mutex mutex_;
  63. };
  64. }
  65. typedef std::shared_ptr<hv::WebSocketChannel> WebSocketChannelPtr;
  66. #endif // HV_WEBSOCKET_CHANNEL_H_