DelegatedConnectionListenerFactory.cs

usingSystem.Buffers;
usingSystem.Buffers.Binary;
usingSystem.Diagnostics;
usingSystem.IO.Pipelines;
usingSystem.Net;
usingSystem.Net.Sockets;
usingSystem.Threading.Channels;
usingMicrosoft.AspNetCore.Connections;
usingMicrosoft.AspNetCore.Server.Kestrel.Transport.Sockets;

// Custom transport that accepts client sockets over the listen socket as data
internalclassDelegatedConnectionListenerFactory:IConnectionListenerFactory
{
privateILogger<DelegatedConnectionListenerFactory>_logger;

publicDelegatedConnectionListenerFactory(ILogger<DelegatedConnectionListenerFactory>logger)
{
_logger=logger;
}

publicasyncValueTask<IConnectionListener>BindAsync(EndPointendpoint,CancellationTokencancellationToken=default)
{
// Front end connections connect to this socket
varsocket=newSocket(endpoint.AddressFamily,SocketType.Stream,ProtocolType.Tcp);
socket.Bind(endpoint);
socket.Listen();

returnnewSocketListener(socket,_logger);
}

classSocketListener:IConnectionListener
{
privateSocket_socket;
privatereadonlyTask_acceptTask;
// This is equivalent to our connection backlog
privatereadonlyChannel<ConnectionContext>_channel=Channel.CreateBounded<ConnectionContext>(20);
privatereadonlySocketConnectionContextFactory_contextFactory;

publicSocketListener(Socketsocket,ILoggerlogger)
{
_socket=socket;
_acceptTask=AcceptServerSocketConnections();
_contextFactory=new(new(),logger);
}

publicEndPointEndPoint=>_socket.LocalEndPoint!;

publicasyncValueTask<ConnectionContext?>AcceptAsync(CancellationTokencancellationToken=default)
{
varconnection=await_channel.Reader.ReadAsync(cancellationToken);

returnconnection;
}

publicasyncValueTaskDisposeAsync()
{
_socket.Dispose();

await_acceptTask;
}

publicValueTaskUnbindAsync(CancellationTokencancellationToken=default)
{
_socket.Dispose();

returndefault;
}

privateasyncTaskAcceptServerSocketConnections()
{
asyncTaskAcceptForwardedConnection(SocketfrontEndSocket)
{
varreader=PipeReader.Create(newNetworkStream(frontEndSocket));
while(true)
{
// Make sure we have the length at least
varresult=awaitreader.ReadAtLeastAsync(2);
varbuffer=result.Buffer;

if(buffer.Length<2)
{
Debug.Assert(result.IsCompleted);
break;
}

varlengthBytes=buffer.Slice(0,2);
varpayloadLength=BinaryPrimitives.ReadInt16LittleEndian(lengthBytes.IsSingleSegment?lengthBytes.FirstSpan:lengthBytes.ToArray());

// Look at the remaining buffer
buffer=buffer.Slice(2);

// We didn't get the payload in the same message so read it
if(buffer.Length<payloadLength)
{
// Advance 2 bytes
reader.AdvanceTo(lengthBytes.End);

// Try reading a new message the length of the payload
result=awaitreader.ReadAtLeastAsync(payloadLength);
buffer=result.Buffer;

// We still didn't get enough, it means we're done reading
if(buffer.Length<payloadLength)
{
Debug.Assert(result.IsCompleted);
break;
}
}

varmessage=buffer.Slice(0,payloadLength);

varsi=newSocketInformation
{
ProtocolInformation=message.ToArray()
};

reader.AdvanceTo(message.End);

varsocket=newSocket(si);

varconnectionContext=_contextFactory.Create(socket);

await_channel.Writer.WriteAsync(connectionContext);
}

awaitreader.CompleteAsync();
}

try
{
while(true)
{
// This loop waits for new front end connections to receive
// forwarded sockets from
varclient=await_socket.AcceptAsync();

// Kick off the listen for sockets passed over this channel
_=Task.Run(()=>AcceptForwardedConnection(client));
}
}
catch(ObjectDisposedException)
{

}
}
}
}