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.
 
 
 
 

390 regels
11 KiB

  1. @cython.no_gc_clear
  2. @cython.freelist(DEFAULT_FREELIST_SIZE)
  3. cdef class _UDPSendContext:
  4. # used to hold additional write request information for uv_write
  5. cdef:
  6. uv.uv_udp_send_t req
  7. uv.uv_buf_t uv_buf
  8. Py_buffer py_buf
  9. UDPTransport udp
  10. bint closed
  11. cdef close(self):
  12. if self.closed:
  13. return
  14. self.closed = 1
  15. PyBuffer_Release(&self.py_buf) # void
  16. self.req.data = NULL
  17. self.uv_buf.base = NULL
  18. Py_DECREF(self)
  19. self.udp = None
  20. @staticmethod
  21. cdef _UDPSendContext new(UDPTransport udp, object data):
  22. cdef _UDPSendContext ctx
  23. ctx = _UDPSendContext.__new__(_UDPSendContext)
  24. ctx.udp = None
  25. ctx.closed = 1
  26. ctx.req.data = <void*> ctx
  27. Py_INCREF(ctx)
  28. PyObject_GetBuffer(data, &ctx.py_buf, PyBUF_SIMPLE)
  29. ctx.uv_buf.base = <char*>ctx.py_buf.buf
  30. ctx.uv_buf.len = ctx.py_buf.len
  31. ctx.udp = udp
  32. ctx.closed = 0
  33. return ctx
  34. def __dealloc__(self):
  35. if UVLOOP_DEBUG:
  36. if not self.closed:
  37. raise RuntimeError(
  38. 'open _UDPSendContext is being deallocated')
  39. self.udp = None
  40. @cython.no_gc_clear
  41. cdef class UDPTransport(UVBaseTransport):
  42. def __cinit__(self):
  43. self._family = uv.AF_UNSPEC
  44. self.__receiving = 0
  45. cdef _init(self, Loop loop, unsigned int family):
  46. cdef int err
  47. self._start_init(loop)
  48. self._handle = <uv.uv_handle_t*>PyMem_RawMalloc(sizeof(uv.uv_udp_t))
  49. if self._handle is NULL:
  50. self._abort_init()
  51. raise MemoryError()
  52. err = uv.uv_udp_init_ex(loop.uvloop,
  53. <uv.uv_udp_t*>self._handle,
  54. family)
  55. if err < 0:
  56. self._abort_init()
  57. raise convert_error(err)
  58. if family in (uv.AF_INET, uv.AF_INET6):
  59. self._family = family
  60. self._finish_init()
  61. cdef _connect(self, system.sockaddr* addr, size_t addr_len):
  62. cdef int err
  63. err = uv.uv_udp_connect(<uv.uv_udp_t*>self._handle, addr)
  64. if err < 0:
  65. exc = convert_error(err)
  66. raise exc
  67. cdef open(self, int family, int sockfd):
  68. if family in (uv.AF_INET, uv.AF_INET6, uv.AF_UNIX):
  69. self._family = family
  70. else:
  71. raise ValueError(
  72. 'cannot open a UDP handle, invalid family {}'.format(family))
  73. cdef int err
  74. err = uv.uv_udp_open(<uv.uv_udp_t*>self._handle,
  75. <uv.uv_os_sock_t>sockfd)
  76. if err < 0:
  77. exc = convert_error(err)
  78. raise exc
  79. cdef _bind(self, system.sockaddr* addr, bint reuse_addr):
  80. cdef:
  81. int err
  82. int flags = 0
  83. self._ensure_alive()
  84. if reuse_addr:
  85. flags |= uv.UV_UDP_REUSEADDR
  86. err = uv.uv_udp_bind(<uv.uv_udp_t*>self._handle, addr, flags)
  87. if err < 0:
  88. exc = convert_error(err)
  89. raise exc
  90. cdef _set_broadcast(self, bint on):
  91. cdef int err
  92. self._ensure_alive()
  93. err = uv.uv_udp_set_broadcast(<uv.uv_udp_t*>self._handle, on)
  94. if err < 0:
  95. exc = convert_error(err)
  96. raise exc
  97. cdef size_t _get_write_buffer_size(self):
  98. if self._handle is NULL:
  99. return 0
  100. return (<uv.uv_udp_t*>self._handle).send_queue_size
  101. cdef bint _is_reading(self):
  102. return self.__receiving
  103. cdef _start_reading(self):
  104. cdef int err
  105. if self.__receiving:
  106. return
  107. self._ensure_alive()
  108. err = uv.uv_udp_recv_start(<uv.uv_udp_t*>self._handle,
  109. __loop_alloc_buffer,
  110. __uv_udp_on_receive)
  111. if err < 0:
  112. exc = convert_error(err)
  113. self._fatal_error(exc, True)
  114. return
  115. else:
  116. # UDPTransport must live until the read callback is called
  117. self.__receiving_started()
  118. cdef _stop_reading(self):
  119. cdef int err
  120. if not self.__receiving:
  121. return
  122. self._ensure_alive()
  123. err = uv.uv_udp_recv_stop(<uv.uv_udp_t*>self._handle)
  124. if err < 0:
  125. exc = convert_error(err)
  126. self._fatal_error(exc, True)
  127. return
  128. else:
  129. self.__receiving_stopped()
  130. cdef inline __receiving_started(self):
  131. if self.__receiving:
  132. return
  133. self.__receiving = 1
  134. Py_INCREF(self)
  135. cdef inline __receiving_stopped(self):
  136. if not self.__receiving:
  137. return
  138. self.__receiving = 0
  139. Py_DECREF(self)
  140. cdef _new_socket(self):
  141. if self._family not in (uv.AF_INET, uv.AF_INET6, uv.AF_UNIX):
  142. raise RuntimeError(
  143. 'UDPTransport.family is undefined; '
  144. 'cannot create python socket')
  145. fileno = self._fileno()
  146. return PseudoSocket(self._family, uv.SOCK_DGRAM, 0, fileno)
  147. cdef _send(self, object data, object addr):
  148. cdef:
  149. _UDPSendContext ctx
  150. system.sockaddr_storage saddr_st
  151. system.sockaddr *saddr
  152. Py_buffer try_pybuf
  153. uv.uv_buf_t try_uvbuf
  154. self._ensure_alive()
  155. if self._family not in (uv.AF_INET, uv.AF_INET6, uv.AF_UNIX):
  156. raise RuntimeError('UDPTransport.family is undefined; cannot send')
  157. if addr is None:
  158. saddr = NULL
  159. else:
  160. try:
  161. __convert_pyaddr_to_sockaddr(self._family, addr,
  162. <system.sockaddr*>&saddr_st)
  163. except (ValueError, TypeError):
  164. raise
  165. except Exception:
  166. raise ValueError(
  167. f'{addr!r}: socket family mismatch or '
  168. f'a DNS lookup is required')
  169. saddr = <system.sockaddr*>(&saddr_st)
  170. if self._get_write_buffer_size() == 0:
  171. PyObject_GetBuffer(data, &try_pybuf, PyBUF_SIMPLE)
  172. try_uvbuf.base = <char*>try_pybuf.buf
  173. try_uvbuf.len = try_pybuf.len
  174. err = uv.uv_udp_try_send(<uv.uv_udp_t*>self._handle,
  175. &try_uvbuf,
  176. 1,
  177. saddr)
  178. PyBuffer_Release(&try_pybuf)
  179. else:
  180. err = uv.UV_EAGAIN
  181. if err == uv.UV_EAGAIN:
  182. ctx = _UDPSendContext.new(self, data)
  183. err = uv.uv_udp_send(&ctx.req,
  184. <uv.uv_udp_t*>self._handle,
  185. &ctx.uv_buf,
  186. 1,
  187. saddr,
  188. __uv_udp_on_send)
  189. if err < 0:
  190. ctx.close()
  191. exc = convert_error(err)
  192. self._fatal_error(exc, True)
  193. else:
  194. self._maybe_pause_protocol()
  195. else:
  196. if err < 0:
  197. exc = convert_error(err)
  198. self._fatal_error(exc, True)
  199. else:
  200. self._on_sent(None)
  201. cdef _on_receive(self, bytes data, object exc, object addr):
  202. if exc is None:
  203. self._protocol.datagram_received(data, addr)
  204. else:
  205. self._protocol.error_received(exc)
  206. cdef _on_sent(self, object exc):
  207. if exc is not None:
  208. if isinstance(exc, OSError):
  209. self._protocol.error_received(exc)
  210. else:
  211. self._fatal_error(
  212. exc, False, 'Fatal write error on datagram transport')
  213. self._maybe_resume_protocol()
  214. if not self._get_write_buffer_size():
  215. if self._closing:
  216. self._schedule_call_connection_lost(None)
  217. # === Public API ===
  218. def sendto(self, data, addr=None):
  219. if not data:
  220. # Replicating asyncio logic here.
  221. return
  222. if self._conn_lost:
  223. # Replicating asyncio logic here.
  224. if self._conn_lost >= LOG_THRESHOLD_FOR_CONNLOST_WRITES:
  225. aio_logger.warning('socket.send() raised exception.')
  226. self._conn_lost += 1
  227. return
  228. self._send(data, addr)
  229. cdef void __uv_udp_on_receive(uv.uv_udp_t* handle,
  230. ssize_t nread,
  231. const uv.uv_buf_t* buf,
  232. const system.sockaddr* addr,
  233. unsigned flags) with gil:
  234. if __ensure_handle_data(<uv.uv_handle_t*>handle,
  235. "UDPTransport receive callback") == 0:
  236. return
  237. cdef:
  238. UDPTransport udp = <UDPTransport>handle.data
  239. Loop loop = udp._loop
  240. bytes data
  241. object pyaddr
  242. # It's OK to free the buffer early, since nothing will
  243. # be able to touch it until this method is done.
  244. __loop_free_buffer(loop)
  245. if udp._closed:
  246. # The handle was closed, there is no reason to
  247. # do any work now.
  248. udp.__receiving_stopped() # Just in case.
  249. return
  250. if addr is NULL and nread == 0:
  251. # From libuv docs:
  252. # addr: struct sockaddr* containing the address
  253. # of the sender. Can be NULL. Valid for the duration
  254. # of the callback only.
  255. # [...]
  256. # The receive callback will be called with
  257. # nread == 0 and addr == NULL when there is
  258. # nothing to read, and with nread == 0 and
  259. # addr != NULL when an empty UDP packet is
  260. # received.
  261. return
  262. if addr is NULL:
  263. pyaddr = None
  264. else:
  265. try:
  266. pyaddr = __convert_sockaddr_to_pyaddr(addr)
  267. except BaseException as exc:
  268. udp._error(exc, False)
  269. return
  270. if nread < 0:
  271. exc = convert_error(nread)
  272. udp._on_receive(None, exc, pyaddr)
  273. return
  274. if pyaddr is None:
  275. udp._fatal_error(
  276. RuntimeError(
  277. 'uv_udp.receive callback: addr is NULL and nread >= 0'),
  278. False)
  279. return
  280. if nread == 0:
  281. data = b''
  282. else:
  283. data = loop._recv_buffer[:nread]
  284. try:
  285. udp._on_receive(data, None, pyaddr)
  286. except BaseException as exc:
  287. udp._error(exc, False)
  288. cdef void __uv_udp_on_send(uv.uv_udp_send_t* req, int status) with gil:
  289. if req.data is NULL:
  290. # Shouldn't happen as:
  291. # - _UDPSendContext does an extra INCREF in its 'init()'
  292. # - _UDPSendContext holds a ref to the relevant UDPTransport
  293. aio_logger.error(
  294. 'UVStream.write callback called with NULL req.data, status=%r',
  295. status)
  296. return
  297. cdef:
  298. _UDPSendContext ctx = <_UDPSendContext> req.data
  299. UDPTransport udp = <UDPTransport>ctx.udp
  300. ctx.close()
  301. if status < 0:
  302. exc = convert_error(status)
  303. print(exc)
  304. else:
  305. exc = None
  306. try:
  307. udp._on_sent(exc)
  308. except BaseException as exc:
  309. udp._error(exc, False)