o
    Rf                     @   s   d dl Z d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dl	Zd dl
Zd dlZd dlmZmZmZmZmZ ee drCe jZne jZG dd deZG dd deZG dd	 d	eZdS )
    N)QUIC_MAX_DATAGRAMBaseQuicConnectionBaseQuicManagerBaseQuicStreamUnexpectedEOFPollSelectorc                       sX   e Zd Z fddZdd ZdddZdd	d
Zdd Zdd Zdd Z	dd Z
  ZS )SyncQuicStreamc                    s&   t  || t | _t | _d S N)super__init__	threading	Condition_wake_upLock_lock)self
connection	stream_id	__class__ G/home/ubuntu/webapp/venv/lib/python3.10/site-packages/dns/quic/_sync.pyr       s   
zSyncQuicStream.__init__c                 C   s   	 |  |}| j | j|r	 W d    d S || _W d    n1 s&w   Y  | j | j|s9tjj	W d    n1 sCw   Y  d| _q)NTr   )
_timeout_from_expirationr   _bufferhave
_expectingr   waitdns	exceptionTimeout)r   amount
expirationtimeoutr   r   r   wait_for%   s   
zSyncQuicStream.wait_forNc                 C   s   |  |}| d| | j td| jd\}W d    n1 s$w   Y  | || | j | j|W  d    S 1 sCw   Y  d S )N   z!H)_expiration_from_timeoutr#   r   structunpackr   get)r   r"   r!   sizer   r   r   receive1   s   

$zSyncQuicStream.receiveFc                 C   s    |  |}| j| j|| d S r	   )_encapsulate_connectionwrite
_stream_id)r   datagramis_enddatar   r   r   send:   s   
zSyncQuicStream.sendc                 C   sF   |  ||r!| j | j  W d    d S 1 sw   Y  d S d S r	   )_common_add_inputr   notify)r   r1   r0   r   r   r   
_add_input>   s
   "zSyncQuicStream._add_inputc                 C   s4   | j  |   W d    d S 1 sw   Y  d S r	   )r   _closer   r   r   r   closeC   s   
"zSyncQuicStream.closec                 C      | S r	   r   r7   r   r   r   	__enter__G      zSyncQuicStream.__enter__c                 C   s>   |    | j | j  W d    dS 1 sw   Y  dS NF)r8   r   r4   )r   exc_typeexc_valexc_tbr   r   r   __exit__J   s   
zSyncQuicStream.__exit__r	   F)__name__
__module____qualname__r   r#   r*   r2   r5   r8   r:   r@   __classcell__r   r   r   r   r      s    

	r   c                       sl   e Zd Z fddZdd Zdd Zdd Zd	d
 ZdddZdd Z	dddZ
 fddZdd Z  ZS )SyncQuicConnectionc                    s   t  |||||| t| jtjd| _| jd ur7z| jtj	
| j| j W n ty6   | j   w | j| j t \| _| _| jd | jd t | _d | _t | _d S )Nr   F)r
   r   socket_af
SOCK_DGRAM_socket_sourcebindr   inetlow_level_address_tuple	Exceptionr8   connect_peer
socketpair_send_wakeup_receive_wakeupsetblockingr   Event_handshake_complete_worker_threadr   r   )r   r   addressportsourcesource_portmanagerr   r   r   r   R   s$   


zSyncQuicConnection.__init__c                 C   s   d}|dk rA|d7 }z| j t}W n
 ty   Y d S w | j | j|| jt  W d    n1 s6w   Y  |dk sd S d S )Nr   
      )	rJ   recvr   BlockingIOErrorr   r,   receive_datagramrQ   time)r   countr/   r   r   r   _reade   s   zSyncQuicConnection._readc                 C   s*   	 z| j d W n
 ty   Y d S w q)NT    )rT   r`   ra   r7   r   r   r   _drain_wakeupp   s   z SyncQuicConnection._drain_wakeupc           	      C   s|  zt  }|| jtj| j || jtj| j | js| 	d\}}|
|}|D ]\}}|  q)| j | | W d    n1 sEw   Y  |   | j | jt }W d    n1 sdw   Y  |D ]\}}z| j| W qk ty   Y qkw | jrW | j d| _W d    n1 sw   Y  | j  d S | j d| _W d    n1 sw   Y  | j  w )NFT)_selector_classregisterrJ   	selectors
EVENT_READre   rT   rg   _done_get_timer_valuesselectr1   r   _handle_timer_handle_eventsr,   datagrams_to_sendrc   r2   ra   rW   set)	r   selr!   intervalitemskey_	datagramsr/   r   r   r   _workerw   s@   

zSyncQuicConnection._workerc                 C   sH  	 | j  | j }W d    n1 sw   Y  |d u rd S t|tjjjrL| j  | j	|j
}W d    n1 s<w   Y  |rK||j|j nWt|tjjjrZ| j  nIt|tjjjry| j  d| _W d    n1 ssw   Y  n*t|tjjjr| j  | j	|j
}W d    n1 sw   Y  |r|dd q)NT    )r   r,   
next_event
isinstanceaioquicquiceventsStreamDataReceived_streamsr(   r   r5   r1   
end_streamHandshakeCompletedrW   rr   ConnectionTerminatedrl   StreamReset)r   eventstreamr   r   r   rp      s6   z!SyncQuicConnection._handle_eventsFc                 C   sF   | j  | j||| W d    n1 sw   Y  | jd d S )N   )r   r,   send_stream_datarS   r2   )r   r   r1   r0   r   r   r   r-      s   zSyncQuicConnection.writec                 C   s(   | j rd S tj| jd| _| j  d S )N)target)_closedr   Threadry   rX   startr7   r   r   r   run   s   zSyncQuicConnection.runNc                 C   sj   | j |s
tjj| j | jrt| j	d}t
| |}|| j|< W d    |S 1 s.w   Y  |S r<   )rW   r   r   r   r   r   rl   r   r,   get_next_available_stream_idr   r   )r   r"   r   r   r   r   r   make_stream   s   

zSyncQuicConnection.make_streamc                    s8   | j  t | W d    d S 1 sw   Y  d S r	   )r   r
   close_stream)r   r   r   r   r   r      s   "zSyncQuicConnection.close_streamc                 C   s   | j / | jr	 W d    d S | j| jd | jd  d| _| j  | jd W d    n1 s5w   Y  | j	
  d S )Nr   r_   Tr   )r   r   _managerclosedrQ   r,   r8   rS   r2   rX   joinr7   r   r   r   r8      s   
zSyncQuicConnection.closerA   r	   )rB   rC   rD   r   re   rg   ry   rp   r-   r   r   r   r8   rE   r   r   r   r   rF   Q   s    

rF   c                       sZ   e Zd Zdejdf fdd	Z	dddZ fd	d
Z fddZdd Z	dd Z
  ZS )SyncQuicManagerNc                    s    t  ||t| t | _d S r	   )r
   r   rF   r   r   r   )r   confverify_modeserver_namer   r   r   r      s   zSyncQuicManager.__init__U  r   Tc                 C   sP   | j  | |||||\}}|r|  |W  d    S 1 s!w   Y  d S r	   )r   _connectr   )r   rY   rZ   r[   r\   want_session_ticketr   r   r   r   r   rP      s   
$zSyncQuicManager.connectc                    s:   | j  t || W d    d S 1 sw   Y  d S r	   )r   r
   r   )r   rY   rZ   r   r   r   r      s   "zSyncQuicManager.closedc                    s<   | j  t ||| W d    d S 1 sw   Y  d S r	   )r   r
   save_session_ticket)r   rY   rZ   ticketr   r   r   r      s   "z#SyncQuicManager.save_session_ticketc                 C   r9   r	   r   r7   r   r   r   r:      r;   zSyncQuicManager.__enter__c                 C   s$   t | j }|D ]}|  q	dS r<   )list_connectionsvaluesr8   )r   r=   r>   r?   connectionsr   r   r   r   r@      s   
zSyncQuicManager.__exit__)r   Nr   T)rB   rC   rD   sslCERT_REQUIREDr   rP   r   r   r:   r@   rE   r   r   r   r   r      s    
r   )rj   rG   r   r&   r   rc   aioquic.quic.configurationr}   aioquic.quic.connectionaioquic.quic.eventsdns.exceptionr   dns.inetdns.quic._commonr   r   r   r   r   hasattrr   rh   SelectSelectorr   rF   r   r   r   r   r   <module>   s$   
	2|