o
    ;ic                     @   s  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 d dlmZm	Z	 d dl
mZmZmZmZmZmZmZ d dl
mZmZ d dl mZmZmZmZmZ d dl mZ d dlmZ d d	lmZ d d
lmZ d dl m!Z! d dl"mZ# d dl"m$Z$m%Z%m&Z&m'Z' d dl"m(Z(m)Z)m*Z* d dl"m+Z+m,Z,m-Z- d dl"m.Z.m/Z/m0Z0 d dl1m2Z2 d dl3m4Z4m5Z5 d dl6m7Z7m8Z8m9Z9 ddl:m;Z; ddl<m=Z= ddl>m?Z? erd dl
m@Z@ eAeBZCeDedreDedrdndZEdejddfdd ZFn
dejddfd!d ZFG d"d# d#eZGG d$d% d%eZHG d&d' d'ZIG d(d) d)ZJG d*d+ d+ZKeeeLeLf  ZMG d,d- d-ZNG d.d/ d/eZOeePeNf ZQG d0d1 d1ZRG d2d3 d3eZSdS )4    N)BytesIO)ABCabstractmethod)OptionalListTupleDict
NamedTupleCallableAny)castTYPE_CHECKING)	TransportProtocolEventBaseTransportTimerHandle)Queue)partial)deque)
ErrorCodes)H2Configuration)r   )RequestReceivedDataReceivedStreamEndedWindowUpdated)ConnectionTerminatedRemoteSettingsChangedStreamReset)SettingsAcknowledgedResponseReceivedTrailersReceived)PriorityUpdatedPingReceivedPingAckReceived)SettingCodes)H2ConnectionConnectionState)ProtocolErrorTooManyStreamsErrorStreamClosedError   )Wrapper)Configuration)StreamTerminatedError)DequeTCP_NODELAYSOCK_NONBLOCK   l    sockreturnc                 C   sP   | j tjtjhv r"| jt@ tjkr$| jtjkr&| 	tjtj
d d S d S d S d S Nr+   )familysocketAF_INETAF_INET6type_sock_type_maskSOCK_STREAMprotoIPPROTO_TCP
setsockoptr0   r3    rA   D/home/ubuntu/.local/lib/python3.10/site-packages/grpclib/protocol.py_set_nodelay*   s   rC   c                 C      d S NrA   r@   rA   rA   rB   rC   2      c                   @   s&   e Zd ZU eed< eed< eed< dS )UnackedDatadata	data_sizeack_sizeN)__name__
__module____qualname__bytes__annotations__intrA   rA   rA   rB   rG   6   s   
 rG   c                   @   s   e Zd ZU eed< eed< dS )	AckedDatarH   rI   N)rK   rL   rM   
memoryviewrO   rP   rA   rA   rA   rB   rQ   <   s   
 rQ   c                   @   sh   e Zd Zdeegdf ddfddZdededdfdd	Zdd
dZdedefddZ	defddZ
dS )Bufferack_callbackNr4   c                 C   s&   || _ d| _t | _t | _d| _d S )NFr   )_ack_callback_eofr   _unackedr   _acked_acked_size)selfrT   rA   rA   rB   __init__C   s
   
zBuffer.__init__rH   rJ   c                 C   s   | j t|t|| d S rE   )rW   
put_nowaitrG   len)rZ   rH   rJ   rA   rA   rB   addJ   s   z
Buffer.addc                 C   s   | j tddd d| _d S )N    r   T)rW   r\   rG   rV   rZ   rA   rA   rB   eofM   s   
z
Buffer.eofsizec           
         sf  |dks	J d|dkrdS | j r| j sF| j|k rF| j I d H \}}}|s*n| jtt|| |  j|7  _| 	| | j|k s| j rP| jdkrPdS | j|k rYt
dg }d}||k r| jd \}}|| |kr}|| ||7 }| j  n || }	||d |	  ||	7 }t||	d  ||	 d| jd< ||k sa|  j|8  _||ksJ d|S )Nr   zSize can not be negativer_   z Received less data than expected)rH   rI   )rV   rW   emptyrY   getrX   appendrQ   rR   rU   AssertionErrorpopleftjoin)
rZ   rb   rH   rI   rJ   chunkschunks_size
next_chunknext_chunk_sizeoffsetrA   rA   rB   readQ   sH   






zBuffer.readc                    s    t  fddt j D S )Nc                 3   s    | ]	} j  jV  qd S rE   )rW   
get_nowaitrJ   ).0_r`   rA   rB   	<genexpr>z   s    z&Buffer.unacked_size.<locals>.<genexpr>)sumrangerW   qsizer`   rA   r`   rB   unacked_sizey   s   zBuffer.unacked_sizer4   N)rK   rL   rM   r
   rP   r[   rN   r^   ra   rn   rv   rA   rA   rA   rB   rS   A   s    
(rS   c                   @   sV   e Zd ZdZdeddfddZdeeee	f  fddZ
deeeef  fd	d
ZdS )Peerz=
    Represents an information about a connection's peer
    	transportr4   Nc                 C   s
   || _ d S rE   )
_transport)rZ   ry   rA   rA   rB   r[         
zPeer.__init__c                 C   s   | j dS )z4Returns the remote address to which we are connectedpeername)rz   get_extra_infor`   rA   rA   rB   addr   s   z	Peer.addrc                 C   s    | j d}|dur| S dS )zhReturns the peer certificate

        Result of the :py:meth:`python:ssl.SSLSocket.getpeercert`
        
ssl_objectN)rz   r}   getpeercert)rZ   r   rA   rA   rB   cert   s   z	Peer.cert)rK   rL   rM   __doc__r   r[   r   r   strrP   r~   r   r   r   rA   rA   rA   rB   rx   ~   s
    rx   c                   @   s  e Zd ZU dZdZdZdZdZdZdZ	dZ
dZee ed< dZee ed< dZee ed< dZee ed< dZee ed< dZee ed	< dZeed
< dZee ed< dZee ed< dedededdfddZdedee fddZ dededdfddZ!d7ddZ"d7ddZ#ddddee dee$ dd fd!d"Z%d7d#d$Z&d7d%d&Z'de(fd'd(Z)de*fd)d*Z+d7d+d,Z,de*fd-d.Z-d7d/d0Z.d7d1d2Z/d7d3d4Z0d7d5d6Z1dS )8
Connectionzh
    Holds connection state (write_ready), and manages
    H2Connection <-> Transport communication
    r   Nlast_stream_createdlast_data_sentlast_data_receivedlast_message_sentlast_message_receivedlast_ping_sentping_count_in_sequence_ping_handle_close_by_ping_handler
connectionry   configr4   c                C   s0   || _ || _|| _t | _| j  t | _d S rE   )_connectionrz   _configr   write_readysetstream_close_waiter)rZ   r   ry   r   rA   rA   rB   r[      s   
zConnection.__init__rH   c                 C   s   | j |S rE   )r   receive_datarZ   rH   rA   rA   rB   feed   s   zConnection.feed	stream_idrb   c                 C   s"   |r| j || |   d S d S rE   )r   acknowledge_received_dataflush)rZ   r   rb   rA   rA   rB   ack   s   zConnection.ackc                 C      | j   d S rE   )r   clearr`   rA   rA   rB   pause_writing      zConnection.pause_writingc                 C   r   rE   )r   r   r`   rA   rA   rB   resume_writing   r   zConnection.resume_writingr   wrapperr   Streamc                C   s   t | | j| j||dS )Nr   )r   r   rz   )rZ   r   r   rA   rA   rB   create_stream   s   zConnection.create_streamc                 C   s"   | j  }|r| j| d S d S rE   )r   data_to_sendrz   writer   rA   rA   rB   r      s   
zConnection.flushc                 C   s,   | j jd urt | j j| j| _d S d S rE   )r   _keepalive_timeasyncioget_event_loop
call_later_pingr   r`   rA   rA   rB   
initialize   s   
zConnection.initializec                 C   s
   t | jS rE   )rx   rz   r`   rA   rA   rB   get_peer   r{   zConnection.get_peerc                 C   s   t | dr
| j S dS )Nrz   T)hasattrrz   
is_closingr`   rA   rA   rB   r      s   

zConnection.is_closingc                 C   sZ   t | dr| j  | `t | jdr| j`| jd ur| j  | jd ur+| j  d S d S )Nrz   _frame_dispatch_table)r   rz   closer   r   r   cancelr   r`   rA   rA   rB   r      s   




zConnection.closec                 C   s|   | j jd usJ | j jstdd | jj D sdS | j jdkr*| j| j jkr*dS | j	d ur<t
 | j	 | j jk r<dS dS )Nc                 s   s    | ]}|j V  qd S rE   )open)rp   srA   rA   rB   rr      s    z0Connection._is_need_send_ping.<locals>.<genexpr>Fr   T)r   r   _keepalive_permit_without_callsanyr   streamsvalues_http2_max_pings_without_datar   r   time	monotonic*_http2_min_sent_ping_interval_without_datar`   rA   rA   rB   _is_need_send_ping   s   
zConnection._is_need_send_pingc                 C   s   | j jd usJ |  rDtd tdtt	 d }| j
| |   t	 | _|  jd7  _| jd u rDt | j j| j| _t | j j| j| _d S )Nz	send pingz!Qi@B r+   )r   r   r   logdebugstructpackrP   r   r   r   pingr   r   r   r   r   r   r   _keepalive_timeoutr   r   r   r   rA   rA   rB   r     s&   



zConnection._pingc                 C   s
   d| _ d S Nr   )r   r`   rA   rA   rB   headers_send_process  r{   zConnection.headers_send_processc                 C   s   d| _ t | _d S r   )r   r   r   r   r`   rA   rA   rB   data_send_process  s   zConnection.data_send_processc                 C   s"   | j d ur| j   d | _ d S d S rE   )r   r   r`   rA   rA   rB   ping_ack_process"  s   


zConnection.ping_ack_processrw   )2rK   rL   rM   r   streams_startedstreams_succeededstreams_failed	data_sentdata_receivedmessages_sentmessages_receivedr   r   floatrO   r   r   r   r   r   r   rP   r   r   r   r&   r   r-   r[   rN   r   H2Eventr   r   r   r   r,   r   r   r   rx   r   boolr   r   r   r   r   r   r   rA   rA   rA   rB   r      sd   
 




	




r   c                   @   sp  e Zd ZU dZdZee ed< dZee	 ed< dZ
dZddddeded	ed
ee dee ddfddZd
ededdfddZdefddZdedefddZdefddZ	d2dededddeg df fddZ	d2dededdfdd Zd2d!ededdfd"d#Zd3d$d%Zejfd&eddfd'd(Zejfd&eddfd)d*Z d3d+d,Z!d-e"ddfd.d/Z#e$defd0d1Z%dS )4r   zL
    API for working with streams, used by clients and request handlers
    Nidcreatedr   r   r   h2_connectionry   r   r   r4   c                C   sV   || _ || _|| _|| _|d ur| || j  t | _d | _t | _d | _	t | _
d S rE   )r   _h2_connectionrz   r   init_streamr   window_updatedheadersheaders_receivedtrailerstrailers_received)rZ   r   r   ry   r   r   rA   rA   rB   r[   6  s   	zStream.__init__c                 C   s@   || _ tt|j| j | _| j jd7  _t  | _	| j_
d S r5   )r   rS   r   r   bufferr   r   r   r   r   r   )rZ   r   r   rA   rA   rB   r   M  s   zStream.init_streamc                    0   | j d u r| j I d H  | j d usJ | j S rE   )r   r   waitr`   rA   rA   rB   recv_headersT  
   
zStream.recv_headersrb   c                    s   | j |I d H S rE   )r   rn   )rZ   rb   rA   rA   rB   	recv_dataZ  s   zStream.recv_datac                    r   rE   )r   r   r   r`   rA   rA   rB   recv_trailers]  r   zStream.recv_trailersFr   
end_stream
_processorEventsProcessorc                   s   | j d u sJ | j 	 | jj I d H  | j }z| jj|||d W n ty=   | jj	  | jj I d H  Y qw | 
|| j || }| j| j  | j  |S )NTr   )r   r   r   r   r   get_next_available_stream_idsend_headersr)   r   r   r   registerrz   r   r   r   )rZ   r   r   r   r   release_streamrA   rA   rB   send_requestc  s&   




zStream.send_requestc                    sn   | j d usJ | jj I d H  | j | jjvrt| j | jj| j ||d | j	| j
  | j  d S )Nr   )r   r   r   r   r   r   r*   r   rz   r   r   r   )rZ   r   r   rA   rA   rB   r     s   
zStream.send_headersrH   c           
         s@  | j d usJ t|}dt|}}	 | jj I d H  | j| j }|dks6| j	  | j I d H  q| jj
}|t|||| }t|}	| }||krz| jj| j ||d | j| j  |  j|	7  _| j j|	7  _| j  d S | j| j | | j| j  |  j|	7  _| j j|	7  _| j  q)Nr   Tr   )r   r   r]   r   r   r   r   local_flow_control_windowr   r   max_outbound_frame_sizern   mintell	send_datarz   r   r   r   r   )
rZ   rH   r   ff_posf_lastwindowmax_frame_sizef_chunkf_chunk_lenrA   rA   rB   r     s<   


zStream.send_datac                    sF   | j d usJ | jj I d H  | j| j  | j| j  d S rE   )	r   r   r   r   r   r   rz   r   r   r`   rA   rA   rB   end  s
   z
Stream.end
error_codec                    sJ   | j d usJ | jj I d H  | jj| j |d | j| j  d S N)r   )	r   r   r   r   r   reset_streamrz   r   r   rZ   r   rA   rA   rB   reset  s
   zStream.resetc                 C   sF   | j d usJ | jj| j |d | jj r!| j| j  d S d S r   )	r   r   r   r   r   is_setrz   r   r   r   rA   rA   rB   reset_nowait  s
   zStream.reset_nowaitc                 C   r   rE   )r   ra   r`   rA   rA   rB   	__ended__  r   zStream.__ended__reasonc                 C   s"   | j d ur| j t| d S d S rE   )r   r   r.   )rZ   r  rA   rA   rB   __terminated__  s   
zStream.__terminated__c                 C   sT   | j d usJ | j rdS | jjjtju rdS | jj	| j }|d u r&dS |j
 S )NF)r   rz   r   r   state_machinestater'   CLOSEDr   rd   closedrZ   streamrA   rA   rB   closable  s   
zStream.closable)Frw   )&rK   rL   rM   r   r   r   rP   rO   r   r   r   r   r   r&   r   r,   r[   r   _Headersr   rN   r   r   r   r
   r   r   r   r   r   NO_ERRORr   r   r   r   r  propertyr	  rA   rA   rA   rB   r   +  sp   
 
	

-

%

	r   c                	   @   sV   e Zd Zedededeg df ddfddZededdfdd	Zedd
dZ	dS )AbstractHandlerr  r   r   Nr4   c                 C   rD   rE   rA   )rZ   r  r   r   rA   rA   rB   accept  s   zAbstractHandler.acceptc                 C   rD   rE   rA   r  rA   rA   rB   r        zAbstractHandler.cancelc                 C   rD   rE   rA   r`   rA   rA   rB   r     r  zAbstractHandler.closerw   )
rK   rL   rM   r   r   r
  r
   r  r   r   rA   rA   rA   rB   r    s    
r  c                   @   sP  e Zd ZdZdededdfddZdedeg df fd	d
Z	d,de
ddfddZdeddfddZdeddfddZdeddfddZdeddfddZdeddfddZdeddfddZdeddfddZdeddfddZdeddfd d!Zdeddfd"d#Zde ddfd$d%Z!de"ddfd&d'Z#de$ddfd(d)Z%de&ddfd*d+Z'dS )-r   zP
    H2 events processor, synchronous, not doing any IO, as hyper-h2 itself
    handlerr   r4   Nc                 C   sj   || _ || _t| jt| jt| jt| j	t
| jt| jt| jt| jt| jt| jt| jt| jt| ji| _i | _d S rE   )r  r   r   process_request_receivedr    process_response_receivedr   process_remote_settings_changedr   process_settings_acknowledgedr   process_data_receivedr   process_window_updatedr!   process_trailers_receivedr   process_stream_endedr   process_stream_resetr"   process_priority_updatedr   process_connection_terminatedr#   process_ping_receivedr$   process_ping_ack_received
processorsr   )rZ   r  r   rA   rA   rB   r[   	  s"   
zEventsProcessor.__init__r  c                    s>   j d usJ  jj <  jddtdd f fdd}|S )N)_streamsr  r4   c                    sN   j d usJ | j } jj   j s% jj |j  d S d S rE   )	r   popr   r   r   r   r   r   rv   )r  _streamr  rA   rB   r   '  s   
z0EventsProcessor.register.<locals>.release_stream)r   r   _Streams)rZ   r  r   rA   r  rB   r   #  s    zEventsProcessor.registerConnection closedr  c                 C   sD   | j   | j  | j D ]}|| qt| dr | `d S d S )Nr  )r   r   r  r   r   r  r   r  )rZ   r  r  rA   rA   rB   r   0  s   


zEventsProcessor.closeeventc                 C   sF   z| j |j }W n ty   t| ty   Y d S w || d S rE   )r  	__class__KeyErrorNotImplementedErrorAttributeError)rZ   r$  procrA   rA   rB   process9  s   zEventsProcessor.processc                 C   s0   | j j|jd}| |}| j||j| d S )N)r   )r   r   r   r   r  r  r   )rZ   r$  r  r   rA   rA   rB   r  C  s   
z(EventsProcessor.process_request_receivedc                 C   s0   | j |j}|d ur|j|_|j  d S d S rE   )r   rd   r   r   r   r   rZ   r$  r  rA   rA   rB   r  M  
   z)EventsProcessor.process_response_receivedc                 C   s.   t j|jv r| j D ]	}|j  qd S d S rE   )r%   INITIAL_WINDOW_SIZEchanged_settingsr   r   r   r   r+  rA   rA   rB   r  S  s
   z/EventsProcessor.process_remote_settings_changedc                 C   rD   rE   rA   rZ   r$  rA   rA   rB   r  [  s   z-EventsProcessor.process_settings_acknowledgedc                 C   st   t |j}| j|j}|d ur!|j|j|j | j|7  _n	| j	
|j|j | j	 j|7  _t | j	_d S rE   )r]   rH   r   rd   r   r   r^   flow_controlled_lengthr   r   r   r   r   r   )rZ   r$  rb   r  rA   rA   rB   r  a  s   
z%EventsProcessor.process_data_receivedc                 C   sP   |j dkr| j D ]}|j  q
d S | j|j }|d ur&|j  d S d S r   )r   r   r   r   r   rd   )rZ   r$  valuer  rA   rA   rB   r  r  s   
z&EventsProcessor.process_window_updatedc                 C   s0   | j |j}|d ur|j|_|j  d S d S rE   )r   rd   r   r   r   r   r   r+  rA   rA   rB   r  {  r,  z)EventsProcessor.process_trailers_receivedc                 C   s2   | j |j}|d ur|  | j jd7  _d S r5   )r   rd   r   r   r   r   r+  rA   rA   rB   r    s   z$EventsProcessor.process_stream_endedc                 C   sX   | j |j}|d ur"|jrd|j}nd}|| | j| | j	 j
d7  _
d S )Nz,Stream reset by remote party, error_code: {}Protocol errorr+   )r   rd   r   remote_resetformatr   r  r  r   r   r   )rZ   r$  r  msgrA   rA   rB   r    s   
z$EventsProcessor.process_stream_resetc                 C   rD   rE   rA   r/  rA   rA   rB   r    rF   z(EventsProcessor.process_priority_updatedc                 C   s   | j d|jd d S )Nz9Received GOAWAY frame, closing connection; error_code: {}r  )r   r4  r   r/  rA   rA   rB   r    s   
z-EventsProcessor.process_connection_terminatedc                 C   rD   rE   rA   r/  rA   rA   rB   r    rF   z%EventsProcessor.process_ping_receivedc                 C   r   rE   )r   r   r/  rA   rA   rB   r    r   z)EventsProcessor.process_ping_ack_received)r#  )(rK   rL   rM   r   r  r   r[   r   r
   r   r   r   r   r*  r   r  r    r  r   r  r   r  r   r  r   r  r!   r  r   r  r   r  r"   r  r   r  r#   r  r$   r  rA   rA   rA   rB   r     sJ    
	



	
	r   c                   @   s   e Zd ZU eed< eed< dedededdfdd	Z	d
e
ddfddZdeddfddZdddZdddZdee ddfddZdS )
H2Protocolr   	processorr  r   	h2_configr4   Nc                 C   s   || _ || _|| _d S rE   )r  r   r9  )rZ   r  r   r9  rA   rA   rB   r[     s   
zH2Protocol.__init__ry   c                 C   s   | d}|d urt| t| jd}|  |jj}| jj| }| jj	| }|r.|
| |r:|tj| jj	i t|tt|| jd| _| j  | j  t| j| j| _d S )Nr7   )r   )r}   rC   r&   r9  initiate_connectionlocal_settingsinitial_window_sizer   http2_connection_window_sizehttp2_stream_window_sizeincrement_flow_control_windowupdate_settingsr%   r-  r   r   r   r   r   r   r   r  r8  )rZ   ry   r3   h2_conninitial
conn_deltastream_deltarA   rA   rB   connection_made  s.   



zH2Protocol.connection_maderH   c                 C   sn   z| j |}W n ty   tjddd | jd Y d S w | j   |D ]}| j| q'| j   d S )Nr2  T)exc_info)	r   r   r(   r   r   r8  r   r   r*  )rZ   rH   eventsr$  rA   rA   rB   r     s   
zH2Protocol.data_receivedc                 C   r   rE   )r   r   r`   rA   rA   rB   r     r   zH2Protocol.pause_writingc                 C   r   rE   )r   r   r`   rA   rA   rB   r     r   zH2Protocol.resume_writingexcc                 C   s   | j jdd d S )NzConnection lostr6  )r8  r   )rZ   rH  rA   rA   rB   connection_lost  s   zH2Protocol.connection_lostrw   )rK   rL   rM   r   rO   r   r  r-   r   r[   r   rE  rN   r   r   r   r   BaseExceptionrI  rA   rA   rA   rB   r7    s"   
 



r7  )Tr   r   r   r7   loggingior   abcr   r   typingr   r   r   r   r	   r
   r   r   r   r   r   r   r   r   r   	functoolsr   collectionsr   	h2.errorsr   	h2.configr   	h2.eventsr   r   r   r   r   r   r   r   r   r    r!   r"   r#   r$   h2.settingsr%   h2.connectionr&   r'   h2.exceptionsr(   r)   r*   utilsr,   r   r-   
exceptionsr.   r/   	getLoggerrK   r   r   r;   rC   rG   rQ   rS   rx   r   r   r
  r   r  rP   r"  r   r7  rA   rA   rA   rB   <module>   s^    $

=  D $