ASP.NET close connection before sending all websocket data to client

Question

I wrote a simple asp.net websocket handler as a gateway between a remote data processing server and clients. I tested in my local machine (win8, IIS EXPRESS 8) and everything worked well. But in azure website, ASP.NET closes connection before sending all websocket data to client.

Following is my data transfer code:

internal class WebSocketStreamTransfer{

    public WebSocketStreamTransfer(CancellationToken disconnectionToken){
        DisconnectionToken = disconnectionToken;
    }

    private CancellationToken DisconnectionToken{
        get;
        set;
    }

    public async Task AcceptWebSocketConnection(WebSocketContext context) {
        if (context == null)
            throw new ArgumentNullException("context");
        WebSocket websocket = context.WebSocket;
        if (websocket == null)
            throw new SocksOverHttpException("Null websocket");
        using(IConnection conn = ConnectionManagerFactory.ConnectionManager.CreateConnection(Guid.NewGuid().ToString())) {
            try {
                DisconnectionToken.Register(conn.Close);
                TaskCompletionSource<bool> tcs = new TaskCompletionSource<bool>(null);
                await Task.WhenAny(SendDataToRemoteServer(conn, websocket, DisconnectionToken, tcs), SendDataToClient(conn, websocket, DisconnectionToken, tcs.Task));
            } catch(Exception e) {
                Logger.LogException(e);
            }
        }
    }

    internal static async Task SendDataToRemoteServer(IConnection conn, WebSocket websocket, CancellationToken cancelToken, TaskCompletionSource<bool> tcs) {
        try {
            ArraySegment<byte> buffer = new ArraySegment<byte>(new byte[ApplicationConfiguration.GetDefaultBufferSize()]);
            while (IsConnected(conn, cancelToken, websocket)) {
                WebSocketReceiveResult result = await websocket.ReceiveAsync(buffer, cancelToken);
                if (websocket.State == WebSocketState.Open) {
                    if (result.MessageType == WebSocketMessageType.Binary) {
                        if (result.Count > 0) {
                            if (IsConnected(conn, cancelToken, websocket)) {
                                int numRead = await conn.SendData(buffer.Array, 0, result.Count, cancelToken);
                                if (numRead > 0) {
                                    tcs.TrySetResult(true); // Notify SendDataToClient can continue
                                }else{
                                    Logger.LogError("Client not send enough data for remote connection built");
                                    return;
                                }
                            } else {
                                Logger.LogInformation("SendDataToRemoteServer: Cancel send data to remote server due to connection closed");
                            }
                        } else
                            Logger.LogInformation("Receive empty binary message");
                    } else if (result.MessageType == WebSocketMessageType.Text) {
                        Logger.LogError("Receive unexpected text message");
                        return;
                    } else {
                        Logger.LogInformation("Receive close message");
                        await websocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Close Connection", cancelToken);
                        return;
                    }
                } else {
                    Logger.LogInformation("SendDataToRemoteServer: WebSocket connection closed by client");
                    return;
                }
            }
        }finally{
            tcs.TrySetResult(true);
        }
    }

    internal static async Task SendDataToClient(IConnection conn, WebSocket websocket, CancellationToken cancelToken, Task connectedTask) {
        await connectedTask;
        while (IsConnected(conn, cancelToken, websocket)) {
            byte[] data = await conn.ReceiveData(cancelToken);
            if (data.Length <= 0) {
                Logger.LogInformation("SendDataToClient: Get empty data from remote server");
                return;
            }
            if (IsConnected(conn, cancelToken, websocket)) {
                await websocket.SendAsync(new ArraySegment<byte>(data), WebSocketMessageType.Binary, true, cancelToken);
            } else {
                Logger.LogInformation("SendDataToClient: Cancel send data to client due to connection closed");
            }
        }
    }

    internal static bool IsConnected(IConnection conn, CancellationToken cancelToken, WebSocket websocket) {
        bool socketConnected = websocket.State == WebSocketState.Open;
        return socketConnected && conn.Connected && !cancelToken.IsCancellationRequested;
    }
}

Problem scenario:

1. SendDataToRemoteServer waiting for client data and client has not data to send yet

2. SendDataToClient receive empty data from remote server, means remote server start closing connection. So finish SendDataToClient

3. AcceptWebSocketConnection finish because Task.WhenAny(SendDataToRemoteServer(conn, websocket, DisconnectionToken, tcs), SendDataToClient(conn, websocket, DisconnectionToken, tcs.Task))

4. Expect ASP.NET send all data before close tcp connection but ASP.NET close connection immediately (azure).

Answer 1

A WebSocket message can be split in different frames. You are not checking if the message is completed. The code for sending information form a WS connection to other should look like this:

WebSocketReceiveResult result = null;
do
{
    result = await source.ReceiveAsync(buffer, CancellationToken.None);
    var sendBuffer = new ArraySegment<Byte>(buffer.Array, buffer.Offset, result.Count);

    await target.SendAsync(sendBuffer, result.MessageType, result.EndOfMessage, CancellationToken.None);
}
while (!result.EndOfMessage);

You have to check the EndOfMessage property, and continue reading while the message is not completed.

It works in your local computer because locally you are no affected by the buffering in the same way, or because the messages you were trying were smaller.