U
    î¤ùd  ã                   @   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ƒr†e 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__© ú2/tmp/pip-unpacked-wheel-pk5slln3/dns/quic/_sync.pyr       s    
zSyncQuicStream.__init__c              	   C   sn   |   |¡}| j& | j |¡r,W 5 Q R £ d S || _W 5 Q R X | j | j |¡sXtjj	‚W 5 Q R X d| _q d S )Nr   )
Z_timeout_from_expirationr   Ú_bufferZhaveZ
_expectingr   ÚwaitÚdnsÚ	exceptionÚTimeout)r   ÚamountÚ
expirationÚtimeoutr   r   r   Úwait_for%   s    
zSyncQuicStream.wait_forNc              
   C   st   |   |¡}|  d|¡ | j t d| j d¡¡\}W 5 Q R X |  ||¡ | j | j |¡W  5 Q R £ S Q R X d S )Né   z!H)Z_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	   )Z_encapsulateÚ_connectionÚwriteZ
_stream_id)r   ÚdatagramÚis_endÚdatar   r   r   Úsend:   s    
zSyncQuicStream.sendc              	   C   s,   |   ||¡r(| j | j ¡  W 5 Q R X d S r	   )Z_common_add_inputr   Únotify)r   r+   r*   r   r   r   Ú
_add_input>   s    zSyncQuicStream._add_inputc              	   C   s   | j  |  ¡  W 5 Q R X d S r	   )r   Ú_close©r   r   r   r   ÚcloseC   s    zSyncQuicStream.closec                 C   s   | S r	   r   r0   r   r   r   Ú	__enter__G   s    zSyncQuicStream.__enter__c              	   C   s(   |   ¡  | j | j ¡  W 5 Q R X dS ©NF)r1   r   r-   )r   Úexc_typeÚexc_valÚexc_tbr   r   r   Ú__exit__J   s    zSyncQuicStream.__exit__)N)F)Ú__name__Ú
__module__Ú__qualname__r   r    r&   r,   r.   r1   r2   r7   Ú__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 | j¡ t ¡ \| _	| _
| j
 d¡ | j d¡ | jd k	rªz| j tj | j| j¡¡ W n  tk
r¨   | j ¡  ‚ Y nX t ¡ | _d | _t ¡ | _d S )Nr   F)r
   r   ÚsocketZ_afÚ
SOCK_DGRAMÚ_socketÚconnectÚ_peerÚ
socketpairÚ_send_wakeupÚ_receive_wakeupÚsetblockingÚ_sourceÚbindr   ZinetZlow_level_address_tupleÚ	Exceptionr1   r   ÚEventÚ_handshake_completeÚ_worker_threadr   r   )r   r   ÚaddressÚportÚsourceÚsource_portÚmanagerr   r   r   r   R   s"    
ÿ

zSyncQuicConnection.__init__c              	   C   sn   d}|dk rj|d7 }z| j  t¡}W n tk
r:   Y d S X | j  | j || jd t ¡ ¡ W 5 Q R X qd S )Nr   é
   é   )	r?   Úrecvr   ÚBlockingIOErrorr   r'   Zreceive_datagramrA   Útime)r   Úcountr)   r   r   r   Ú_reade   s    zSyncQuicConnection._readc                 C   s.   z| j  d¡ W q  tk
r&   Y d S X q d S )Né    )rD   rS   rT   r0   r   r   r   Ú_drain_wakeupp   s    z SyncQuicConnection._drain_wakeupc           	      C   sú   zÐtƒ }| | jtj| j	¡ | | j
tj| j¡ | jsÎ|  d¡\}}| |¡}|D ]\}}| ¡  qR| j   |  |¡ | j t ¡ ¡}W 5 Q R X |D ].\}}z| j |¡ W q” tk
rÀ   Y q”X q”|  ¡  q0W 5 | j  d| _W 5 Q R X | j ¡  X d S )NTF)r   Ú_donerJ   ÚsetÚ_selector_classÚregisterr?   Ú	selectorsÚ
EVENT_READrW   rD   rY   Z_get_timer_valuesÚselectr+   Z_handle_timerr'   Zdatagrams_to_sendrU   r,   rT   Ú_handle_events)	r   Úselr   ÚintervalÚitemsÚkeyÚ_Z	datagramsr)   r   r   r   Ú_workerw   s*    


zSyncQuicConnection._workerc              	   C   sÈ   | j  | j ¡ }W 5 Q R X |d kr(d S t|tjjjƒrn| j  | j 	|j
¡}W 5 Q R X |rÂ| |j|j¡ q t|tjjjƒrŠ| j ¡  q t|tjjjƒsªt|tjjjƒr | j  d| _W 5 Q R X q d S )NT)r   r'   Z
next_eventÚ
isinstanceÚaioquicZquicÚeventsZStreamDataReceivedÚ_streamsr$   r   r.   r+   Z
end_streamZHandshakeCompletedrJ   r[   ZConnectionTerminatedZStreamResetrZ   )r   ÚeventÚstreamr   r   r   ra   ‘   s&     ÿþz!SyncQuicConnection._handle_eventsFc              	   C   s2   | j  | j |||¡ W 5 Q R X | j d¡ d S )Nó   )r   r'   Zsend_stream_datarC   r,   )r   rm   r+   r*   r   r   r   r(   ¤   s    zSyncQuicConnection.writec                 C   s(   | j r
d S tj| jd| _| j ¡  d S )N)Útarget)Ú_closedr   ÚThreadrg   rK   Ústartr0   r   r   r   Úrun©   s    zSyncQuicConnection.runNc              	   C   sT   | j  |¡stjj‚| j0 | jr&t‚| j 	d¡}t
| |ƒ}|| j|< W 5 Q R X |S r3   )rJ   r   r   r   r   r   rZ   r   r'   Zget_next_available_stream_idr   rk   )r   r   r   rm   r   r   r   Úmake_stream¯   s    
zSyncQuicConnection.make_streamc              	      s"   | j  tƒ  |¡ W 5 Q R X d S r	   )r   r
   Úclose_stream)r   r   r   r   r   ru   º   s    zSyncQuicConnection.close_streamc              	   C   sj   | j P | jrW 5 Q R £ d S | j | jd | jd ¡ d| _| j ¡  | j d¡ W 5 Q R X | j	 
¡  d S )Nr   rR   Trn   )r   rp   Ú_managerÚclosedrA   r'   r1   rC   r,   rK   Újoinr0   r   r   r   r1   ¾   s    
zSyncQuicConnection.close)F)N)r8   r9   r:   r   rW   rY   rg   ra   r(   rs   rt   ru   r1   r;   r   r   r   r   r<   Q   s   

r<   c                       sL   e Zd Zdejdf‡ fdd„	Zd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   r<   r   r   r   )r   ÚconfÚverify_modeZserver_namer   r   r   r   Ê   s    zSyncQuicManager.__init__éU  r   c              
   C   sB   | j 2 |  ||||¡\}}|r(| ¡  |W  5 Q R £ S Q R X d S r	   )r   Ú_connectrs   )r   rL   rM   rN   rO   r   rr   r   r   r   r@   Î   s
    zSyncQuicManager.connectc              	      s$   | j  tƒ  ||¡ W 5 Q R X d S r	   )r   r
   rw   )r   rL   rM   r   r   r   rw   Õ   s    zSyncQuicManager.closedc                 C   s   | S r	   r   r0   r   r   r   r2   Ù   s    zSyncQuicManager.__enter__c                 C   s$   t | j ¡ ƒ}|D ]}| ¡  qdS r3   )ÚlistZ_connectionsÚvaluesr1   )r   r4   r5   r6   Úconnectionsr   r   r   r   r7   Ü   s    
zSyncQuicManager.__exit__)r|   Nr   )r8   r9   r:   ÚsslÚCERT_REQUIREDr   r@   rw   r2   r7   r;   r   r   r   r   ry   É   s
   
ry   )r^   r=   r   r"   r   rU   Zaioquic.quic.configurationri   Zaioquic.quic.connectionZaioquic.quic.eventsZdns.exceptionr   Zdns.inetZdns.quic._commonr   r   r   r   r   Úhasattrr   r\   ÚSelectSelectorr   r<   ry   r   r   r   r   Ú<module>   s"   	
2x