1+ #if FEATURE_TAP
2+ using System ;
3+ using System . Net ;
4+ using System . Net . Sockets ;
5+ using System . Runtime . CompilerServices ;
6+ using System . Threading ;
7+ using System . Threading . Tasks ;
8+
9+ namespace Renci . SshNet . Abstractions
10+ {
11+ // Async helpers based on https://devblogs.microsoft.com/pfxteam/awaiting-socket-operations/
12+
13+ internal static class SocketExtensions
14+ {
15+ sealed class SocketAsyncEventArgsAwaitable : SocketAsyncEventArgs , INotifyCompletion
16+ {
17+ private readonly static Action SENTINEL = ( ) => { } ;
18+
19+ private bool isCancelled ;
20+ private Action continuationAction ;
21+
22+ public SocketAsyncEventArgsAwaitable ( )
23+ {
24+ Completed += delegate { SetCompleted ( ) ; } ;
25+ }
26+
27+ public SocketAsyncEventArgsAwaitable ExecuteAsync ( Func < SocketAsyncEventArgs , bool > func )
28+ {
29+ if ( ! func ( this ) )
30+ {
31+ SetCompleted ( ) ;
32+ }
33+ return this ;
34+ }
35+
36+ public void SetCompleted ( )
37+ {
38+ IsCompleted = true ;
39+ var continuation = continuationAction ?? Interlocked . CompareExchange ( ref continuationAction , SENTINEL , null ) ;
40+ if ( continuation != null )
41+ {
42+ continuation ( ) ;
43+ }
44+ }
45+
46+ public void SetCancelled ( )
47+ {
48+ isCancelled = true ;
49+ SetCompleted ( ) ;
50+ }
51+
52+ public SocketAsyncEventArgsAwaitable GetAwaiter ( ) { return this ; }
53+
54+ public bool IsCompleted { get ; private set ; }
55+
56+ void INotifyCompletion . OnCompleted ( Action continuation )
57+ {
58+ if ( continuationAction == SENTINEL || Interlocked . CompareExchange ( ref continuationAction , continuation , null ) == SENTINEL )
59+ {
60+ // We have already completed; run continuation asynchronously
61+ Task . Run ( continuation ) ;
62+ }
63+ }
64+
65+ public void GetResult ( )
66+ {
67+ if ( isCancelled )
68+ {
69+ throw new TaskCanceledException ( ) ;
70+ }
71+ else if ( IsCompleted )
72+ {
73+ if ( SocketError != SocketError . Success )
74+ {
75+ throw new SocketException ( ( int ) SocketError ) ;
76+ }
77+ }
78+ else
79+ {
80+ // We don't support sync/async
81+ throw new InvalidOperationException ( "The asynchronous operation has not yet completed." ) ;
82+ }
83+ }
84+ }
85+
86+ public static async Task ConnectAsync ( this Socket socket , IPEndPoint remoteEndpoint , CancellationToken cancellationToken )
87+ {
88+ cancellationToken . ThrowIfCancellationRequested ( ) ;
89+
90+ using ( var args = new SocketAsyncEventArgsAwaitable ( ) )
91+ {
92+ args . RemoteEndPoint = remoteEndpoint ;
93+
94+ using ( cancellationToken . Register ( o => ( ( SocketAsyncEventArgsAwaitable ) o ) . SetCancelled ( ) , args , false ) )
95+ {
96+ await args . ExecuteAsync ( socket . ConnectAsync ) ;
97+ }
98+ }
99+ }
100+
101+ public static async Task < int > ReceiveAsync ( this Socket socket , byte [ ] buffer , int offset , int length , CancellationToken cancellationToken )
102+ {
103+ cancellationToken . ThrowIfCancellationRequested ( ) ;
104+
105+ using ( var args = new SocketAsyncEventArgsAwaitable ( ) )
106+ {
107+ args . SetBuffer ( buffer , offset , length ) ;
108+
109+ using ( cancellationToken . Register ( o => ( ( SocketAsyncEventArgsAwaitable ) o ) . SetCancelled ( ) , args , false ) )
110+ {
111+ await args . ExecuteAsync ( socket . ReceiveAsync ) ;
112+ }
113+
114+ return args . BytesTransferred ;
115+ }
116+ }
117+ }
118+ }
119+ #endif
0 commit comments