sockhandler.h 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. #pragma once
  2. #include <library/cpp/logger/all.h>
  3. #include <util/generic/buffer.h>
  4. #include <util/generic/map.h>
  5. #include <util/generic/vector.h>
  6. #include <util/network/address.h>
  7. #include <util/network/ip.h>
  8. #include <util/network/socket.h>
  9. #include <util/system/mutex.h>
  10. #include <util/system/yassert.h>
  11. #include <cerrno>
  12. #include <util/generic/noncopyable.h>
  13. class TAddrList: public TVector<NAddr::IRemoteAddrRef> {
  14. private:
  15. using TBase = TVector<NAddr::IRemoteAddrRef>;
  16. public:
  17. //msvc doesn't support base class constructor inheritance
  18. TAddrList() = default;
  19. template <typename T>
  20. TAddrList(T&& arg)
  21. : TBase(std::forward<T>(arg))
  22. {
  23. }
  24. template <typename T1, typename T2>
  25. TAddrList(T1&& arg1, T2&& arg2)
  26. : TBase(std::forward<T1>(arg1), std::forward<T2>(arg2))
  27. {
  28. }
  29. TAddrList(std::initializer_list<NAddr::IRemoteAddrRef> list)
  30. : TBase(list)
  31. {
  32. }
  33. static TAddrList MakeV4Addr(ui32 ip, TIpPort port) {
  34. return TAddrList({new NAddr::TIPv4Addr(TIpAddress(htonl(ip), htons(port)))});
  35. }
  36. std::pair<ui32, TIpPort> GetV4Addr() const {
  37. for (const auto& addrRef : *this) {
  38. const sockaddr* sa = addrRef->Addr();
  39. if (sa->sa_family == AF_INET) {
  40. const sockaddr_in* sin = reinterpret_cast<const sockaddr_in*>(sa);
  41. return std::make_pair(ntohl(sin->sin_addr.s_addr), ntohs(sin->sin_port));
  42. }
  43. }
  44. return std::make_pair(0, 0);
  45. }
  46. };
  47. class TSimpleSocketHandler {
  48. public:
  49. TSimpleSocketHandler() = default;
  50. int Good() const {
  51. return static_cast<bool>(Socket);
  52. }
  53. int Connect(const TAddrList& addrs, TDuration timeout) {
  54. try {
  55. for (const auto& item : addrs) {
  56. const sockaddr* sa = item->Addr();
  57. TSocketHolder s(socket(sa->sa_family, SOCK_STREAM, 0));
  58. if (s.Closed()) {
  59. continue;
  60. }
  61. #ifndef WIN32
  62. if (fcntl(s, F_SETFD, FD_CLOEXEC)) // no inherit on fork()/exec()
  63. return errno ? errno : EBADF;
  64. #endif
  65. if (connect(s, sa, item->Len())) {
  66. s.Close();
  67. continue;
  68. }
  69. Socket.Reset(new TSocket(s.Release()));
  70. Socket->SetSocketTimeout(timeout.Seconds(), timeout.MilliSecondsOfSecond());
  71. Socket->SetZeroLinger();
  72. Socket->SetKeepAlive(true);
  73. return 0;
  74. }
  75. } catch (...) {
  76. return EBADF;
  77. }
  78. return errno ? errno : EBADF;
  79. }
  80. void Disconnect() {
  81. if (!Socket)
  82. return;
  83. Socket->ShutDown(SHUT_RDWR);
  84. Socket.Destroy();
  85. }
  86. void SetSocket(SOCKET fd) {
  87. Socket.Reset(new TSocket(fd));
  88. }
  89. void shutdown() {
  90. Socket->ShutDown(SHUT_WR);
  91. }
  92. int send(const void* message, size_t messlen) {
  93. return ((ssize_t)messlen == Socket->Send(message, messlen));
  94. }
  95. int peek() {
  96. char buf[1];
  97. return (1 == recv(*Socket, buf, 1, MSG_PEEK));
  98. }
  99. ssize_t read(void* buffer, size_t buflen) {
  100. return Socket->Recv(buffer, buflen);
  101. }
  102. THolder<TSocket> PickOutSocket() {
  103. return std::move(Socket);
  104. }
  105. protected:
  106. THolder<TSocket> Socket;
  107. };