/*
 * Main include file
 * This file contains various routines and globals
 *
 * Copyright (c) 2019 Virtuozzo International GmbH
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met :
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and / or other materials provided with the distribution.
 * 3. Neither the names of the copyright holders nor the names of their contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */
#if !defined(VIOSOCK_H)
#define VIOSOCK_H
#include "public.h"

#define VIOSOCK_DRIVER_MEMORY_TAG (ULONG)'cosV'

#pragma pack (push)
#pragma pack (1)

typedef enum _VIRTIO_VSOCK_OP {
    VIRTIO_VSOCK_OP_INVALID = 0,

    /* Connect operations */
    VIRTIO_VSOCK_OP_REQUEST = 1,
    VIRTIO_VSOCK_OP_RESPONSE = 2,
    VIRTIO_VSOCK_OP_RST = 3,
    VIRTIO_VSOCK_OP_SHUTDOWN = 4,

    /* To send payload */
    VIRTIO_VSOCK_OP_RW = 5,

    /* Tell the peer our credit info */
    VIRTIO_VSOCK_OP_CREDIT_UPDATE = 6,
    /* Request the peer to send the credit info to us */
    VIRTIO_VSOCK_OP_CREDIT_REQUEST = 7,
}VIRTIO_VSOCK_OP;

/* VIRTIO_VSOCK_OP_SHUTDOWN flags values */
enum virtio_vsock_shutdown {
    VIRTIO_VSOCK_SHUTDOWN_RCV = 1,
    VIRTIO_VSOCK_SHUTDOWN_SEND = 2,
    VIRTIO_VSOCK_SHUTDOWN_MASK = VIRTIO_VSOCK_SHUTDOWN_RCV | VIRTIO_VSOCK_SHUTDOWN_SEND,
};

typedef struct _VIRTIO_VSOCK_HDR {
    ULONG64 src_cid;
    ULONG64 dst_cid;
    ULONG32 src_port;
    ULONG32 dst_port;
    ULONG32 len;
    USHORT  type;
    USHORT  op;
    ULONG32 flags;
    ULONG32 buf_alloc;
    ULONG32 fwd_cnt;
}VIRTIO_VSOCK_HDR, *PVIRTIO_VSOCK_HDR;

typedef enum _VIRTIO_VSOCK_EVENT_ID {
    VIRTIO_VSOCK_EVENT_TRANSPORT_RESET = 0,
}VIRTIO_VSOCK_EVENT_ID;

typedef struct _VIRTIO_VSOCK_EVENT {
    ULONG32 id;
}VIRTIO_VSOCK_EVENT, *PVIRTIO_VSOCK_EVENT;

#pragma pack (pop)

typedef struct virtqueue VIOSOCK_VQ, *PVIOSOCK_VQ;
typedef struct VirtIOBufferDescriptor VIOSOCK_SG_DESC, *PVIOSOCK_SG_DESC;

#define VIOSOCK_VQ_RX  0
#define VIOSOCK_VQ_TX  1
#define VIOSOCK_VQ_EVT 2
#define VIOSOCK_VQ_MAX 3

#define VIRTIO_VSOCK_DEFAULT_RX_BUF_SIZE	(1024 * 4)
#define VIRTIO_VSOCK_MAX_PKT_BUF_SIZE		(1024 * 64)

#define VSOCK_CLOSE_TIMEOUT                 SEC_TO_NANO(8)
#define VSOCK_DEFAULT_CONNECT_TIMEOUT       SEC_TO_NANO(2)
#define VSOCK_DEFAULT_BUFFER_SIZE           (1024 * 256)
#define VSOCK_DEFAULT_BUFFER_MAX_SIZE       (1024 * 256)
#define VSOCK_DEFAULT_BUFFER_MIN_SIZE       128

#define VIRTIO_VSOCK_MAX_EVENTS 8

#define LAST_RESERVED_PORT  1023
#define MAX_PORT_RETRIES    24
#define INVALID_THREAD_ID   LONG64_ERROR
//////////////////////////////////////////////////////////////////////////
#define VIOSOCK_TIMER_TOLERANCE MSEC_TO_NANO(50)
typedef struct _VIOSOCK_TIMER
{
    WDFTIMER    Timer;
    LONGLONG    StartTime; //ticks when timer started
    LONGLONG    Timeout;   //timeout in 100ns
    ULONG       StartRefs;
}VIOSOCK_TIMER,*PVIOSOCK_TIMER;

#define VIOSOCK_DEVICE_NAME L"\\Device\\Viosock"

typedef struct _DEVICE_CONTEXT {

    VIRTIO_WDF_DRIVER           VDevice;

    WDFDEVICE                   ThisDevice;

    //Recv packets
    WDFQUEUE                    ReadQueue;

    WDFSPINLOCK                 RxLock;
    _Guarded_by_(RxLock) PVIOSOCK_VQ                 RxVq;
    PVOID                       RxPktVA;        //contiguous array of VIOSOCK_RX_PKT
    PHYSICAL_ADDRESS            RxPktPA;
    _Guarded_by_(RxLock) SINGLE_LIST_ENTRY           RxPktList;      //postponed requests
    ULONG                       RxPktNum;
    ULONG                       RxCbBuffersNum;
    WDFLOOKASIDE                RxCbBufferMemoryList;
    _Guarded_by_(RxLock) SINGLE_LIST_ENTRY           RxCbBuffers;    //list or Rx buffers

    //Send packets
    WDFQUEUE                    WriteQueue;

    WDFSPINLOCK                 TxLock;
    _Guarded_by_(TxLock) PVIOSOCK_VQ                 TxVq;
    _Guarded_by_(TxLock) PVIRTIO_DMA_MEMORY_SLICED   TxPktSliced;
    LONG                       TxPktNum;       //Num of slices in TxPktSliced
    _Guarded_by_(TxLock) LONG                        TxQueuedReply;
    _Guarded_by_(TxLock) LIST_ENTRY                  TxList;
    _Guarded_by_(TxLock) VIOSOCK_TIMER               TxTimer;
    WDFLOOKASIDE                TxMemoryList;
    _Interlocked_ volatile LONG               TxEnqueued;
    _Interlocked_ volatile LONG               TxPktAllocated;

    WDFLOOKASIDE                AcceptMemoryList;

    //Events
    PVIOSOCK_VQ                 EvtVq;
    PVIRTIO_VSOCK_EVENT         EvtVA;
    PHYSICAL_ADDRESS            EvtPA;
    ULONG                       EvtRstOccured;

    WDFSPINLOCK                 BoundLock;
    _Guarded_by_(BoundLock) WDFCOLLECTION               BoundList;

    WDFSPINLOCK                 ConnectedLock;
    _Guarded_by_(ConnectedLock) WDFCOLLECTION               ConnectedList;

    WDFWAITLOCK                 SelectLock;
    _Guarded_by_(SelectLock) LIST_ENTRY                  SelectList;
    _Interlocked_ volatile LONG               SelectInProgress;
    WDFWORKITEM                 SelectWorkitem;
    _Guarded_by_(SelectLock) VIOSOCK_TIMER               SelectTimer;

    WDFQUEUE                    IoCtlQueue;

    WDFINTERRUPT                WdfInterrupt;

    _Interlocked_ volatile LONG             SocketId; //for debug

    VIRTIO_VSOCK_CONFIG         Config;
} DEVICE_CONTEXT, *PDEVICE_CONTEXT;

WDF_DECLARE_CONTEXT_TYPE_WITH_NAME(DEVICE_CONTEXT, GetDeviceContext);

typedef enum _VIOSOCK_STATE
{
    VIOSOCK_STATE_CLOSE = 0,
    VIOSOCK_STATE_CONNECTING = 1,
    VIOSOCK_STATE_CONNECTED = 2,
    VIOSOCK_STATE_CLOSING = 3,
    VIOSOCK_STATE_LISTEN = 4,
}VIOSOCK_STATE;

#define SOCK_CONTROL    0x01
#define SOCK_BOUND      0x02
#define SOCK_LINGER     0x04
#define SOCK_NON_BLOCK  0x08
#define SOCK_LOOPBACK   0x10

typedef struct _VIOSOCK_ACCEPT_ENTRY
{
    LIST_ENTRY      ListEntry;
    WDFMEMORY       Memory;
    WDFFILEOBJECT   ConnectSocket;
    ULONG32         dst_cid;
    ULONG32         dst_port;
    ULONG32         peer_buf_alloc;
    ULONG32         peer_fwd_cnt;
}VIOSOCK_ACCEPT_ENTRY, *PVIOSOCK_ACCEPT_ENTRY;

typedef struct _SOCKET_CONTEXT {

    WDFFILEOBJECT   ThisSocket;

    _Interlocked_ volatile LONG             Flags;
    LONG            SocketId; //for debug

    VIRTIO_VSOCK_TYPE  type;

    ULONG32         dst_cid;
    ULONG32         src_port;
    ULONG32         dst_port;

    WDFSPINLOCK     StateLock;
    _Interlocked_ volatile VIOSOCK_STATE    State;
    LONGLONG        ConnectTimeout;
    ULONG           SendTimeout;
    ULONG           RecvTimeout;
    ULONG32         BufferMinSize;
    ULONG32         BufferMaxSize;
    ULONG32         PeerShutdown;
    ULONG32         Shutdown;

    KEVENT          CloseEvent;

    _Guarded_by_(StateLock) PKEVENT         EventObject;
    _Guarded_by_(StateLock) ULONG           EventsMask;
    _Guarded_by_(StateLock) ULONG           Events;
    _Guarded_by_(StateLock) NTSTATUS        EventsStatus[FD_MAX_EVENTS];

    WDFSPINLOCK     RxLock;         //accept list lock for listen socket
    LONG64 RxProcessingThreadId; // Prevents recursion in
    _Guarded_by_(RxLock) LIST_ENTRY      RxCbList;
    _Guarded_by_(RxLock) volatile ULONG           RxBytes;        //used bytes in rx buffer
    _Guarded_by_(RxLock) ULONG           RxBuffers;      //used rx buffers (for debug)

    WDFQUEUE        ReadQueue;
    _Guarded_by_(RxLock) VIOSOCK_TIMER   ReadTimer;

    _Guarded_by_(RxLock) WDFREQUEST      PendedRequest;
    VIOSOCK_TIMER        PendedTimer;

    _Guarded_by_(RxLock) LIST_ENTRY      AcceptList;
    LONG            Backlog;
    _Interlocked_ volatile LONG   AcceptPended;

    ULONG32         buf_alloc;
    _Guarded_by_(RxLock) ULONG32         fwd_cnt;
    ULONG32         last_fwd_cnt;

    _Guarded_by_(StateLock) ULONG32         peer_buf_alloc;
    _Guarded_by_(StateLock) ULONG32         peer_fwd_cnt;
    _Guarded_by_(StateLock) ULONG32         tx_cnt;

    USHORT          LingerTime;
    WDFFILEOBJECT   LoopbackSocket;

    volatile LONG   SelectRefs[FDSET_MAX];
} SOCKET_CONTEXT, *PSOCKET_CONTEXT;

WDF_DECLARE_CONTEXT_TYPE_WITH_NAME(SOCKET_CONTEXT, GetSocketContext);

#define VIOSockIsFlag(s,f) ((s)->Flags & (f))
#define VIOSockSetFlag(s,f) (InterlockedOr(&(s)->Flags, (f)) & (f))
#define VIOSockResetFlag(s,f) (InterlockedAnd(&(s)->Flags, ~(f)) & (f))

#define VIOSockIsNonBlocking(s) VIOSockIsFlag((s), SOCK_NON_BLOCK)

#define GetSocketContextFromRequest(r) GetSocketContext(WdfRequestGetFileObject((r)))

#define GetDeviceContextFromRequest(r) GetDeviceContext(WdfFileObjectGetDevice(WdfRequestGetFileObject((r))))

#define GetDeviceContextFromSocket(s) GetDeviceContext(WdfFileObjectGetDevice((s)->ThisSocket))

#define IsControlRequest(r) VIOSockIsFlag(GetSocketContextFromRequest(r), SOCK_CONTROL)

#define IsLoopbackSocket(s) (VIOSockIsFlag(s,SOCK_LOOPBACK))

//////////////////////////////////////////////////////////////////////////
//Device functions

EVT_WDF_DRIVER_DEVICE_ADD   VIOSockEvtDeviceAdd;

NTSTATUS
VIOSockInterruptInit(
    IN WDFDEVICE hDevice
);

//////////////////////////////////////////////////////////////////////////
//Socket functions
EVT_WDF_DEVICE_FILE_CREATE  VIOSockCreateStub;
EVT_WDF_FILE_CLEANUP          VIOSockCleanup;

NTSTATUS
VIOSockDeviceControl(
    IN WDFREQUEST Request,
    IN ULONG      IoControlCode,
    IN OUT size_t *pLength
);

WDFFILEOBJECT
VIOSockGetSocketFromHandle(
    IN PDEVICE_CONTEXT pContext,
    IN ULONGLONG       uSocket,
    IN BOOLEAN         bIs32BitProcess
);

VOID
VIOSockHandleTransportReset(
    IN PDEVICE_CONTEXT pContext
);

NTSTATUS
VIOSockBoundListInit(
    IN WDFDEVICE hDevice
);

NTSTATUS
VIOSockBoundAdd(
    IN PSOCKET_CONTEXT pSocket,
    IN ULONG32         svm_port
);

PSOCKET_CONTEXT
VIOSockBoundFindByPort(
    IN PDEVICE_CONTEXT pContext,
    IN ULONG32         ulSrcPort
);

PSOCKET_CONTEXT
VIOSockBoundFindByPortUnlocked(
    IN PDEVICE_CONTEXT pContext,
    IN ULONG32         ulSrcPort
);

PSOCKET_CONTEXT
VIOSockBoundFindByFile(
    IN PDEVICE_CONTEXT pContext,
    IN PFILE_OBJECT pFileObject
);

NTSTATUS
VIOSockConnectedListInit(
    IN WDFDEVICE hDevice
);

PSOCKET_CONTEXT
VIOSockConnectedFindByRxPkt(
    IN PDEVICE_CONTEXT      pContext,
    IN PVIRTIO_VSOCK_HDR    pPkt
);

_Requires_lock_held_(pSocket->StateLock)
VIOSOCK_STATE
VIOSockStateSet(
    IN PSOCKET_CONTEXT pSocket,
    IN VIOSOCK_STATE   NewState
);

_Requires_lock_not_held_(pSocket->StateLock)
VIOSOCK_STATE
VIOSockStateSetLocked(
    IN PSOCKET_CONTEXT pSocket,
    IN VIOSOCK_STATE   NewState
);

#define VIOSockStateGet(s) ((s)->State)

_Requires_lock_not_held_(pSocket->StateLock)
BOOLEAN
VIOSockShutdownFromPeer(
    PSOCKET_CONTEXT pSocket,
    ULONG uFlags
);

VOID
VIOSockDoClose(
    PSOCKET_CONTEXT pSocket
);

__inline
BOOLEAN
VIOSockIsDone(
    PSOCKET_CONTEXT pSocket
)
{
    return !!KeReadStateEvent(&pSocket->CloseEvent);
}

_Requires_lock_held_(pSocket->RxLock)
NTSTATUS
VIOSockPendedRequestSetEx(
    IN PSOCKET_CONTEXT  pSocket,
    IN WDFREQUEST Request,
    IN LONGLONG Timeout,
    IN BOOLEAN Resume
);

#define VIOSockPendedRequestSet(s,r,t) VIOSockPendedRequestSetEx(s,r,t,FALSE)

_Requires_lock_not_held_(pSocket->RxLock)
NTSTATUS
VIOSockPendedRequestSetLocked(
    IN PSOCKET_CONTEXT  pSocket,
    IN WDFREQUEST Request,
    IN LONGLONG Timeout
);

#define VIOSockPendedRequestSetResume(s,r) VIOSockPendedRequestSetEx(s,r,0,TRUE)

_Requires_lock_not_held_(pSocket->RxLock)
NTSTATUS
VIOSockPendedRequestSetResumeLocked(
    IN PSOCKET_CONTEXT  pSocket,
    IN WDFREQUEST       Request
);

_Requires_lock_held_(pSocket->RxLock)
NTSTATUS
VIOSockPendedRequestGet(
    IN PSOCKET_CONTEXT  pSocket,
    OUT WDFREQUEST      *Request
);


_Requires_lock_not_held_(pSocket->RxLock)
__declspec(noinline)
NTSTATUS
VIOSockPendedRequestGetLocked(
    IN PSOCKET_CONTEXT  pSocket,
    OUT WDFREQUEST      *Request
);

NTSTATUS
VIOSockAcceptInitSocket(
    PSOCKET_CONTEXT pAcceptSocket,
    PSOCKET_CONTEXT pListenSocket
);

_Requires_lock_not_held_(pListenSocket->RxLock)
NTSTATUS
VIOSockAcceptEnqueuePkt(
    IN PSOCKET_CONTEXT      pListenSocket,
    IN PVIRTIO_VSOCK_HDR    pPkt
);

_Requires_lock_not_held_(pListenSocket->RxLock)
VOID
VIOSockAcceptRemovePkt(
    IN PSOCKET_CONTEXT      pListenSocket,
    IN PVIRTIO_VSOCK_HDR    pPkt
);

_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
VIOSockSelectRun(
    IN PSOCKET_CONTEXT pSocket
);

/*
 * WinSock 2 extension -- bit values and indices for FD_XXX network events
 */
#define FD_READ_BIT      0
#define FD_READ          (1 << FD_READ_BIT)
#define FD_WRITE_BIT     1
#define FD_WRITE         (1 << FD_WRITE_BIT)
#define FD_OOB_BIT       2
#define FD_OOB           (1 << FD_OOB_BIT)
#define FD_ACCEPT_BIT    3
#define FD_ACCEPT        (1 << FD_ACCEPT_BIT)
#define FD_CONNECT_BIT   4
#define FD_CONNECT       (1 << FD_CONNECT_BIT)
#define FD_CLOSE_BIT     5
#define FD_CLOSE         (1 << FD_CLOSE_BIT)
#define FD_QOS_BIT       6
#define FD_QOS           (1 << FD_QOS_BIT)
#define FD_GROUP_QOS_BIT 7
#define FD_GROUP_QOS     (1 << FD_GROUP_QOS_BIT)
#define FD_ROUTING_INTERFACE_CHANGE_BIT 8
#define FD_ROUTING_INTERFACE_CHANGE     (1 << FD_ROUTING_INTERFACE_CHANGE_BIT)
#define FD_ADDRESS_LIST_CHANGE_BIT 9
#define FD_ADDRESS_LIST_CHANGE     (1 << FD_ADDRESS_LIST_CHANGE_BIT)

_Requires_lock_not_held_(pSocket->StateLock)
VOID
VIOSockEventSetBit(
    IN PSOCKET_CONTEXT pSocket,
    IN ULONG uSetBit,
    IN NTSTATUS Status
);

_Requires_lock_not_held_(pSocket->StateLock)
VOID
VIOSockEventSetBitLocked(
    IN PSOCKET_CONTEXT pSocket,
    IN ULONG uSetBit,
    IN NTSTATUS Status
);

_Requires_lock_not_held_(pSocket->StateLock)
VOID
VIOSockEventClearBit(
    IN PSOCKET_CONTEXT pSocket,
    IN ULONG uClearBit
);

//////////////////////////////////////////////////////////////////////////
//Tx functions

NTSTATUS
VIOSockWriteQueueInit(
    IN WDFDEVICE hDevice
);

NTSTATUS
VIOSockTxVqInit(
    IN PDEVICE_CONTEXT pContext
);

_Requires_lock_not_held_(pContext->TxLock)
VOID
VIOSockTxVqCleanup(
    IN PDEVICE_CONTEXT pContext
);

_Requires_lock_not_held_(pContext->TxLock)
VOID
VIOSockTxVqProcess(
    IN PDEVICE_CONTEXT pContext
);

_Requires_lock_not_held_(pSocket->StateLock)
NTSTATUS
VIOSockStateValidate(
    PSOCKET_CONTEXT pSocket,
    BOOLEAN         bTx
);

NTSTATUS
VIOSockTxEnqueue(
    IN PSOCKET_CONTEXT  pSocket,
    IN VIRTIO_VSOCK_OP  Op,
    IN ULONG32          Flags OPTIONAL,
    IN BOOLEAN          Reply,
    IN WDFREQUEST       Request OPTIONAL
);

#define VIOSockSendCreditUpdate(s) VIOSockTxEnqueue(s, VIRTIO_VSOCK_OP_CREDIT_UPDATE, 0, FALSE, WDF_NO_HANDLE)

#define VIOSockSendConnect(s) VIOSockTxEnqueue(s, VIRTIO_VSOCK_OP_REQUEST, 0, FALSE, WDF_NO_HANDLE)

#define VIOSockSendShutdown(s, f) VIOSockTxEnqueue(s, VIRTIO_VSOCK_OP_SHUTDOWN, f, FALSE, WDF_NO_HANDLE)

#define VIOSockSendWrite(s, rq) VIOSockTxEnqueue(s, VIRTIO_VSOCK_OP_RW, 0, FALSE, rq)

#define VIOSockSendReset(s, r) VIOSockTxEnqueue(s, VIRTIO_VSOCK_OP_RST, 0, r, WDF_NO_HANDLE)

#define VIOSockSendResponse(s) VIOSockTxEnqueue(s, VIRTIO_VSOCK_OP_RESPONSE, 0, TRUE, WDF_NO_HANDLE)

_Requires_lock_not_held_(pContext->TxLock)
NTSTATUS
VIOSockSendResetNoSock(
    IN PDEVICE_CONTEXT pContext,
    IN PVIRTIO_VSOCK_HDR pHeader
);

_Requires_lock_not_held_(pSocket->StateLock)
__inline
ULONG32
VIOSockTxGetCredit(
    IN PSOCKET_CONTEXT pSocket,
    IN ULONG32 uCredit

)
{
    ULONG32 uRet;

    WdfSpinLockAcquire(pSocket->StateLock);
    uRet = pSocket->peer_buf_alloc - (pSocket->tx_cnt - pSocket->peer_fwd_cnt);
    if (uRet > uCredit)
        uRet = uCredit;
    pSocket->tx_cnt += uRet;
    WdfSpinLockRelease(pSocket->StateLock);

    return uRet;
}

_Requires_lock_not_held_(pSocket->StateLock)
__inline
VOID
VIOSockTxPutCredit(
    IN PSOCKET_CONTEXT pSocket,
    IN ULONG32 uCredit

)
{
    WdfSpinLockAcquire(pSocket->StateLock);
    pSocket->tx_cnt -= uCredit;
    WdfSpinLockRelease(pSocket->StateLock);
}

_Requires_lock_held_(pSocket->StateLock)
__inline
LONG
VIOSockTxHasSpace(
    IN PSOCKET_CONTEXT pSocket
)
{
    LONG lBytes = (LONG)pSocket->peer_buf_alloc - (pSocket->tx_cnt - pSocket->peer_fwd_cnt);
    if (lBytes < 0)
        lBytes = 0;
    return lBytes;
}

_Requires_lock_not_held_(pSocket->StateLock)
__inline
LONG
VIOSockTxSpaceUpdate(
    IN PSOCKET_CONTEXT pSocket,
    IN PVIRTIO_VSOCK_HDR pPkt
)
{
    LONG uSpace;

    WdfSpinLockAcquire(pSocket->StateLock);
    pSocket->peer_buf_alloc = pPkt->buf_alloc;
    pSocket->peer_fwd_cnt = pPkt->fwd_cnt;
    uSpace = VIOSockTxHasSpace(pSocket);
    WdfSpinLockRelease(pSocket->StateLock);

    return uSpace;
}

__inline
BOOLEAN
VIOSockTxMoreReplies(
    IN PDEVICE_CONTEXT  pContext
)
{
    return pContext->TxQueuedReply < (LONG)pContext->RxPktNum;
}

_Requires_lock_not_held_(pContext->TxLock)
VOID
VIOSockTxCleanup(
    PDEVICE_CONTEXT pContext,
    WDFFILEOBJECT   Socket,
    NTSTATUS        Status
);

VOID
VIOSockWriteIoSuspend(
    IN PDEVICE_CONTEXT pContext
);

VOID
VIOSockWriteIoRestart(
    IN PDEVICE_CONTEXT pContext
);

//////////////////////////////////////////////////////////////////////////
//Rx functions

NTSTATUS
VIOSockRxVqInit(
    IN PDEVICE_CONTEXT pContext
);

_Requires_lock_not_held_(pContext->RxLock)
VOID
VIOSockRxVqCleanup(
    IN PDEVICE_CONTEXT pContext
);

_Requires_lock_not_held_(pSocket->StateLock)
__inline
VOID
VIOSockRxIncTxPkt(
    IN PSOCKET_CONTEXT pSocket,
    IN OUT PVIRTIO_VSOCK_HDR pPkt
)
{
    WdfSpinLockAcquire(pSocket->StateLock);
    pSocket->last_fwd_cnt = pSocket->fwd_cnt;
    pPkt->fwd_cnt = pSocket->fwd_cnt;
    pPkt->buf_alloc = pSocket->buf_alloc;
    WdfSpinLockRelease(pSocket->StateLock);
}

_Requires_lock_not_held_(pContext->RxLock)
VOID
VIOSockRxVqProcess(
    IN PDEVICE_CONTEXT pContext
);

_Requires_lock_not_held_(pSocket->RxLock)
NTSTATUS
VIOSockRxRequestEnqueueCb(
    IN PSOCKET_CONTEXT  pSocket,
    IN WDFREQUEST       Request,
    IN ULONG            Length
);

__inline
ULONG
VIOSockRxHasData(
    IN PSOCKET_CONTEXT pSocket
)
{
    return pSocket->RxBytes;
}

NTSTATUS
VIOSockReadQueueInit(
    IN WDFDEVICE hDevice
);

NTSTATUS
VIOSockReadSocketQueueInit(
    IN PSOCKET_CONTEXT pSocket
);

_Requires_lock_not_held_(pSocket->RxLock)
BOOLEAN
VIOSockReadDequeueCb(
    IN PSOCKET_CONTEXT  pSocket
);

_Requires_lock_not_held_(pSocket->RxLock)
__inline
VIOSockReadProcessDequeueCb(
    IN PSOCKET_CONTEXT pSocket
)
{
    while (VIOSockReadDequeueCb(pSocket));
}

_Requires_lock_not_held_(pSocket->RxLock)
VOID
VIOSockReadCleanupCb(
    IN PSOCKET_CONTEXT pSocket
);

NTSTATUS
VIOSockReadWithFlags(
    IN WDFREQUEST Request
);

//////////////////////////////////////////////////////////////////////////
//Event functions
_IRQL_requires_max_(PASSIVE_LEVEL)
NTSTATUS
VIOSockEvtVqInit(
    IN PDEVICE_CONTEXT pContext
);

_IRQL_requires_max_(PASSIVE_LEVEL)
VOID
VIOSockEvtVqCleanup(
    IN PDEVICE_CONTEXT pContext
);

_IRQL_requires_max_(DISPATCH_LEVEL)
VOID
VIOSockEvtVqProcess(
    IN PDEVICE_CONTEXT pContext
);

//////////////////////////////////////////////////////////////////////////
BOOLEAN
VIOSockLoopbackAcceptDequeue(
    IN PSOCKET_CONTEXT pAcceptSocket,
    IN PVIOSOCK_ACCEPT_ENTRY pAcceptEntry
);

NTSTATUS
VIOSockLoopbackTxEnqueue(
    IN PSOCKET_CONTEXT  pSocket,
    IN VIRTIO_VSOCK_OP  Op,
    IN ULONG32          Flags OPTIONAL,
    IN WDFREQUEST       Request OPTIONAL,
    IN ULONG            Length OPTIONAL
);

//////////////////////////////////////////////////////////////////////////
__inline
NTSTATUS
VIOSockTimerCreate(
    IN PVIOSOCK_TIMER   pTimer,
    IN WDFOBJECT        ParentObject,
    IN PFN_WDF_TIMER    EvtTimerFunc
)
{
    WDF_OBJECT_ATTRIBUTES   Attributes;
    WDF_TIMER_CONFIG        timerConfig;

    pTimer->Timeout = 0;
    pTimer->StartTime = 0;
    pTimer->StartRefs = 0;

    WDF_TIMER_CONFIG_INIT(&timerConfig, EvtTimerFunc);

    WDF_OBJECT_ATTRIBUTES_INIT(&Attributes);
    Attributes.ParentObject = ParentObject;

    return WdfTimerCreate(&timerConfig, &Attributes, &pTimer->Timer);
}

VOID
VIOSockTimerStart(
    IN PVIOSOCK_TIMER   pTimer,
    IN LONGLONG         Timeout
);

__inline
VOID
VIOSockTimerSet(
    IN PVIOSOCK_TIMER   pTimer,
    IN LONGLONG         Timeout
)
{
    LARGE_INTEGER liTicks;

    if (!Timeout || Timeout == LONGLONG_MAX)
    {
        ASSERT(!pTimer->StartRefs);
        pTimer->StartTime = 0;
        pTimer->Timeout = 0;
        return;
    }

    ASSERT(Timeout > VIOSOCK_TIMER_TOLERANCE);
    if (Timeout <= VIOSOCK_TIMER_TOLERANCE)
        Timeout = VIOSOCK_TIMER_TOLERANCE + 1;

    KeQueryTickCount(&liTicks);

    pTimer->StartTime = liTicks.QuadPart;
    pTimer->Timeout = Timeout;
    WdfTimerStart(pTimer->Timer, -Timeout);
}

__inline
VOID
VIOSockTimerCancel(
    IN PVIOSOCK_TIMER pTimer
)
{
    WdfTimerStop(pTimer->Timer, FALSE);
    pTimer->Timeout = 0;
    pTimer->StartTime = 0;
}

__inline
VOID
VIOSockTimerDeref(
    IN PVIOSOCK_TIMER   pTimer,
    IN BOOLEAN          bStop
)
{
    ASSERT(pTimer->StartRefs);
    if (--pTimer->StartRefs == 0 && bStop)
        VIOSockTimerCancel(pTimer);
}

__inline
LONGLONG
VIOSockTimerPassed(
    IN PVIOSOCK_TIMER pTimer
)
{
    LARGE_INTEGER liTicks;

    KeQueryTickCount(&liTicks);

    return (liTicks.QuadPart - pTimer->StartTime) * KeQueryTimeIncrement();
}

__inline
BOOLEAN
VIOSockTimerSuspend(
    IN PVIOSOCK_TIMER pTimer
)
{
    if (WdfTimerStop(pTimer->Timer, FALSE))
    {
        return pTimer->Timeout > VIOSockTimerPassed(pTimer) + VIOSOCK_TIMER_TOLERANCE;
    }
    return TRUE;
}

__inline
BOOLEAN
VIOSockTimerResume(
    IN PVIOSOCK_TIMER pTimer
)
{
    LONGLONG llTimeout = VIOSockTimerPassed(pTimer);

    if (pTimer->Timeout > llTimeout + VIOSOCK_TIMER_TOLERANCE)
    {
        VIOSockTimerSet(pTimer, pTimer->Timeout - llTimeout);
        return TRUE;
    }
    return FALSE;
}

//////////////////////////////////////////////////////////////////////////

__inline
VOID
VioSockReference(
    IN WDFFILEOBJECT Socket
)
{
    ObReferenceObject(WdfFileObjectWdmGetFileObject(Socket));
    WdfObjectReference(Socket);

    return;
}


__inline
VOID
VioSockDereference(
    IN WDFFILEOBJECT Socket
)
{
    WdfObjectDereference(Socket);
    ObDereferenceObject(WdfFileObjectWdmGetFileObject(Socket));

    return;
}


#endif /* VIOSOCK_H */
