Hermit/src/discover.c

#ifndef _DEFAULT_SOURCE
#define _DEFAULT_SOURCE
#endif
#include "cst_discover.h"
#include "cJSON.h"
#include "cst_realm.h"
#include "cst_util.h"
#include "os_compat.h"

#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#ifndef _WIN32
#include <sys/file.h>
#endif

static pthread_mutex_t local_mu = PTHREAD_MUTEX_INITIALIZER;

#ifdef _WIN32
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#include <winsock2.h>
#include <ws2tcpip.h>
typedef SOCKET disc_sock_t;
#define DISC_BAD INVALID_SOCKET
static void disc_close(disc_sock_t s) { closesocket(s); }
#else
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
typedef int disc_sock_t;
#define DISC_BAD (-1)
static void disc_close(disc_sock_t s) { close(s); }
#endif

#define PEER_CAP 64
#define BEACON_MAX 512

struct cst_discover {
  char realm_fp[33];
  char instance_id[64];
  char name[64];
  int peer_port;
  volatile int stop;
  int thread_on;
  pthread_t thread;
  pthread_mutex_t mu;
  cst_peer_info_t peers[PEER_CAP];
  int npeers;
  disc_sock_t sock;
};

static int local_path(char *out, size_t n) {
  char dir[OS_PATH_MAX];
  if (cst_realm_dir(dir, sizeof(dir)) != 0)
    return -1;
  if (os_mkdir_p(dir, 0700) != 0)
    return -1;
  return os_path_join(out, n, dir, "local-instances.json");
}

static cJSON *load_local(void) {
  char path[OS_PATH_MAX];
  if (local_path(path, sizeof(path)) != 0)
    return cJSON_CreateObject();
  FILE *f = fopen(path, "rb");
  if (!f)
    return cJSON_CreateObject();
  cst_buf_t b = {0};
  char chunk[4096];
  size_t n;
  while ((n = fread(chunk, 1, sizeof(chunk), f)) && b.len < 256 * 1024)
    cst_buf_append(&b, chunk, n);
  fclose(f);
  cJSON *o = cJSON_Parse(b.buf ? b.buf : "{}");
  cst_buf_free(&b);
  if (!o)
    o = cJSON_CreateObject();
  if (!cJSON_GetObjectItemCaseSensitive(o, "instances"))
    cJSON_AddArrayToObject(o, "instances");
  return o;
}

static int save_local(cJSON *o) {
  char path[OS_PATH_MAX], tmp[OS_PATH_MAX + 8];
  if (local_path(path, sizeof(path)) != 0)
    return -1;
  if (snprintf(tmp, sizeof(tmp), "%s.tmp", path) >= (int)sizeof(tmp))
    return -1;
  char *txt = cJSON_PrintUnformatted(o);
  if (!txt)
    return -1;
  FILE *f = fopen(tmp, "wb");
  int ok = 0;
  if (f) {
    ok = fputs(txt, f) >= 0 && fputc('\n', f) >= 0;
    if (fclose(f) != 0)
      ok = 0;
  }
  free(txt);
  if (!ok) {
    remove(tmp);
    return -1;
  }
  if (os_rename_replace(tmp, path) != 0) {
    remove(tmp);
    return -1;
  }
  return 0;
}

static void local_remove_iid(cJSON *arr, const char *instance_id) {
  if (!cJSON_IsArray(arr) || !instance_id)
    return;
  for (int i = cJSON_GetArraySize(arr) - 1; i >= 0; i--) {
    cJSON *it = cJSON_GetArrayItem(arr, i);
    const cJSON *id = cJSON_GetObjectItemCaseSensitive(it, "iid");
    if (cJSON_IsString(id) && !strcmp(id->valuestring, instance_id))
      cJSON_DeleteItemFromArray(arr, i);
  }
}

int cst_discover_local_put(const char *realm_fp, const char *instance_id,
                           const char *name, int peer_port) {
  if (!realm_fp || !instance_id || peer_port <= 0)
    return -1;
  pthread_mutex_lock(&local_mu);
  cJSON *root = load_local();
  cJSON *arr = cJSON_GetObjectItemCaseSensitive(root, "instances");
  if (!cJSON_IsArray(arr)) {
    arr = cJSON_CreateArray();
    cJSON_AddItemToObject(root, "instances", arr);
  }
  local_remove_iid(arr, instance_id);
  cJSON *row = cJSON_CreateObject();
  cJSON_AddStringToObject(row, "rid", realm_fp);
  cJSON_AddStringToObject(row, "iid", instance_id);
  cJSON_AddStringToObject(row, "name", name ? name : "");
  cJSON_AddNumberToObject(row, "port", peer_port);
  cJSON_AddNumberToObject(row, "pid", (double)os_getpid());
  cJSON_AddNumberToObject(row, "updated", (double)time(NULL));
  cJSON_AddItemToArray(arr, row);
  int rc = save_local(root);
  cJSON_Delete(root);
  pthread_mutex_unlock(&local_mu);
  return rc;
}

int cst_discover_local_del(const char *instance_id) {
  if (!instance_id)
    return -1;
  pthread_mutex_lock(&local_mu);
  cJSON *root = load_local();
  cJSON *arr = cJSON_GetObjectItemCaseSensitive(root, "instances");
  local_remove_iid(arr, instance_id);
  int rc = save_local(root);
  cJSON_Delete(root);
  pthread_mutex_unlock(&local_mu);
  return rc;
}

int cst_discover_local_list(const char *realm_fp, cst_peer_info_t *out, int max) {
  if (!out || max <= 0)
    return 0;
  pthread_mutex_lock(&local_mu);
  cJSON *root = load_local();
  cJSON *arr = cJSON_GetObjectItemCaseSensitive(root, "instances");
  int n = 0;
  const cJSON *it;
  time_t now = time(NULL);
  cJSON_ArrayForEach(it, arr) {
    if (n >= max)
      break;
    const cJSON *rid = cJSON_GetObjectItemCaseSensitive(it, "rid");
    const cJSON *iid = cJSON_GetObjectItemCaseSensitive(it, "iid");
    const cJSON *name = cJSON_GetObjectItemCaseSensitive(it, "name");
    const cJSON *port = cJSON_GetObjectItemCaseSensitive(it, "port");
    const cJSON *updated = cJSON_GetObjectItemCaseSensitive(it, "updated");
    if (!cJSON_IsString(rid) || !cJSON_IsString(iid) || !cJSON_IsNumber(port))
      continue;
    if (realm_fp && strcmp(rid->valuestring, realm_fp) != 0)
      continue;
    if (cJSON_IsNumber(updated) && now - (time_t)updated->valuedouble > 120)
      continue;
    memset(&out[n], 0, sizeof(out[n]));
    snprintf(out[n].realm_fp, sizeof(out[n].realm_fp), "%s", rid->valuestring);
    snprintf(out[n].instance_id, sizeof(out[n].instance_id), "%s",
             iid->valuestring);
    if (cJSON_IsString(name))
      snprintf(out[n].name, sizeof(out[n].name), "%s", name->valuestring);
    snprintf(out[n].host, sizeof(out[n].host), "%s", "127.0.0.1");
    out[n].port = port->valueint;
    n++;
  }
  cJSON_Delete(root);
  pthread_mutex_unlock(&local_mu);
  return n;
}

static void remember(cst_discover_t *d, const cst_peer_info_t *p) {
  if (!strcmp(p->instance_id, d->instance_id))
    return;
  if (strcmp(p->realm_fp, d->realm_fp) != 0)
    return;
  pthread_mutex_lock(&d->mu);
  for (int i = 0; i < d->npeers; i++) {
    if (!strcmp(d->peers[i].instance_id, p->instance_id)) {
      d->peers[i] = *p;
      pthread_mutex_unlock(&d->mu);
      return;
    }
  }
  if (d->npeers < PEER_CAP)
    d->peers[d->npeers++] = *p;
  pthread_mutex_unlock(&d->mu);
}

static void ingest_local(cst_discover_t *d) {
  cst_peer_info_t list[PEER_CAP];
  int n = cst_discover_local_list(d->realm_fp, list, PEER_CAP);
  for (int i = 0; i < n; i++)
    remember(d, &list[i]);
}

static int send_beacon(cst_discover_t *d) {
  if (d->sock == DISC_BAD)
    return -1;
  char body[BEACON_MAX];
  int n = snprintf(body, sizeof(body),
                   "{\"v\":1,\"rid\":\"%s\",\"iid\":\"%s\",\"name\":\"%s\","
                   "\"port\":%d}",
                   d->realm_fp, d->instance_id, d->name, d->peer_port);
  if (n <= 0 || n >= (int)sizeof(body))
    return -1;
  struct sockaddr_in dst;
  memset(&dst, 0, sizeof(dst));
  dst.sin_family = AF_INET;
  dst.sin_port = htons(CST_DISCOVER_PORT);
  inet_pton(AF_INET, CST_DISCOVER_GROUP, &dst.sin_addr);
  return sendto(d->sock, body, (size_t)n, 0, (struct sockaddr *)&dst,
                sizeof(dst)) < 0
             ? -1
             : 0;
}

static void handle_datagram(cst_discover_t *d, const char *buf, size_t n,
                            const struct sockaddr_in *from) {
  char tmp[BEACON_MAX];
  if (n >= sizeof(tmp))
    n = sizeof(tmp) - 1;
  memcpy(tmp, buf, n);
  tmp[n] = 0;
  cJSON *o = cJSON_Parse(tmp);
  if (!o)
    return;
  const cJSON *rid = cJSON_GetObjectItemCaseSensitive(o, "rid");
  const cJSON *iid = cJSON_GetObjectItemCaseSensitive(o, "iid");
  const cJSON *name = cJSON_GetObjectItemCaseSensitive(o, "name");
  const cJSON *port = cJSON_GetObjectItemCaseSensitive(o, "port");
  if (cJSON_IsString(rid) && cJSON_IsString(iid) && cJSON_IsNumber(port)) {
    cst_peer_info_t p;
    memset(&p, 0, sizeof(p));
    snprintf(p.realm_fp, sizeof(p.realm_fp), "%s", rid->valuestring);
    snprintf(p.instance_id, sizeof(p.instance_id), "%s", iid->valuestring);
    if (cJSON_IsString(name))
      snprintf(p.name, sizeof(p.name), "%s", name->valuestring);
    p.port = port->valueint;
    inet_ntop(AF_INET, &from->sin_addr, p.host, sizeof(p.host));
    remember(d, &p);
  }
  cJSON_Delete(o);
}

static void *discover_main(void *arg) {
  cst_discover_t *d = arg;
  int ticks = 0;
  while (!d->stop) {
    if ((ticks++ % 10) == 0) {
      send_beacon(d);
      (void)cst_discover_local_put(d->realm_fp, d->instance_id, d->name,
                                   d->peer_port);
      ingest_local(d);
    }
    if (d->sock == DISC_BAD) {
      struct timespec ts = {0, 100 * 1000 * 1000};
      nanosleep(&ts, NULL);
      continue;
    }
    struct sockaddr_in from;
    socklen_t flen = sizeof(from);
    char buf[BEACON_MAX];
#ifdef _WIN32
    DWORD tv = 100;
    setsockopt(d->sock, SOL_SOCKET, SO_RCVTIMEO, (const char *)&tv, sizeof(tv));
#else
    struct timeval tv = {0, 100000};
    setsockopt(d->sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
#endif
    int got = (int)recvfrom(d->sock, buf, sizeof(buf) - 1, 0,
                            (struct sockaddr *)&from, &flen);
    if (got > 0)
      handle_datagram(d, buf, (size_t)got, &from);
  }
  (void)cst_discover_local_del(d->instance_id);
  return NULL;
}

static disc_sock_t open_multicast(void) {
#ifdef _WIN32
  os_net_init();
#endif
  disc_sock_t s = socket(AF_INET, SOCK_DGRAM, 0);
  if (s == DISC_BAD)
    return DISC_BAD;
  int yes = 1;
  setsockopt(s, SOL_SOCKET, SO_REUSEADDR, (const void *)&yes, sizeof(yes));
#ifdef SO_REUSEPORT
  setsockopt(s, SOL_SOCKET, SO_REUSEPORT, &yes, sizeof(yes));
#endif
  struct sockaddr_in addr;
  memset(&addr, 0, sizeof(addr));
  addr.sin_family = AF_INET;
  addr.sin_port = htons(CST_DISCOVER_PORT);
  addr.sin_addr.s_addr = htonl(INADDR_ANY);
  if (bind(s, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
    disc_close(s);
    return DISC_BAD;
  }
  struct ip_mreq mreq;
  memset(&mreq, 0, sizeof(mreq));
  inet_pton(AF_INET, CST_DISCOVER_GROUP, &mreq.imr_multiaddr);
  mreq.imr_interface.s_addr = htonl(INADDR_ANY);
  if (setsockopt(s, IPPROTO_IP, IP_ADD_MEMBERSHIP, (const void *)&mreq,
                 sizeof(mreq)) != 0) {
    /* Multicast may be unavailable in some sandboxes; local file still works. */
  }
  unsigned char ttl = 1;
  setsockopt(s, IPPROTO_IP, IP_MULTICAST_TTL, (const void *)&ttl, sizeof(ttl));
  return s;
}

cst_discover_t *cst_discover_start(const char *realm_fp, const char *instance_id,
                                   const char *name, int peer_port) {
  if (!realm_fp || !instance_id || peer_port <= 0)
    return NULL;
  cst_discover_t *d = calloc(1, sizeof(*d));
  if (!d)
    return NULL;
  snprintf(d->realm_fp, sizeof(d->realm_fp), "%s", realm_fp);
  snprintf(d->instance_id, sizeof(d->instance_id), "%s", instance_id);
  snprintf(d->name, sizeof(d->name), "%s", name ? name : "");
  d->peer_port = peer_port;
  d->sock = DISC_BAD;
  pthread_mutex_init(&d->mu, NULL);
  d->sock = open_multicast();
  if (cst_discover_local_put(realm_fp, instance_id, name, peer_port) != 0) {
    cst_discover_stop(d);
    return NULL;
  }
  if (pthread_create(&d->thread, NULL, discover_main, d) != 0) {
    cst_discover_stop(d);
    return NULL;
  }
  d->thread_on = 1;
  return d;
}

void cst_discover_stop(cst_discover_t *d) {
  if (!d)
    return;
  d->stop = 1;
  if (d->thread_on)
    pthread_join(d->thread, NULL);
  d->thread_on = 0;
  if (d->sock != DISC_BAD)
    disc_close(d->sock);
  (void)cst_discover_local_del(d->instance_id);
  pthread_mutex_destroy(&d->mu);
  free(d);
}

int cst_discover_peers(cst_discover_t *d, cst_peer_info_t *out, int max) {
  if (!d || !out || max <= 0)
    return 0;
  ingest_local(d);
  pthread_mutex_lock(&d->mu);
  int n = d->npeers < max ? d->npeers : max;
  memcpy(out, d->peers, (size_t)n * sizeof(cst_peer_info_t));
  pthread_mutex_unlock(&d->mu);
  return n;
}