/**

 * FreeRDP: A Remote Desktop Protocol Implementation

 * Dynamic channel

 *

 * Copyright 2022 David Fort <contact@hardening-consulting.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.

 */

#include <freerdp/config.h>

#include <freerdp/log.h>

#include <freerdp/client/channels.h>

#define TAG FREERDP_TAG("genericdynvc")

static BOOL generic_update_connected_state(GENERIC_DYNVC_PLUGIN* plugin, BOOL val)

{

  if (!plugin || !plugin->dynvc_name)

    return FALSE;

  const size_t len = strlen(plugin->dynvc_name);

  plugin->dynvc_name[len + 1] = (char)val;

  return TRUE;

}

WINPR_ATTR_NODISCARD

static BOOL generic_get_connected_state(GENERIC_DYNVC_PLUGIN* plugin)

{

  if (!plugin || !plugin->dynvc_name)

    return FALSE;

  const size_t len = strlen(plugin->dynvc_name);

  return plugin->dynvc_name[len + 1];

}

static UINT generic_on_close(IWTSVirtualChannelCallback* pChannelCallback)

{

  GENERIC_CHANNEL_CALLBACK* callback = (GENERIC_CHANNEL_CALLBACK*)pChannelCallback;

  WINPR_ASSERT(callback);

  GENERIC_DYNVC_PLUGIN* plugin = (GENERIC_DYNVC_PLUGIN*)callback->plugin;

  WINPR_ASSERT(plugin);

  /* We're wrapping the originally set callback to update the channel state.

   * call the originally intended OnClose from the plugin

   */

  WINPR_ASSERT(plugin->channel_callbacks);

  const UINT rc =

      IFCALLRESULT(CHANNEL_RC_OK, plugin->channel_callbacks->OnClose, pChannelCallback);

  generic_update_connected_state(plugin, FALSE);

  return rc;

}

static UINT generic_on_new_channel_connection(IWTSListenerCallback* pListenerCallback,

                                              IWTSVirtualChannel* pChannel,

                                              WINPR_ATTR_UNUSED BYTE* Data,

                                              WINPR_ATTR_UNUSED BOOL* pbAccept,

                                              IWTSVirtualChannelCallback** ppCallback)

{

  GENERIC_LISTENER_CALLBACK* listener_callback = (GENERIC_LISTENER_CALLBACK*)pListenerCallback;

  if (!listener_callback || !listener_callback->plugin)

    return ERROR_INTERNAL_ERROR;

  GENERIC_DYNVC_PLUGIN* plugin = (GENERIC_DYNVC_PLUGIN*)listener_callback->plugin;

  WLog_Print(plugin->log, WLOG_TRACE, "...");

  if (generic_get_connected_state(plugin))

  {

    WLog_Print(plugin->log, WLOG_ERROR, "Channel already connected, terminating.");

    return ERROR_DEVICE_ALREADY_ATTACHED;

  }

  GENERIC_CHANNEL_CALLBACK* callback =

      (GENERIC_CHANNEL_CALLBACK*)calloc(1, plugin->channelCallbackSize);

  if (!callback)

  {

    WLog_Print(plugin->log, WLOG_ERROR, "calloc failed!");

    return CHANNEL_RC_NO_MEMORY;

  }

  /* implant configured channel callbacks */

  callback->iface = *plugin->channel_callbacks;

  callback->iface.OnClose = generic_on_close;

  callback->plugin = listener_callback->plugin;

  callback->channel_mgr = listener_callback->channel_mgr;

  callback->channel = pChannel;

  listener_callback->channel_callback = callback;

  listener_callback->channel = pChannel;

  *ppCallback = &callback->iface;

  generic_update_connected_state(plugin, TRUE);

  return CHANNEL_RC_OK;

}

static UINT generic_dynvc_plugin_initialize(IWTSPlugin* pPlugin,

                                            IWTSVirtualChannelManager* pChannelMgr)

{

  UINT rc = 0;

  GENERIC_DYNVC_PLUGIN* plugin = (GENERIC_DYNVC_PLUGIN*)pPlugin;

  if (!plugin)

    return CHANNEL_RC_BAD_CHANNEL_HANDLE;

  if (!pChannelMgr)

    return ERROR_INVALID_PARAMETER;

  if (plugin->initialized)

  {

    WLog_ERR(TAG, "[%s] channel initialized twice, aborting", plugin->dynvc_name);

    return ERROR_INVALID_DATA;

  }

  WLog_Print(plugin->log, WLOG_TRACE, "...");

  GENERIC_LISTENER_CALLBACK* listener_callback =

      (GENERIC_LISTENER_CALLBACK*)calloc(1, sizeof(GENERIC_LISTENER_CALLBACK));

  if (!listener_callback)

  {

    WLog_Print(plugin->log, WLOG_ERROR, "calloc failed!");

    return CHANNEL_RC_NO_MEMORY;

  }

  plugin->listener_callback = listener_callback;

  listener_callback->iface.OnNewChannelConnection = generic_on_new_channel_connection;

  listener_callback->plugin = pPlugin;

  listener_callback->channel_mgr = pChannelMgr;

  rc = pChannelMgr->CreateListener(pChannelMgr, plugin->dynvc_name, 0, &listener_callback->iface,

                                   &plugin->listener);

  plugin->listener->pInterface = plugin->iface.pInterface;

  plugin->initialized = (rc == CHANNEL_RC_OK);

  return rc;

}

static UINT generic_plugin_terminated(IWTSPlugin* pPlugin)

{

  GENERIC_DYNVC_PLUGIN* plugin = (GENERIC_DYNVC_PLUGIN*)pPlugin;

  UINT error = CHANNEL_RC_OK;

  if (!plugin)

    return CHANNEL_RC_BAD_CHANNEL_HANDLE;

  WLog_Print(plugin->log, WLOG_TRACE, "...");

  /* some channels (namely rdpei), look at initialized to see if they should continue to run */

  plugin->initialized = FALSE;

  if (plugin->terminatePluginFn)

    plugin->terminatePluginFn(plugin);

  if (plugin->listener_callback)

  {

    IWTSVirtualChannelManager* mgr = plugin->listener_callback->channel_mgr;

    if (mgr)

      IFCALL(mgr->DestroyListener, mgr, plugin->listener);

  }

  free(plugin->listener_callback);

  free(plugin->dynvc_name);

  free(plugin);

  return error;

}

static UINT generic_dynvc_plugin_attached(IWTSPlugin* pPlugin)

{

  GENERIC_DYNVC_PLUGIN* pluginn = (GENERIC_DYNVC_PLUGIN*)pPlugin;

  UINT error = CHANNEL_RC_OK;

  if (!pluginn)

    return CHANNEL_RC_BAD_CHANNEL_HANDLE;

  pluginn->attached = TRUE;

  return error;

}

static UINT generic_dynvc_plugin_detached(IWTSPlugin* pPlugin)

{

  GENERIC_DYNVC_PLUGIN* plugin = (GENERIC_DYNVC_PLUGIN*)pPlugin;

  UINT error = CHANNEL_RC_OK;

  if (!plugin)

    return CHANNEL_RC_BAD_CHANNEL_HANDLE;

  plugin->attached = FALSE;

  return error;

}

UINT freerdp_generic_DVCPluginEntry(IDRDYNVC_ENTRY_POINTS* pEntryPoints, const char* logTag,

                                    const char* name, size_t pluginSize, size_t channelCallbackSize,

                                    const IWTSVirtualChannelCallback* channel_callbacks,

                                    DYNVC_PLUGIN_INIT_FN initPluginFn,

                                    DYNVC_PLUGIN_TERMINATE_FN terminatePluginFn)

{

  GENERIC_DYNVC_PLUGIN* plugin = nullptr;

  UINT error = CHANNEL_RC_INITIALIZATION_ERROR;

  WINPR_ASSERT(pEntryPoints);

  WINPR_ASSERT(pEntryPoints->GetPlugin);

  WINPR_ASSERT(logTag);

  WINPR_ASSERT(name);

  WINPR_ASSERT(pluginSize >= sizeof(*plugin));

  WINPR_ASSERT(channelCallbackSize >= sizeof(GENERIC_CHANNEL_CALLBACK));

  plugin = (GENERIC_DYNVC_PLUGIN*)pEntryPoints->GetPlugin(pEntryPoints, name);

  if (plugin != nullptr)

    return CHANNEL_RC_ALREADY_INITIALIZED;

  plugin = (GENERIC_DYNVC_PLUGIN*)calloc(1, pluginSize);

  if (!plugin)

  {

    WLog_ERR(TAG, "calloc failed!");

    return CHANNEL_RC_NO_MEMORY;

  }

  plugin->log = WLog_Get(logTag);

  plugin->attached = TRUE;

  plugin->channel_callbacks = channel_callbacks;

  plugin->channelCallbackSize = channelCallbackSize;

  plugin->iface.Initialize = generic_dynvc_plugin_initialize;

  plugin->iface.Connected = nullptr;

  plugin->iface.Disconnected = nullptr;

  plugin->iface.Terminated = generic_plugin_terminated;

  plugin->iface.Attached = generic_dynvc_plugin_attached;

  plugin->iface.Detached = generic_dynvc_plugin_detached;

  plugin->terminatePluginFn = terminatePluginFn;

  if (initPluginFn)

  {

    rdpSettings* settings = pEntryPoints->GetRdpSettings(pEntryPoints);

    rdpContext* context = pEntryPoints->GetRdpContext(pEntryPoints);

    error = initPluginFn(plugin, context, settings);

    if (error != CHANNEL_RC_OK)

      goto error;

  }

  /* Since we do not have a private struct append the connected state after the name of the

   * channel */

  const size_t len = strlen(name);

  plugin->dynvc_name = calloc(len + 2, sizeof(char));

  if (!plugin->dynvc_name)

    goto error;

  strncpy(plugin->dynvc_name, name, len);

  error = pEntryPoints->RegisterPlugin(pEntryPoints, name, &plugin->iface);

  if (error == CHANNEL_RC_OK)

    return error;

error:

  generic_plugin_terminated(&plugin->iface);

  return error;

}