/*++

    Copyright (c) Microsoft Corporation.
    Licensed under the MIT License.

Abstract:

    MsQuic Stream Wrapper

--*/

#include "precomp.h"
#ifdef QUIC_CLOG
#include "TestStream.cpp.clog.h"
#endif

TestStream::TestStream(
    _In_ HQUIC Handle,
    _In_opt_ STREAM_SHUTDOWN_CALLBACK_HANDLER StreamShutdownHandler,
    _In_ bool IsUnidirectional,
    _In_ bool IsPingSource
    ) :
    QuicStream(Handle),
    IsUnidirectional(IsUnidirectional), IsPingSource(IsPingSource), UsedZeroRtt(false),
    AllDataSent(IsUnidirectional && !IsPingSource), AllDataReceived(IsUnidirectional && IsPingSource),
    SendShutdown(IsUnidirectional && !IsPingSource), RecvShutdown(IsUnidirectional && IsPingSource),
    IsShutdown(false), ConnectionShutdown(false), ConnectionShutdownByApp(false), ConnectionClosedRemotely(false), ConnectionErrorCode(0), ConnectionCloseStatus(0),
    BytesToSend(0), OutstandingSendRequestCount(0), BytesReceived(0),
    StreamShutdownCallback(StreamShutdownHandler), Context(nullptr)
{
    CxPlatEventInitialize(&EventSendShutdownComplete, TRUE, (IsUnidirectional && !IsPingSource) ? TRUE : FALSE);
    CxPlatEventInitialize(&EventRecvShutdownComplete, TRUE, (IsUnidirectional && IsPingSource) ? TRUE : FALSE);
    if (QuicStream == nullptr) {
        TEST_FAILURE("Invalid handle passed into TestStream.");
    }
}

TestStream*
TestStream::FromStreamHandle(
    _In_ HQUIC QuicStreamHandle,
    _In_opt_ STREAM_SHUTDOWN_CALLBACK_HANDLER StreamShutdownHandler,
    _In_ QUIC_STREAM_OPEN_FLAGS Flags
    )
{
    auto IsUnidirectionalStream = !!(Flags & QUIC_STREAM_OPEN_FLAG_UNIDIRECTIONAL);
    auto Stream = new(std::nothrow) TestStream(QuicStreamHandle, StreamShutdownHandler, IsUnidirectionalStream, false);
    if (Stream == nullptr || !Stream->IsValid()) {
        TEST_FAILURE("Failed to create new TestStream.");
        delete Stream;
        return nullptr;
    }
    MsQuic->SetCallbackHandler(QuicStreamHandle, reinterpret_cast<void*>(QuicStreamHandler), Stream);
    return Stream;
}

TestStream*
TestStream::FromConnectionHandle(
    _In_ HQUIC QuicConnectionHandle,
    _In_opt_ STREAM_SHUTDOWN_CALLBACK_HANDLER StreamShutdownHandler,
    _In_ QUIC_STREAM_OPEN_FLAGS Flags
    )
{
    auto IsUnidirectionalStream = !!(Flags & QUIC_STREAM_OPEN_FLAG_UNIDIRECTIONAL);
    HQUIC QuicStreamHandle;
    QUIC_STATUS Status =
        MsQuic->StreamOpen(
            QuicConnectionHandle,
            Flags,
            QuicStreamHandler,
            nullptr,
            &QuicStreamHandle);
    if (QUIC_FAILED(Status)) {
        TEST_FAILURE("MsQuic->StreamOpen failed, 0x%x.", Status);
        return nullptr;
    }
    auto Stream = new(std::nothrow) TestStream(QuicStreamHandle, StreamShutdownHandler, IsUnidirectionalStream, true);
    if (Stream == nullptr || !Stream->IsValid()) {
        TEST_FAILURE("Failed to create new TestStream.");
        if (Stream == nullptr) {
            MsQuic->StreamClose(QuicStreamHandle);
        }
        delete Stream;
        return nullptr;
    }
    MsQuic->SetContext(QuicStreamHandle, Stream);
    return Stream;
}

TestStream::~TestStream()
{
    MsQuic->StreamClose(QuicStream);
    CxPlatEventUninitialize(EventRecvShutdownComplete);
    CxPlatEventUninitialize(EventSendShutdownComplete);
}

QUIC_STATUS
TestStream::Shutdown(
    _In_ QUIC_STREAM_SHUTDOWN_FLAGS Flags,
    _In_ QUIC_UINT62 ErrorCode // Application defined error code
    )
{
    return
        MsQuic->StreamShutdown(
            QuicStream,
            Flags,
            ErrorCode);
}

QUIC_STATUS
TestStream::Start(
    _In_ QUIC_STREAM_START_FLAGS Flags
    )
{
    return
        MsQuic->StreamStart(
            QuicStream,
            Flags);
}

bool
TestStream::StartPing(
    _In_ uint64_t PayloadLength,
    _In_ bool SendFin
    )
{
    BytesToSend = PayloadLength / MaxSendBuffers;

    do {
        auto SendBufferLength = (uint32_t)CXPLAT_MIN(BytesToSend, (int64_t)MaxSendLength);
        auto SendBuffer = new(std::nothrow) QuicSendBuffer(MaxSendBuffers, SendBufferLength);
        if (SendBuffer == nullptr) {
            TEST_FAILURE("Failed to alloc QuicSendBuffer");
            return false;
        }

        auto resultingBytesLeft = InterlockedSubtract64(&BytesToSend, SendBufferLength);

        QUIC_SEND_FLAGS Flags = QUIC_SEND_FLAG_ALLOW_0_RTT;
        if (PayloadLength == 0) {
            Flags |= QUIC_SEND_FLAG_START;
        }
        if (SendFin && resultingBytesLeft == 0) {
            Flags |= QUIC_SEND_FLAG_FIN;
        }

        InterlockedIncrement(&OutstandingSendRequestCount);
#ifdef CXPLAT_RAISE_IRQL
        CXPLAT_RAISE_IRQL();
#endif
        QUIC_STATUS Status =
            MsQuic->StreamSend(
                QuicStream,
                SendBuffer->Buffers,
                SendBuffer->BufferCount,
                Flags,
                SendBuffer);
#ifdef CXPLAT_RAISE_IRQL
        CXPLAT_LOWER_IRQL();
#endif
        if (QUIC_FAILED(Status)) {
            InterlockedDecrement(&OutstandingSendRequestCount);
            delete SendBuffer;
            TEST_FAILURE("MsQuic->StreamSend failed, 0x%x.", Status);
            return false;
        }
        if (resultingBytesLeft == 0) {
            // On the finish packet if it succeeds, the instance
            // we are executing in will be deleted. Return
            // so we don't execute the while on a deleted instance.
            return true;
        }

    } while (BytesToSend != 0 && OutstandingSendRequestCount < (int64_t)MaxSendRequestQueue);

    return true;
}

bool
TestStream::WaitForSendShutdownComplete()
{
    if (!CxPlatEventWaitWithTimeout(EventSendShutdownComplete, TestWaitTimeout)) {
        TEST_FAILURE("WaitForSendShutdownComplete timed out after %u ms.", TestWaitTimeout);
        return false;
    }
    return true;
}

bool
TestStream::WaitForRecvShutdownComplete()
{
    if (!CxPlatEventWaitWithTimeout(EventRecvShutdownComplete, TestWaitTimeout)) {
        TEST_FAILURE("WaitForRecvShutdownComplete timed out after %u ms.", TestWaitTimeout);
        return false;
    }
    return true;
}

//
// Stream Parameters
//

uint64_t
TestStream::GetStreamID()
{
    uint64_t value;
    uint32_t valueSize = sizeof(value);
    QUIC_STATUS Status =
        MsQuic->GetParam(
            QuicStream,
            QUIC_PARAM_STREAM_ID,
            &valueSize,
            &value);
    if (QUIC_FAILED(Status)) {
        value = 0;
        TEST_FAILURE("MsQuic->GetParam(QUIC_PARAM_STREAM_ID) failed, 0x%x.", Status);
    }
    return value;
}

QUIC_STATUS
TestStream::SetReceiveEnabled(
    bool value
    )
{
    return MsQuic->StreamReceiveSetEnabled(QuicStream, value ? TRUE : FALSE);
}

void
TestStream::HandleStreamRecv(
    _In_reads_(Length)
        const uint8_t * Buffer,
    _In_ uint32_t Length,
    _In_ QUIC_RECEIVE_FLAGS Flags
    )
{
    if (Buffer == nullptr) {
        TEST_FAILURE("Null Buffer");
        return;
    }
    if (Length == 0) {
        TEST_FAILURE("Zero Length Buffer");
        return;
    }

    BytesReceived += Length;

    if (!IsPingSource) {
        if (!!(Flags & QUIC_RECEIVE_FLAG_0_RTT)) {
            UsedZeroRtt = true;
        }

        if (!IsUnidirectional) {
            auto SendBuffer = new(std::nothrow) QuicSendBuffer(Length, Buffer);

#ifdef CXPLAT_RAISE_IRQL
            CXPLAT_RAISE_IRQL();
#endif
            QUIC_STATUS Status =
                MsQuic->StreamSend(
                    QuicStream,
                    SendBuffer->Buffers,
                    SendBuffer->BufferCount,
                    QUIC_SEND_FLAG_NONE,
                    SendBuffer);
#ifdef CXPLAT_RAISE_IRQL
            CXPLAT_LOWER_IRQL();
#endif

            if (QUIC_FAILED(Status)) {
                delete SendBuffer;
            }
            if (!SendShutdown) {
                if (QUIC_FAILED(Status)) {
                    TEST_FAILURE("MsQuic->StreamSend failed, 0x%x.", Status);
                }
            }
        }
    }
}

void
TestStream::HandleStreamSendComplete(
    _In_ bool Canceled,
    _In_ QuicSendBuffer* SendBuffer
    )
{
    if (IsPingSource) {
        if (BytesToSend == 0 || Canceled) {
            InterlockedDecrement(&OutstandingSendRequestCount);
            delete SendBuffer;
        } else {
            QUIC_SEND_FLAGS Flags = QUIC_SEND_FLAG_NONE;
            auto SendBufferLength = (uint32_t)CXPLAT_MIN(BytesToSend, (int64_t)MaxSendLength);
            for (uint32_t i = 0; i < SendBuffer->BufferCount; ++i) {
                SendBuffer->Buffers[i].Length = SendBufferLength;
            }
            if (InterlockedSubtract64(&BytesToSend, SendBufferLength) == 0) {
                Flags |= QUIC_SEND_FLAG_FIN;
            }
#ifdef CXPLAT_RAISE_IRQL
            CXPLAT_RAISE_IRQL();
#endif
            QUIC_STATUS Status =
                MsQuic->StreamSend(
                    QuicStream,
                    SendBuffer->Buffers,
                    SendBuffer->BufferCount,
                    Flags,
                    SendBuffer);
#ifdef CXPLAT_RAISE_IRQL
            CXPLAT_LOWER_IRQL();
#endif
            if (QUIC_FAILED(Status)) {
                InterlockedDecrement(&OutstandingSendRequestCount);
                delete SendBuffer;
            }
        }
    } else {
        delete SendBuffer;
    }
}

QUIC_STATUS
TestStream::HandleStreamEvent(
    _Inout_ QUIC_STREAM_EVENT* Event
    )
{
    uint64_t Param = 0;
    uint32_t ParamLength = sizeof(Param);

    switch (Event->Type) {

    case QUIC_STREAM_EVENT_RECEIVE:
        for (uint32_t i = 0; i < Event->RECEIVE.BufferCount; ++i) {
            HandleStreamRecv(
                Event->RECEIVE.Buffers[i].Buffer,
                Event->RECEIVE.Buffers[i].Length,
                Event->RECEIVE.Flags);
        }
        break;

    case QUIC_STREAM_EVENT_SEND_COMPLETE:
        HandleStreamSendComplete(
            Event->SEND_COMPLETE.Canceled,
            (QuicSendBuffer*)Event->SEND_COMPLETE.ClientContext);
        break;

    case QUIC_STREAM_EVENT_PEER_SEND_SHUTDOWN:
        AllDataReceived = true;
        RecvShutdown = true;
        CxPlatEventSet(EventRecvShutdownComplete);
        if (!IsPingSource) {
            Shutdown(QUIC_STREAM_SHUTDOWN_FLAG_GRACEFUL, QUIC_TEST_NO_ERROR);
        }
        break;

    case QUIC_STREAM_EVENT_PEER_SEND_ABORTED:
        AllDataReceived = false;
        RecvShutdown = true;
        CxPlatEventSet(EventRecvShutdownComplete);
        break;

    case QUIC_STREAM_EVENT_PEER_RECEIVE_ABORTED:
        break;

    case QUIC_STREAM_EVENT_SEND_SHUTDOWN_COMPLETE:
        if (Event->SEND_SHUTDOWN_COMPLETE.Graceful) {
            AllDataSent = true;
        }
        SendShutdown = true;
        CxPlatEventSet(EventSendShutdownComplete);
        break;

    case QUIC_STREAM_EVENT_SHUTDOWN_COMPLETE:
        IsShutdown = true;
        ConnectionShutdown = Event->SHUTDOWN_COMPLETE.ConnectionShutdown;
        ConnectionShutdownByApp = Event->SHUTDOWN_COMPLETE.ConnectionShutdownByApp;
        ConnectionClosedRemotely = Event->SHUTDOWN_COMPLETE.ConnectionClosedRemotely;
        ConnectionErrorCode = Event->SHUTDOWN_COMPLETE.ConnectionErrorCode;
        ConnectionCloseStatus = Event->SHUTDOWN_COMPLETE.ConnectionCloseStatus;
        if (QUIC_SUCCEEDED(
            MsQuic->GetParam(
                QuicStream,
                QUIC_PARAM_STREAM_0RTT_LENGTH,
                &ParamLength,
                &Param)) &&
            Param > 0) {
            UsedZeroRtt = true;
        }
        if (StreamShutdownCallback != nullptr) {
            StreamShutdownCallback(this);
        }
        break;

    default:
        break;
    }

    return QUIC_STATUS_SUCCESS;
}
