1- use std:: { env, future:: Future , str:: FromStr , time:: Instant } ;
1+ use std:: {
2+ env,
3+ future:: Future ,
4+ str:: FromStr ,
5+ time:: { Duration , Instant } ,
6+ } ;
27
38use clap:: Parser ;
49use data_encoding:: HEXLOWER ;
510use iroh:: {
611 EndpointId , SecretKey ,
7- endpoint:: { Connecting , Connection } ,
12+ endpoint:: { Connecting , Connection , ZRTTConnection } ,
813} ;
914use n0_future:: { StreamExt , future} ;
1015use n0_snafu:: ResultExt ;
1116use n0_watcher:: Watcher ;
17+ use quinn:: VarInt ;
1218use tracing:: { info, trace} ;
1319
1420const PINGPONG_ALPN : & [ u8 ] = b"0rtt-pingpong" ;
@@ -49,7 +55,7 @@ pub fn get_or_generate_secret_key() -> n0_snafu::Result<SecretKey> {
4955/// read the response immediately. Otherwise, the stream pair is bad and we need
5056/// to open a new stream pair.
5157async fn pingpong (
52- connection : & Connection ,
58+ connection : & Conn ,
5359 proceed : impl Future < Output = bool > ,
5460 x : u64 ,
5561) -> n0_snafu:: Result < ( ) > {
@@ -73,16 +79,46 @@ async fn pingpong(
7379 Ok ( ( ) )
7480}
7581
76- async fn pingpong_0rtt ( connecting : Connecting , i : u64 ) -> n0_snafu:: Result < Connection > {
82+ enum Conn {
83+ ZRTT ( ZRTTConnection ) ,
84+ Full ( Connection ) ,
85+ }
86+
87+ impl Conn {
88+ fn open_bi ( & self ) -> quinn:: OpenBi < ' _ > {
89+ match self {
90+ Conn :: ZRTT ( conn) => conn. open_bi ( ) ,
91+ Conn :: Full ( conn) => conn. open_bi ( ) ,
92+ }
93+ }
94+
95+ fn close ( & self , error_code : VarInt , reason : & [ u8 ] ) {
96+ match self {
97+ Conn :: ZRTT ( conn) => conn. close ( error_code, reason) ,
98+ Conn :: Full ( conn) => conn. close ( error_code, reason) ,
99+ }
100+ }
101+
102+ fn rtt ( & self ) -> Duration {
103+ match self {
104+ Conn :: ZRTT ( conn) => conn. rtt ( ) ,
105+ Conn :: Full ( conn) => conn. rtt ( ) ,
106+ }
107+ }
108+ }
109+
110+ async fn pingpong_0rtt ( connecting : Connecting , i : u64 ) -> n0_snafu:: Result < Conn > {
77111 let connection = match connecting. into_0rtt ( ) {
78112 Ok ( ( connection, accepted) ) => {
79113 trace ! ( "0-RTT possible from our side" ) ;
114+ let connection = Conn :: ZRTT ( connection) ;
80115 pingpong ( & connection, accepted, i) . await ?;
81116 connection
82117 }
83118 Err ( connecting) => {
84119 trace ! ( "0-RTT not possible from our side" ) ;
85120 let connection = connecting. await . e ( ) ?;
121+ let connection = Conn :: Full ( connection) ;
86122 pingpong ( & connection, future:: ready ( true ) , i) . await ?;
87123 connection
88124 }
@@ -105,6 +141,7 @@ async fn connect(args: Args) -> n0_snafu::Result<()> {
105141 . await ?;
106142 let connection = if args. disable_0rtt {
107143 let connection = connecting. await . e ( ) ?;
144+ let connection = Conn :: Full ( connection) ;
108145 trace ! ( "connecting without 0-RTT" ) ;
109146 pingpong ( & connection, future:: ready ( true ) , i) . await ?;
110147 connection
0 commit comments