/*
 * Exports definition for virtio socket WSK interface
 *
 * Copyright (c) 2021 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.
 */

#ifndef _VIOWSK_H
#define _VIOWSK_H

#ifdef _MSC_VER
#pragma once
#endif //_MSC_VER

#define VIOSOCK_WSK_MEMORY_TAG (ULONG)'kswV'

typedef struct _VIOWSK_REG_CONTEXT
{
    PVOID ClientContext;
    WSK_CLIENT_DISPATCH ClientDispatch;
    volatile LONG NPICount;
    ERESOURCE NPILock;
    UNICODE_STRING VIOSockLinkName;
    PDEVICE_OBJECT VIOSockDevice;
    PFILE_OBJECT VIOSockMainFileObject;
}VIOWSK_REG_CONTEXT,*PVIOWSK_REG_CONTEXT;

#define VIOWSK_PROVIDER_VERSION 1

_At_((void*)Irp->IoStatus.Information, __drv_allocatesMem(Mem))
NTSTATUS
WSKAPI
VioWskSocket(
    _In_ PWSK_CLIENT              Client,
    _In_ ADDRESS_FAMILY           AddressFamily,
    _In_ USHORT                   SocketType,
    _In_ ULONG                    Protocol,
    _In_ ULONG                    Flags,
    _In_opt_ PVOID                SocketContext,
    _In_opt_ CONST VOID          *Dispatch,
    _In_opt_ PEPROCESS            OwningProcess,
    _In_opt_ PETHREAD             OwningThread,
    _In_opt_ PSECURITY_DESCRIPTOR SecurityDescriptor,
    _Inout_ PIRP                  Irp
);

_At_(Irp->IoStatus.Information, __drv_allocatesMem(Mem))
NTSTATUS
WSKAPI
VioWskSocketConnect(
    _In_ PWSK_CLIENT                               Client,
    _In_ USHORT                                    SocketType,
    _In_ ULONG                                     Protocol,
    _In_ PSOCKADDR                                 LocalAddress,
    _In_ PSOCKADDR                                 RemoteAddress,
    _Reserved_ ULONG                               Flags,
    _In_opt_ PVOID                                 SocketContext,
    _In_opt_ CONST WSK_CLIENT_CONNECTION_DISPATCH *Dispatch,
    _In_opt_ PEPROCESS                             OwningProcess,
    _In_opt_ PETHREAD                              OwningThread,
    _In_opt_ PSECURITY_DESCRIPTOR                  SecurityDescriptor,
    _Inout_ PIRP                                   Irp
    );

NTSTATUS
WSKAPI
VioWskControlClient(
    _In_ PWSK_CLIENT                    Client,
    _In_ ULONG                          ControlCode,
    _In_ SIZE_T                         InputSize,
    _In_reads_bytes_opt_(InputSize) PVOID    InputBuffer,
    _In_ SIZE_T                         OutputSize,
    _Out_writes_bytes_opt_(OutputSize) PVOID  OutputBuffer,
    _Out_opt_ SIZE_T                   *OutputSizeReturned,
    _Inout_opt_ PIRP                    Irp
);

_At_(*Result, __drv_allocatesMem(Mem))
NTSTATUS
WSKAPI
VioWskGetAddressInfo(
    _In_ PWSK_CLIENT          Client,
    _In_opt_ PUNICODE_STRING  NodeName,
    _In_opt_ PUNICODE_STRING  ServiceName,
    _In_opt_ ULONG            NameSpace,
    _In_opt_ GUID            *Provider,
    _In_opt_ PADDRINFOEXW     Hints,
    _Outptr_ PADDRINFOEXW *Result,
    _In_opt_ PEPROCESS        OwningProcess,
    _In_opt_ PETHREAD         OwningThread,
    _Inout_ PIRP              Irp
);

_At_(AddrInfo, __drv_freesMem(Mem))
VOID
WSKAPI
VioWskFreeAddressInfo(
    _In_ PWSK_CLIENT  Client,
    _In_ PADDRINFOEXW AddrInfo
);

NTSTATUS
WSKAPI
VioWskGetNameInfo(
    _In_ PWSK_CLIENT          Client,
    _In_ PSOCKADDR            SockAddr,
    _In_ ULONG                SockAddrLength,
    _Out_opt_ PUNICODE_STRING NodeName,
    _Out_opt_ PUNICODE_STRING ServiceName,
    _In_ ULONG                Flags,
    _In_opt_ PEPROCESS        OwningProcess,
    _In_opt_ PETHREAD         OwningThread,
    _Inout_ PIRP              Irp
);

__inline
NTSTATUS
VioWskCompleteIrp(
    _In_ PIRP Irp,
    _In_ NTSTATUS Status,
    _In_ ULONG_PTR Information
)
{
    Irp->IoStatus.Information = Information;
    Irp->IoStatus.Status = Status;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return Status;
}

typedef struct _VIOWSK_SOCKET
{
    WSK_SOCKET WskSocket;
    PWSK_CLIENT Client;
    volatile LONG RequestCount;
    IO_REMOVE_LOCK CloseRemoveLock;
    PVOID SocketContext;
    ULONG Type;
    ULONG GuestId;
    HANDLE FileHandle;
    PFILE_OBJECT FileObject;
    WSK_CLIENT_LISTEN_DISPATCH ListenDispatch;
    WSK_CLIENT_CONNECTION_DISPATCH ConnectionDispatch;
    union
    {
        WSK_PROVIDER_BASIC_DISPATCH ProviderBasicDispatch;
        WSK_PROVIDER_LISTEN_DISPATCH ProviderListenDispatch;
        WSK_PROVIDER_CONNECTION_DISPATCH ProviderConnectionDispatch;
#if (NTDDI_VERSION >= NTDDI_WIN10_RS2)
        WSK_PROVIDER_STREAM_DISPATCH ProviderStreamDispatch;
#endif // if (NTDDI_VERSION >= NTDDI_WIN10_RS2)
    }ProviderDispatch;
}VIOWSK_SOCKET,*PVIOWSK_SOCKET;

extern WSK_PROVIDER_BASIC_DISPATCH gBasicDispatch;
extern WSK_PROVIDER_LISTEN_DISPATCH gListenDispatch;
extern WSK_PROVIDER_CONNECTION_DISPATCH gConnectionDispatch;
#if (NTDDI_VERSION >= NTDDI_WIN10_RS2)
extern WSK_PROVIDER_STREAM_DISPATCH gStreamDispatch;
#endif // if (NTDDI_VERSION >= NTDDI_WIN10_RS2)

#endif //_VIOWSK_H
