/*
 * eset_rtp (ESET Real-time file system protection module)
 * Copyright (C) 1992-2025 ESET, spol. s r.o.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 * In case of any questions, you can contact us at ESET, spol. s r.o., Einsteinova 24, 851 01 Bratislava, Slovakia.
 */

#include "ertp_qos.h"

#include <linux/atomic.h>
#include <linux/mutex.h>
#include <linux/sched/sysctl.h>

#include "ertp_array.h"
#include "ertp_cache.h"
#include "ertp_logs.h"
#include "ertp_qos_queue.h"
#include "ertp_stats.h"

#define DEFAULT_SCANNER_QUEUE_SIZE 64

struct ertp_scanner_t {
  atomic_t enabled;
  atomic_t tgid;
};

struct ertp_scanner_queue_t {
  struct mutex lock;
  struct ertp_array events;
};

struct ertp_qos_t {
  struct mutex lock;

  struct ertp_qos_queue_t qos_queue;
  struct ertp_scanner_queue_t scanner_queue;
  struct ertp_scanner_t scanner;
  atomic_t suspend_active;
  atomic_t initialized;

  struct wait_queue_head request_available;

  unsigned int scanner_timeout;
  unsigned int mem_dev_timeout;
  unsigned int delay;
};

static struct ertp_qos_t qos;

static void deinit_item(void *data) {
  struct ertp_event *event = (struct ertp_event *)(data);
  ertp_event_unblock(event, 0);
  ertp_event_unref(event);
}

int ertp_qos_init(void) {
  int err = ertp_array_init(&qos.scanner_queue.events,
                            DEFAULT_SCANNER_QUEUE_SIZE, deinit_item);
  if (err) {
    return err;
  }

  mutex_init(&qos.scanner_queue.lock);
  mutex_init(&qos.lock);

  init_waitqueue_head(&qos.request_available);
  ertp_qos_queue_init(&qos.qos_queue);

  atomic_set(&qos.scanner.enabled, false);
  atomic_set(&qos.scanner.tgid, -1);

  atomic_set(&qos.suspend_active, false);
  atomic_set(&qos.initialized, false);

  qos.scanner_timeout = 0;
  qos.mem_dev_timeout = 0;
  qos.delay = 0;

  return 0;
}

void ertp_qos_deinit(void) {
  ertp_qos_queue_deinit(&qos.qos_queue);
  ertp_array_deinit(&qos.scanner_queue.events);
  mutex_destroy(&qos.scanner_queue.lock);
  mutex_destroy(&qos.lock);
}

bool ertp_qos_enabled(void) {
  return atomic_read(&qos.initialized) && atomic_read(&qos.scanner.enabled) &&
         !atomic_read(&qos.suspend_active);
}

void ertp_qos_scanner_enable(void) {
  mutex_lock(&qos.lock);
  atomic_set(&qos.scanner.enabled, true);
  mutex_unlock(&qos.lock);
}

static void flush_all_events(struct ertp_qos_queue_t *swap) {
  ertp_qos_queue_swap_and_reinit(&qos.qos_queue, swap);

  mutex_lock(&qos.scanner_queue.lock);
  ertp_array_clear(&qos.scanner_queue.events);
  mutex_unlock(&qos.scanner_queue.lock);
}

void ertp_qos_suspend_enable(void) {
  struct ertp_qos_queue_t swap;

  mutex_lock(&qos.lock);

  atomic_set(&qos.suspend_active, true);
  flush_all_events(&swap);

  mutex_unlock(&qos.lock);

  ertp_qos_queue_deinit(&swap);
}

void ertp_qos_scanner_disable(void) {
  struct ertp_qos_queue_t swap;

  mutex_lock(&qos.lock);

  atomic_set(&qos.scanner.enabled, false);
  flush_all_events(&swap);

  mutex_unlock(&qos.lock);

  ertp_qos_queue_deinit(&swap);
}

void ertp_qos_suspend_disable(void) {
  mutex_lock(&qos.lock);
  atomic_set(&qos.suspend_active, false);
  mutex_unlock(&qos.lock);
}

struct ertp_event *ertp_qos_get_next_event(size_t max_size) {
  struct ertp_event *event = NULL;
  mutex_lock(&qos.lock);
  event = ertp_qos_queue_get_next(&qos.qos_queue, max_size);
  mutex_unlock(&qos.lock);
  return event;
}

static bool cmp_events_ids(const void *data, const void *ctx) {
  const struct ertp_event *event = (const struct ertp_event *)data;
  return event->id == *(int32_t *)ctx;
}

static struct ertp_event *scanner_remove_event(int32_t id) {
  struct ertp_event *event = NULL;
  size_t position;

  position = ertp_array_find(&qos.scanner_queue.events, &id, cmp_events_ids);
  if (position != SIZE_MAX) {
    event = ertp_array_at(&qos.scanner_queue.events, position);
    ertp_array_remove(&qos.scanner_queue.events, position);
  } else {
    ertp_pr_edlog("context for event with id %d not found", id);
  }

  return event;
}

static void scanner_remove_event_context(int32_t id) {
  struct ertp_event *event = NULL;

  mutex_lock(&qos.scanner_queue.lock);
  event = scanner_remove_event(id);
  mutex_unlock(&qos.scanner_queue.lock);

  if (event) {
    ertp_event_unref(event);
  }
}

static void drop_event_context(struct ertp_event *event) {
  struct ertp_event *removed_event_context = NULL;
  mutex_lock(&qos.scanner_queue.lock);

  ertp_event_disable_context(event);

  removed_event_context = scanner_remove_event(event->id);
  if (removed_event_context) {
    ertp_event_unref(removed_event_context);
  }

  mutex_unlock(&qos.scanner_queue.lock);
}

static void drop_event_and_unref(struct ertp_event *event,
                                 enum ertp_qos_queue_ids qos_queue_id) {
  mutex_lock(&qos.lock);

  if (!ertp_qos_queue_remove(&qos.qos_queue, qos_queue_id, event)) {
    drop_event_context(event);
  }
  mutex_unlock(&qos.lock);
}

static void delay_processing(struct ertp_event *event) {
  long timeout_jiffies = msecs_to_jiffies(qos.delay);

  wait_for_completion_timeout(&event->wait_scanner, timeout_jiffies);
}

static void wait_for_scanner_start_uninterruptible(struct ertp_event *event,
                                                   long timeout_jiffies) {
  unsigned long ret;

  do {
    ret = wait_for_completion_timeout(&event->wait_qos, timeout_jiffies);

    if (ret > 0) {
      return;
    }

    ertp_pr_debug(
        "event is waiting for too long to be sent to scanner (id: %d)",
        event->id);

  } while (unlikely(ret == 0));
}

static int wait_for_scanner_start(struct ertp_event *event,
                                  enum ertp_qos_queue_ids qos_queue_id) {
  long ret;
  long timeout_jiffies = msecs_to_jiffies(10000);
  do {
    ret = wait_for_completion_interruptible_timeout(&event->wait_qos,
                                                    timeout_jiffies);

    if (ret > 0) {
      return 0;
    }

    if (ret == -ERESTARTSYS) {
      switch (event->signal_handler) {
        case ERTP_SIGNAL_HANDLER_RESTART_INTR:
          drop_event_and_unref(event, qos_queue_id);
          ertp_event_unblock(event, -ERESTARTSYS);
          break;

        case ERTP_SIGNAL_HANDLER_RESTART:
          drop_event_and_unref(event, qos_queue_id);
          ertp_event_unblock(event, -ERESTARTNOINTR);
          break;

        case ERTP_SIGNAL_HANDLER_WAIT:

          drop_event_context(event);

          delay_processing(event);
          ertp_event_unblock(event, 0);
          break;

        case ERTP_SIGNAL_HANDLER_IGNORE:
          wait_for_scanner_start_uninterruptible(event, timeout_jiffies);
          return 0;

        case ERTP_SIGNAL_HANDLER_UNINITIALIZED:
          BUG();
      }

      return -EINTR;
    }

    ertp_pr_debug(
        "event is waiting for too long to be sent to scanner (id: %d)",
        event->id);

  } while (unlikely(ret == 0));

  return 0;
}

static bool wait_for_scanner_end(struct ertp_event *event) {
  long timeout_jiffies = msecs_to_jiffies(qos.scanner_timeout);

  long ret = wait_for_completion_interruptible_timeout(&event->wait_scanner,
                                                       timeout_jiffies);

  if (ret > 0) {
    return 0;
  }

  if (ret == -ERESTARTSYS) {
    switch (event->signal_handler) {
      case ERTP_SIGNAL_HANDLER_RESTART_INTR:
        scanner_remove_event_context(event->id);
        ertp_event_unblock(event, -ERESTARTSYS);
        break;

      case ERTP_SIGNAL_HANDLER_RESTART:
        scanner_remove_event_context(event->id);
        ertp_event_unblock(event, -ERESTARTNOINTR);
        break;

      case ERTP_SIGNAL_HANDLER_WAIT:
        scanner_remove_event_context(event->id);
        delay_processing(event);
        ertp_event_unblock(event, 0);
        break;

      case ERTP_SIGNAL_HANDLER_IGNORE:

        if (wait_for_completion_timeout(&event->wait_scanner,
                                        timeout_jiffies) == 0) {
          goto timeout;
        }
        return 0;

      case ERTP_SIGNAL_HANDLER_UNINITIALIZED:
        BUG();
    }

    return -EINTR;
  }

timeout:

  scanner_remove_event_context(event->id);
  ertp_event_unblock(event, 0);
  return -ETIMEDOUT;
}

static int ertp_qos_handle_new_event(struct ertp_event *event) {
  bool ok = false;
  enum ertp_qos_queue_ids qos_queue_id;

  ertp_qos_event_data_fill(&event->qos_data);

  mutex_lock(&qos.lock);
  if (ertp_qos_enabled()) {
    qos_queue_id = ertp_qos_queue_insert(&qos.qos_queue, event);
    ok = true;
  }
  mutex_unlock(&qos.lock);

  if (ok) {
    int ret = 0;

    wake_up_interruptible(&qos.request_available);

    ret = wait_for_scanner_start(event, qos_queue_id);
    if (ret != 0) {
      return ret;
    }

    ret = wait_for_scanner_end(event);
    if (ret != 0) {
      return ret;
    }
  }

  return 0;
}

int ertp_qos_handle_generic_event(
    const struct ertp_event_generic_info *event_info) {
  int rv;
  int ret;
  struct ertp_event *event_base = ertp_event_generic_new(event_info);
  struct ertp_event_generic *event = ERTP_EVENT_GENERIC(event_base);

  if (IS_ERR(event_base)) {
    return PTR_ERR(event_base);
  }

  ret = ertp_qos_handle_new_event(event_base);
  if (unlikely(ret == -ETIMEDOUT)) {
    ertp_pr_debug(
        "event is waiting for too long for response from scanner (type: "
        "GENERIC %s, id: %d, file: %s, process: %s)",
        ertp_event_type_to_str(event->type), event_base->id,
        event_info->file_path->ptr, event_info->process_path->ptr);
  } else if (unlikely(ret == -EINTR)) {
    ertp_pr_edlog(
        "event handling was interrupted by a signal (type: GENERIC %s, id: %d, "
        "file: %s, process: %s)",
        ertp_event_type_to_str(event->type), event_base->id,
        event_info->file_path->ptr, event_info->process_path->ptr);
  } else if (unlikely(ret != 0)) {
    ertp_pr_edlog(
        "event handling ended with non-zero return value: (errno: %d, type: "
        "GENERIC %s, id: %d, file: %s, process: %s)",
        ret, ertp_event_type_to_str(event->type), event_base->id,
        event_info->file_path->ptr, event_info->process_path->ptr);
  }

  rv = ertp_event_get_result(event_base);

  ertp_event_unref(event_base);
  return rv;
}

int ertp_qos_handle_memory_event(
    const struct ertp_event_memory_info *event_info) {
  int rv;
  int ret;
  struct ertp_event *event_base =
      ertp_event_memory_new(event_info, qos.mem_dev_timeout);
  struct ertp_event_memory *event = ERTP_EVENT_MEMORY(event_base);

  if (IS_ERR(event_base)) {
    return PTR_ERR(event_base);
  }

  ret = ertp_qos_handle_new_event(event_base);
  if (unlikely(ret == -ETIMEDOUT)) {
    ertp_pr_debug(
        "event is waiting for too long for response from scanner (type: MEMORY "
        "%s, id: %d, process: %s, mem_dev file: %s)",
        ertp_event_type_to_str(event->type), event_base->id,
        event->process_path->ptr, event->dev_info->name);
  } else if (unlikely(ret == -EINTR)) {
    ertp_pr_edlog(
        "event handling was interrupted by a signal (type: MEMORY %s, id: %d, "
        "process: %s, mem_dev file: %s)",
        ertp_event_type_to_str(event->type), event_base->id,
        event->process_path->ptr, event->dev_info->name);
  } else if (unlikely(ret != 0)) {
    ertp_pr_edlog(
        "event handling ended with non-zero return value: (errno: %d, type: "
        "MEMORY %s, id: %d, process: %s, mem_dev file: %s)",
        ret, ertp_event_type_to_str(event->type), event_base->id,
        event->process_path->ptr, event->dev_info->name);
  }

  rv = ertp_event_get_result(event_base);

  ertp_mem_dev_destroy(event);
  ertp_event_unref(event_base);
  return rv;
}

void ertp_qos_handle_remove_event(
    const struct ertp_event_remove_info *event_info) {
  struct ertp_event *event_base = ertp_event_remove_new(event_info);
  struct ertp_event_remove *event = ERTP_EVENT_REMOVE(event_base);
  int ret;

  if (IS_ERR(event_base)) {
    return;
  }

  ertp_cache_remove(event->ino, event->dev);

  ret = ertp_qos_handle_new_event(event_base);
  if (unlikely(ret == -ETIMEDOUT)) {
    ertp_pr_debug(
        "event is waiting for too long for response from scanner (type: "
        "REMOVE, id: %d, file: %s, process: %s)",
        event_base->id, event_info->file_path->ptr,
        event_info->process_path->ptr);
  } else if (unlikely(ret == -EINTR)) {
    ertp_pr_edlog(
        "event handling was interrupted by a signal (type: REMOVE, id: %d, "
        "file: %s, process: %s)",
        event_base->id, event_info->file_path->ptr,
        event_info->process_path->ptr);
  } else if (unlikely(ret != 0)) {
    ertp_pr_edlog(
        "event handling ended with non-zero return value: (errno: %d, type: "
        "REMOVE, id: %d, file: %s, process: %s)",
        ret, event_base->id, event_info->file_path->ptr,
        event_info->process_path->ptr);
  }

  ertp_event_unref(event_base);
}

bool ertp_qos_scanner_init(unsigned int scanner_timeout,
                           unsigned int mem_dev_timeout, unsigned int delay) {
  if (atomic_read(&qos.initialized)) {
    return false;
  }

  qos.scanner_timeout = scanner_timeout;
  qos.mem_dev_timeout = mem_dev_timeout;
  qos.delay = delay;

  atomic_set(&qos.initialized, true);
  return true;
}

bool ertp_qos_scanner_check_thread(void) {
  return (atomic_read(&qos.scanner.tgid) == current->tgid);
}

int ertp_qos_scanner_connect(pid_t tgid) {
  int ret = -EEXIST;
  mutex_lock(&qos.lock);
  if (atomic_read(&qos.scanner.tgid) == -1) {
    ertp_pr_edlog("registering new scanner with pid(tgid) %d",
                  atomic_read(&qos.scanner.tgid));
    atomic_set(&qos.scanner.tgid, tgid);
    ret = 0;
  }
  mutex_unlock(&qos.lock);
  return ret;
}

void ertp_qos_scanner_disconnect(void) {
  struct ertp_qos_queue_t swap;

  mutex_lock(&qos.lock);

  atomic_set(&qos.scanner.enabled, false);

  if (atomic_read(&qos.scanner.tgid) != -1) {
    ertp_pr_edlog("unregistering scanner with pid(tgid) %d",
                  atomic_read(&qos.scanner.tgid));
    atomic_set(&qos.scanner.tgid, -1);
  }
  flush_all_events(&swap);

  mutex_unlock(&qos.lock);

  ertp_qos_queue_deinit(&swap);
}

int ertp_qos_scanner_save_sent_event(struct ertp_event *event) {
  int ret = 0;

  mutex_lock(&qos.scanner_queue.lock);
  if (event->save_context) {
    ret = ertp_array_insert(&qos.scanner_queue.events, event);
    if (ret == 0) {
      ertp_event_ref(event);
      ertp_stats_report_scan_started(
          &event->stats, ertp_array_size(&qos.scanner_queue.events));
    }
  }
  mutex_unlock(&qos.scanner_queue.lock);
  return ret;
}

static struct ertp_event *scanner_get_finished_event(int32_t id) {
  struct ertp_event *event = NULL;

  mutex_lock(&qos.scanner_queue.lock);
  event = scanner_remove_event(id);
  mutex_unlock(&qos.scanner_queue.lock);

  wake_up_interruptible(&qos.request_available);

  if (event) {
    ertp_stats_report_scan_finished(&event->stats, event->class);
  }

  return event;
}

static inline enum ertp_event_type convert_event_type(
    enum eset_rtp_cache_data_type type) {
  switch (type) {
    case ERTP_TYPE_CACHE_OPEN:
      return ERTP_EVENT_TYPE_OPEN;
    case ERTP_TYPE_CACHE_CLOSE:
      return ERTP_EVENT_TYPE_CLOSE;
    case ERTP_TYPE_CACHE_EXEC:
      return ERTP_EVENT_TYPE_EXEC;
  }

  BUG();
}

void ertp_qos_scanner_process_scan_response(
    const struct eset_rtp_set_av_status *response) {
  struct ertp_event *event = scanner_get_finished_event(response->id);
  if (event) {
    int result;

    switch (response->result) {
      case ERTP_RESULT_ALLOW:
        result = 0;
        break;
      case ERTP_RESULT_DENY:
        result = -EPERM;
        break;
      default:
        BUG();
    }

    ertp_event_finished(event, result);
    ertp_event_unref(event);
  }

  if (response->result == ERTP_RESULT_ALLOW &&
      (response->flags & ERTP_AV_RESULT_FLAG_SHOULD_CACHE)) {
    ertp_timespec_t ctime;
    ctime.tv_sec = response->cache_data.ctime.tv_sec;
    ctime.tv_nsec = response->cache_data.ctime.tv_nsec;

    ertp_cache_add(response->cache_data.ino, response->cache_data.dev,
                   convert_event_type(response->cache_data.type), ctime);
  }
}

void ertp_qos_scanner_process_event_response(
    const struct eset_rtp_remove_event_finished *response) {
  struct ertp_event *event = scanner_get_finished_event(response->id);
  if (event) {
    ertp_event_finished(event, 0);
    ertp_event_unref(event);
  }
}

static bool event_ready_for_scanner(void) {
  bool ready;
  mutex_lock(&qos.lock);
  ready = ertp_qos_queue_event_ready(&qos.qos_queue);
  mutex_unlock(&qos.lock);
  return ready;
}

unsigned int ertp_qos_scanner_wait(struct file *file, poll_table *wait) {
  unsigned int mask;

  poll_wait(file, &qos.request_available, wait);

  mask = POLLOUT | POLLWRNORM;

  if (event_ready_for_scanner()) {
    mask |= POLLIN | POLLRDNORM;
  }

  return mask;
}