/*
 * eset_wap (ESET Web Access 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 "ewap_tcp_map.h"

#include <linux/socket.h>
#include <net/sock.h>

#include "ewap_helpers.h"

static struct kmem_cache *ewap_tcp_node_allocator = NULL;

static inline int ewap_tcp_key_part_cmp(uint64_t p1, uint64_t p2) {
  if (p1 < p2) {
    return -1;
  }

  if (p1 > p2) {
    return 1;
  }

  return 0;
}

static inline int ewap_tcp_key_cmp(const struct ewap_tcp_key *k1,
                                   const struct ewap_tcp_key *k2) {
  int res;

  res = ewap_tcp_key_part_cmp(k1->port, k2->port);
  if (res != 0) {
    return res;
  }

  res = ewap_tcp_key_part_cmp(k1->saddr_hi, k2->saddr_hi);
  if (res != 0) {
    return res;
  }

  res = ewap_tcp_key_part_cmp(k1->saddr_lo, k2->saddr_lo);
  if (res != 0) {
    return res;
  }

  res = ewap_tcp_key_part_cmp(k1->daddr_hi, k2->daddr_hi);
  if (res != 0) {
    return res;
  }

  res = ewap_tcp_key_part_cmp(k1->daddr_lo, k2->daddr_lo);
  if (res != 0) {
    return res;
  }

  return 0;
}

int ewap_tcp_node_allocator_init(void) {
  ewap_tcp_node_allocator = kmem_cache_create(
      "ewap_tcp_node_allocator", sizeof(struct ewap_tcp_node),
      __alignof__(struct ewap_tcp_node), SLAB_RECLAIM_ACCOUNT, NULL);

  if (!ewap_tcp_node_allocator) {
    return -ENOMEM;
  }

  return 0;
}

void ewap_tcp_node_allocator_deinit(void) {
  if (ewap_tcp_node_allocator) {
    kmem_cache_destroy(ewap_tcp_node_allocator);
  }
}

struct ewap_tcp_node *ewap_tcp_node_new(struct ewap_tcp_key key, pid_t pid,
                                        uid_t uid) {
  struct ewap_path *path;
  struct ewap_tcp_node *node =
      kmem_cache_zalloc(ewap_tcp_node_allocator, GFP_KERNEL);
  if (node == NULL) {
    return ERR_PTR(-ENOMEM);
  }

  path = ewap_path_new();
  if (IS_ERR(path)) {
    long err = PTR_ERR(path);
    if (err == -ENODATA) {
      ewap_pr_log(EWAP_LOG_ERRORS,
                  "cannot obtain process path for pid %d (error: -ENODATA)",
                  pid);
    } else {
      ewap_pr_warning("cannot obtain process path for pid %d (error: %ld)", pid,
                      err);
    }
    path = NULL;
  }

  node->key = key;
  node->pid = pid;
  node->uid = uid;
  node->path = path;
  node->tcp_map = NULL;
  INIT_LIST_HEAD(&node->list);

  ewap_pr_log(EWAP_LOG_CONNECTIONS,
              "tcp_node created (key: %u %llu %llu %llu %llu), pid: %d, uid: "
              "%u, path %s",
              key.port, key.saddr_hi, key.saddr_lo, key.daddr_hi, key.daddr_lo,
              node->pid, node->uid, node->path ? node->path->ptr : "-");

  return node;
}

void ewap_tcp_node_free(struct ewap_tcp_node *node) {
  if (node) {
    ewap_pr_log(
        EWAP_LOG_CONNECTIONS,
        "tcp_node deleted (key: %u %llu %llu %llu %llu), pid: %d, uid: %u",
        node->key.port, node->key.saddr_hi, node->key.saddr_lo,
        node->key.daddr_hi, node->key.daddr_lo, node->pid, node->uid);

    ewap_path_unref(node->path);
    kmem_cache_free(ewap_tcp_node_allocator, node);
  }
}

void ewap_tcp_map_init(struct ewap_tcp_map *map) { map->connections = RB_ROOT; }

void ewap_tcp_map_deinit(struct ewap_tcp_map *map) {
  struct ewap_tcp_node *cursor;
  struct ewap_tcp_node *tmp;

  rbtree_postorder_for_each_entry_safe(cursor, tmp, &map->connections, node) {
    ewap_tcp_node_free(cursor);
  }
}

static struct ewap_tcp_node *tcp_map_find(struct rb_root *root,
                                          const struct ewap_tcp_key *key,
                                          struct rb_node **parent) {
  struct rb_node *p = NULL;
  struct rb_node *node = root->rb_node;
  struct ewap_tcp_node *result = NULL;

  while (node) {
    struct rb_node *child;
    struct ewap_tcp_node *entry = rb_entry(node, struct ewap_tcp_node, node);
    int cmp = ewap_tcp_key_cmp(key, &entry->key);

    if (cmp < 0) {
      child = node->rb_left;
    } else if (cmp > 0) {
      child = node->rb_right;
    } else {
      result = entry;
      break;
    }

    p = node;
    node = child;
  }

  if (parent) {
    *parent = p;
  }

  return result;
}

static void tcp_map_insert(struct rb_root *root, struct ewap_tcp_node *item,
                           struct rb_node *parent) {
  struct rb_node **link;

  if (RB_EMPTY_ROOT(root)) {
    link = &root->rb_node;
  } else {
    struct ewap_tcp_node *parent_entry;
    int cmp;

    if (unlikely(parent == NULL)) {
      BUG();
    }

    parent_entry = rb_entry(parent, struct ewap_tcp_node, node);
    cmp = ewap_tcp_key_cmp(&item->key, &parent_entry->key);

    if (cmp < 0) {
      link = &parent->rb_left;
    } else if (cmp > 0) {
      link = &parent->rb_right;
    } else {
      BUG();
    }
  }

  rb_link_node(&item->node, parent, link);
  rb_insert_color(&item->node, root);
}

bool ewap_tcp_map_add_pid(struct ewap_tcp_map *map,
                          struct ewap_tcp_node *new_node) {
  struct rb_node *parent = NULL;

  struct ewap_tcp_node *node =
      tcp_map_find(&map->connections, &new_node->key, &parent);
  if (unlikely(node)) {
    return false;
  }

  tcp_map_insert(&map->connections, new_node, parent);
  new_node->tcp_map = map;

  return true;
}

struct ewap_tcp_node *ewap_tcp_map_get_node(struct ewap_tcp_map *map,
                                            const struct ewap_tcp_key *key) {
  return tcp_map_find(&map->connections, key, NULL);
}

struct ewap_tcp_node *ewap_tcp_map_extract_node(
    struct ewap_tcp_map *map, const struct ewap_tcp_key *key) {
  struct ewap_tcp_node *node = tcp_map_find(&map->connections, key, NULL);
  if (node) {
    rb_erase(&node->node, &map->connections);
  }

  return node;
}