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.
 
 
 
 

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