/**

 * FreeRDP: A Remote Desktop Protocol Implementation

 * RDPEDISP Virtual Channel Extension

 *

 * Copyright 2019 Kobi Mizrachi <kmizrachi18@gmail.com>

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *	 http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

#ifdef HAVE_CONFIG_H

#include "config.h"

#endif

#include "disp_main.h"

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <winpr/crt.h>

#include <winpr/synch.h>

#include <winpr/thread.h>

#include <winpr/stream.h>

#include <winpr/sysinfo.h>

#include <freerdp/channels/wtsvc.h>

#include <freerdp/channels/log.h>

#include <freerdp/server/disp.h>

#include "../disp_common.h"

#define TAG CHANNELS_TAG("rdpedisp.server")

/**

 * Function description

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static wStream* disp_server_single_packet_new(UINT32 type, UINT32 length)

{

	UINT error;

	DISPLAY_CONTROL_HEADER header;

	wStream* s = Stream_New(NULL, DISPLAY_CONTROL_HEADER_LENGTH + length);

	if (!s)

	{

		WLog_ERR(TAG, "Stream_New failed!");

		goto error;

	}

	header.type = type;

	header.length = length;

	if ((error = disp_write_header(s, &header)))

	{

		WLog_ERR(TAG, "Failed to write header with error %" PRIu32 "!", error);

		goto error;

	}

	return s;

error:

	Stream_Free(s, TRUE);

	return NULL;

}

static BOOL disp_server_is_monitor_layout_valid(DISPLAY_CONTROL_MONITOR_LAYOUT* monitor)

{

	if (monitor->Width < DISPLAY_CONTROL_MIN_MONITOR_WIDTH ||

	    monitor->Width > DISPLAY_CONTROL_MAX_MONITOR_WIDTH)

	{

		WLog_WARN(TAG, "Received invalid value for monitor->Width: %" PRIu32 "", monitor->Width);

		return FALSE;

	}

	if (monitor->Height < DISPLAY_CONTROL_MIN_MONITOR_HEIGHT ||

	    monitor->Height > DISPLAY_CONTROL_MAX_MONITOR_HEIGHT)

	{

		WLog_WARN(TAG, "Received invalid value for monitor->Height: %" PRIu32 "", monitor->Width);

		return FALSE;

	}

	if (monitor->PhysicalWidth < DISPLAY_CONTROL_MIN_PHYSICAL_MONITOR_WIDTH ||

	    monitor->PhysicalWidth > DISPLAY_CONTROL_MAX_PHYSICAL_MONITOR_WIDTH)

	{

		WLog_WARN(TAG, "Received invalid value for monitor->PhysicalWidth: %" PRIu32 "",

		          monitor->PhysicalWidth);

		return FALSE;

	}

	if (monitor->PhysicalHeight < DISPLAY_CONTROL_MIN_PHYSICAL_MONITOR_HEIGHT ||

	    monitor->PhysicalHeight > DISPLAY_CONTROL_MAX_PHYSICAL_MONITOR_HEIGHT)

	{

		WLog_WARN(TAG, "Received invalid value for monitor->Height: %" PRIu32 "",

		          monitor->PhysicalHeight);

		return FALSE;

	}

	switch (monitor->Orientation)

	{

		case ORIENTATION_LANDSCAPE:

		case ORIENTATION_PORTRAIT:

		case ORIENTATION_LANDSCAPE_FLIPPED:

		case ORIENTATION_PORTRAIT_FLIPPED:

			break;

		default:

			WLog_WARN(TAG, "Received incorrect value for monitor->Orientation: %" PRIu32 "",

			          monitor->Orientation);

			return FALSE;

	}

	return TRUE;

}

static UINT disp_recv_display_control_monitor_layout_pdu(wStream* s, DispServerContext* context)

{

	UINT32 error = CHANNEL_RC_OK;

	UINT32 index;

	DISPLAY_CONTROL_MONITOR_LAYOUT_PDU pdu;

	DISPLAY_CONTROL_MONITOR_LAYOUT* monitor;

	if (Stream_GetRemainingLength(s) < 8)

	{

		WLog_ERR(TAG, "not enough data!");

		return ERROR_INVALID_DATA;

	}

	Stream_Read_UINT32(s, pdu.MonitorLayoutSize); /* MonitorLayoutSize (4 bytes) */

	if (pdu.MonitorLayoutSize != DISPLAY_CONTROL_MONITOR_LAYOUT_SIZE)

	{

		WLog_ERR(TAG, "MonitorLayoutSize is set to %" PRIu32 ". expected %" PRIu32 "",

		         pdu.MonitorLayoutSize, DISPLAY_CONTROL_MONITOR_LAYOUT_SIZE);

		return ERROR_INVALID_DATA;

	}

	Stream_Read_UINT32(s, pdu.NumMonitors); /* NumMonitors (4 bytes) */

	if (pdu.NumMonitors > context->MaxNumMonitors)

	{

		WLog_ERR(TAG, "NumMonitors (%" PRIu32 ")> server MaxNumMonitors (%" PRIu32 ")",

		         pdu.NumMonitors, context->MaxNumMonitors);

		return ERROR_INVALID_DATA;

	}

	if (Stream_GetRemainingLength(s) < DISPLAY_CONTROL_MONITOR_LAYOUT_SIZE * pdu.NumMonitors)

	{

		WLog_ERR(TAG, "not enough data!");

		return ERROR_INVALID_DATA;

	}

	pdu.Monitors = (DISPLAY_CONTROL_MONITOR_LAYOUT*)calloc(pdu.NumMonitors,

	                                                       sizeof(DISPLAY_CONTROL_MONITOR_LAYOUT));

	if (!pdu.Monitors)

	{

		WLog_ERR(TAG, "disp_recv_display_control_monitor_layout_pdu(): calloc failed!");

		return CHANNEL_RC_NO_MEMORY;

	}

	WLog_DBG(TAG, "disp_recv_display_control_monitor_layout_pdu: NumMonitors=%" PRIu32 "",

	         pdu.NumMonitors);

	for (index = 0; index < pdu.NumMonitors; index++)

	{

		monitor = &(pdu.Monitors[index]);

		Stream_Read_UINT32(s, monitor->Flags);              /* Flags (4 bytes) */

		Stream_Read_UINT32(s, monitor->Left);               /* Left (4 bytes) */

		Stream_Read_UINT32(s, monitor->Top);                /* Top (4 bytes) */

		Stream_Read_UINT32(s, monitor->Width);              /* Width (4 bytes) */

		Stream_Read_UINT32(s, monitor->Height);             /* Height (4 bytes) */

		Stream_Read_UINT32(s, monitor->PhysicalWidth);      /* PhysicalWidth (4 bytes) */

		Stream_Read_UINT32(s, monitor->PhysicalHeight);     /* PhysicalHeight (4 bytes) */

		Stream_Read_UINT32(s, monitor->Orientation);        /* Orientation (4 bytes) */

		Stream_Read_UINT32(s, monitor->DesktopScaleFactor); /* DesktopScaleFactor (4 bytes) */

		Stream_Read_UINT32(s, monitor->DeviceScaleFactor);  /* DeviceScaleFactor (4 bytes) */

		WLog_DBG(TAG,

		         "\t%d : Flags: 0x%08" PRIX32 " Left/Top: (%" PRId32 ",%" PRId32 ") W/H=%" PRIu32

		         "x%" PRIu32 ")",

		         index, monitor->Flags, monitor->Left, monitor->Top, monitor->Width,

		         monitor->Height);

		WLog_DBG(TAG,

		         "\t   PhysicalWidth: %" PRIu32 " PhysicalHeight: %" PRIu32 " Orientation: %" PRIu32

		         "",

		         monitor->PhysicalWidth, monitor->PhysicalHeight, monitor->Orientation);

		if (!disp_server_is_monitor_layout_valid(monitor))

		{

			error = ERROR_INVALID_DATA;

			goto out;

		}

	}

	if (context)

		IFCALLRET(context->DispMonitorLayout, error, context, &pdu);

out:

	free(pdu.Monitors);

	return error;

}

static UINT disp_server_receive_pdu(DispServerContext* context, wStream* s)

{

	UINT error = CHANNEL_RC_OK;

	size_t beg, end;

	DISPLAY_CONTROL_HEADER header;

	beg = Stream_GetPosition(s);

	if ((error = disp_read_header(s, &header)))

	{

		WLog_ERR(TAG, "disp_read_header failed with error %" PRIu32 "!", error);

		return error;

	}

	switch (header.type)

	{

		case DISPLAY_CONTROL_PDU_TYPE_MONITOR_LAYOUT:

			if ((error = disp_recv_display_control_monitor_layout_pdu(s, context)))

				WLog_ERR(TAG,

				         "disp_recv_display_control_monitor_layout_pdu "

				         "failed with error %" PRIu32 "!",

				         error);

			break;

		default:

			error = CHANNEL_RC_BAD_PROC;

			WLog_WARN(TAG, "Received unknown PDU type: %" PRIu32 "", header.type);

			break;

	}

	end = Stream_GetPosition(s);

	if (end != (beg + header.length))

	{

		WLog_ERR(TAG, "Unexpected DISP pdu end: Actual: %d, Expected: %" PRIu32 "", end,

		         (beg + header.length));

		Stream_SetPosition(s, (beg + header.length));

	}

	return error;

}

static UINT disp_server_handle_messages(DispServerContext* context)

{

	DWORD BytesReturned;

	void* buffer;

	UINT ret = CHANNEL_RC_OK;

	DispServerPrivate* priv = context->priv;

	wStream* s = priv->input_stream;

	/* Check whether the dynamic channel is ready */

	if (!priv->isReady)

	{

		if (WTSVirtualChannelQuery(priv->disp_channel, WTSVirtualChannelReady, &buffer,

		                           &BytesReturned) == FALSE)

		{

			if (GetLastError() == ERROR_NO_DATA)

				return ERROR_NO_DATA;

			WLog_ERR(TAG, "WTSVirtualChannelQuery failed");

			return ERROR_INTERNAL_ERROR;

		}

		priv->isReady = *((BOOL*)buffer);

		WTSFreeMemory(buffer);

	}

	/* Consume channel event only after the disp dynamic channel is ready */

	Stream_SetPosition(s, 0);

	if (!WTSVirtualChannelRead(priv->disp_channel, 0, NULL, 0, &BytesReturned))

	{

		if (GetLastError() == ERROR_NO_DATA)

			return ERROR_NO_DATA;

		WLog_ERR(TAG, "WTSVirtualChannelRead failed!");

		return ERROR_INTERNAL_ERROR;

	}

	if (BytesReturned < 1)

		return CHANNEL_RC_OK;

	if (!Stream_EnsureRemainingCapacity(s, BytesReturned))

	{

		WLog_ERR(TAG, "Stream_EnsureRemainingCapacity failed!");

		return CHANNEL_RC_NO_MEMORY;

	}

	if (WTSVirtualChannelRead(priv->disp_channel, 0, (PCHAR)Stream_Buffer(s), Stream_Capacity(s),

	                          &BytesReturned) == FALSE)

	{

		WLog_ERR(TAG, "WTSVirtualChannelRead failed!");

		return ERROR_INTERNAL_ERROR;

	}

	Stream_SetLength(s, BytesReturned);

	Stream_SetPosition(s, 0);

	while (Stream_GetPosition(s) < Stream_Length(s))

	{

		if ((ret = disp_server_receive_pdu(context, s)))

		{

			WLog_ERR(TAG,

			         "disp_server_receive_pdu "

			         "failed with error %" PRIu32 "!",

			         ret);

			return ret;

		}

	}

	return ret;

}

static DWORD WINAPI disp_server_thread_func(LPVOID arg)

{

	DispServerContext* context = (DispServerContext*)arg;

	DispServerPrivate* priv = context->priv;

	DWORD status;

	DWORD nCount;

	HANDLE events[8];

	UINT error = CHANNEL_RC_OK;

	nCount = 0;

	events[nCount++] = priv->stopEvent;

	events[nCount++] = priv->channelEvent;

	/* Main virtual channel loop. RDPEDISP do not need version negotiation */

	while (TRUE)

	{

		status = WaitForMultipleObjects(nCount, events, FALSE, INFINITE);

		if (status == WAIT_FAILED)

		{

			error = GetLastError();

			WLog_ERR(TAG, "WaitForMultipleObjects failed with error %" PRIu32 "", error);

			break;

		}

		/* Stop Event */

		if (status == WAIT_OBJECT_0)

			break;

		if ((error = disp_server_handle_messages(context)))

		{

			WLog_ERR(TAG, "disp_server_handle_messages failed with error %" PRIu32 "", error);

			break;

		}

	}

	ExitThread(error);

	return error;

}

/**

 * Function description

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static UINT disp_server_open(DispServerContext* context)

{

	UINT rc = ERROR_INTERNAL_ERROR;

	DispServerPrivate* priv = context->priv;

	DWORD BytesReturned = 0;

	PULONG pSessionId = NULL;

	void* buffer;

	buffer = NULL;

	priv->SessionId = WTS_CURRENT_SESSION;

	if (WTSQuerySessionInformationA(context->vcm, WTS_CURRENT_SESSION, WTSSessionId,

	                                (LPSTR*)&pSessionId, &BytesReturned) == FALSE)

	{

		WLog_ERR(TAG, "WTSQuerySessionInformationA failed!");

		rc = ERROR_INTERNAL_ERROR;

		goto out_close;

	}

	priv->SessionId = (DWORD)*pSessionId;

	WTSFreeMemory(pSessionId);

	priv->disp_channel = (HANDLE)WTSVirtualChannelOpenEx(priv->SessionId, DISP_DVC_CHANNEL_NAME,

	                                                     WTS_CHANNEL_OPTION_DYNAMIC);

	if (!priv->disp_channel)

	{

		WLog_ERR(TAG, "WTSVirtualChannelOpenEx failed!");

		rc = GetLastError();

		goto out_close;

	}

	/* Query for channel event handle */

	if (!WTSVirtualChannelQuery(priv->disp_channel, WTSVirtualEventHandle, &buffer,

	                            &BytesReturned) ||

	    (BytesReturned != sizeof(HANDLE)))

	{

		WLog_ERR(TAG,

		         "WTSVirtualChannelQuery failed "

		         "or invalid returned size(%" PRIu32 ")",

		         BytesReturned);

		if (buffer)

			WTSFreeMemory(buffer);

		rc = ERROR_INTERNAL_ERROR;

		goto out_close;

	}

	CopyMemory(&priv->channelEvent, buffer, sizeof(HANDLE));

	WTSFreeMemory(buffer);

	if (priv->thread == NULL)

	{

		if (!(priv->stopEvent = CreateEvent(NULL, TRUE, FALSE, NULL)))

		{

			WLog_ERR(TAG, "CreateEvent failed!");

			rc = ERROR_INTERNAL_ERROR;

		}

		if (!(priv->thread =

		          CreateThread(NULL, 0, disp_server_thread_func, (void*)context, 0, NULL)))

		{

			WLog_ERR(TAG, "CreateEvent failed!");

			CloseHandle(priv->stopEvent);

			priv->stopEvent = NULL;

			rc = ERROR_INTERNAL_ERROR;

		}

	}

	return CHANNEL_RC_OK;

out_close:

	WTSVirtualChannelClose(priv->disp_channel);

	priv->disp_channel = NULL;

	priv->channelEvent = NULL;

	return rc;

}

static UINT disp_server_packet_send(DispServerContext* context, wStream* s)

{

	UINT ret;

	ULONG written;

	if (!WTSVirtualChannelWrite(context->priv->disp_channel, (PCHAR)Stream_Buffer(s),

	                            Stream_GetPosition(s), &written))

	{

		WLog_ERR(TAG, "WTSVirtualChannelWrite failed!");

		ret = ERROR_INTERNAL_ERROR;

		goto out;

	}

	if (written < Stream_GetPosition(s))

	{

		WLog_WARN(TAG, "Unexpected bytes written: %" PRIu32 "/%" PRIuz "", written,

		          Stream_GetPosition(s));

	}

	ret = CHANNEL_RC_OK;

out:

	Stream_Free(s, TRUE);

	return ret;

}

/**

 * Function description

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static UINT disp_server_send_caps_pdu(DispServerContext* context)

{

	wStream* s = disp_server_single_packet_new(DISPLAY_CONTROL_PDU_TYPE_CAPS, 12);

	if (!s)

	{

		WLog_ERR(TAG, "disp_server_single_packet_new failed!");

		return CHANNEL_RC_NO_MEMORY;

	}

	Stream_Write_UINT32(s, context->MaxNumMonitors);        /* MaxNumMonitors (4 bytes) */

	Stream_Write_UINT32(s, context->MaxMonitorAreaFactorA); /* MaxMonitorAreaFactorA (4 bytes) */

	Stream_Write_UINT32(s, context->MaxMonitorAreaFactorB); /* MaxMonitorAreaFactorB (4 bytes) */

	return disp_server_packet_send(context, s);

}

/**

 * Function description

 *

 * @return 0 on success, otherwise a Win32 error code

 */

static UINT disp_server_close(DispServerContext* context)

{

	UINT error = CHANNEL_RC_OK;

	DispServerPrivate* priv = context->priv;

	if (priv->thread)

	{

		SetEvent(priv->stopEvent);

		if (WaitForSingleObject(priv->thread, INFINITE) == WAIT_FAILED)

		{

			error = GetLastError();

			WLog_ERR(TAG, "WaitForSingleObject failed with error %" PRIu32 "", error);

			return error;

		}

		CloseHandle(priv->thread);

		CloseHandle(priv->stopEvent);

		priv->thread = NULL;

		priv->stopEvent = NULL;

	}

	if (priv->disp_channel)

	{

		WTSVirtualChannelClose(priv->disp_channel);

		priv->disp_channel = NULL;

	}

	return error;

}

DispServerContext* disp_server_context_new(HANDLE vcm)

{

	DispServerContext* context;

	DispServerPrivate* priv;

	context = (DispServerContext*)calloc(1, sizeof(DispServerContext));

	if (!context)

	{

		WLog_ERR(TAG, "disp_server_context_new(): calloc DispServerContext failed!");

		goto out_free;

	}

	priv = context->priv = (DispServerPrivate*)calloc(1, sizeof(DispServerPrivate));

	if (!context->priv)

	{

		WLog_ERR(TAG, "disp_server_context_new(): calloc DispServerPrivate failed!");

		goto out_free;

	}

	priv->input_stream = Stream_New(NULL, 4);

	if (!priv->input_stream)

	{

		WLog_ERR(TAG, "Stream_New failed!");

		goto out_free_priv;

	}

	context->vcm = vcm;

	context->Open = disp_server_open;

	context->Close = disp_server_close;

	context->DisplayControlCaps = disp_server_send_caps_pdu;

	priv->isReady = FALSE;

	return context;

out_free_priv:

	free(context->priv);

out_free:

	free(context);

	return NULL;

}

void disp_server_context_free(DispServerContext* context)

{

	if (!context)

		return;

	disp_server_close(context);

	if (context->priv)

	{

		Stream_Free(context->priv->input_stream, TRUE);

		free(context->priv);

	}

	free(context);

}