guestfs-inspect/guestfs-inspectd.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <czmq.h>
#include <zmsg.h>
#include <zactor.h>
#include <guestfs.h>
#include <time.h>

#include "libguestfs-inspect.h"

#define STREQ(a, b) (strcmp((a), (b)) == 0)

static void worker_task(zsock_t*, const char*);

static int compare_key_len(const void*, const void*);
static int count_mountpoints(char *const *argv);

static int compare_key_len(const void *p1, const void *p2) {
  const char *key1 = *(char* const*) p1;
  const char *key2 = *(char* const*) p2;
  return strlen(key1) - strlen(key2);
}

static void print_intro() {
  printf("\n"
"\n"
"   ,ggggggg,\n"
" ,dP\"\"\"\"\"\"Y8b         8I          8I\n"
" d8'    a  Y8         8I          8I\n"
" 88     \"Y8P\'         8I          8I\n"
" `8baaaa              8I          8I\n"
",d8P\"\"\"\"        ,gggg,8I    ,gggg,8I    ,gggg,gg\n"
"d8\"            dP\"  \"Y8I   dP\"  \"Y8I   dP\"  \"Y8I\n"
"Y8,           i8'    ,8I  i8'    ,8I  i8'    ,8I\n"
"`Yba,,_____, ,d8,   ,d8b,,d8,   ,d8b,,d8,   ,d8b,\n"
"  `\"Y8888888 P\"Y8888P\"`Y8P\"Y8888P\"`Y8P\"Y8888P\"`Y8\n"
"\n");

}

static int count_mountpoints(char *const *argv) {
  size_t c;
  for (c = 0; argv[c]; c++) {}
  return c;
}

static void cat_file(guestfs_h *g, const char *file_path, char **file_content, size_t *file_size) {
  assert((*file_content) == NULL);
  if (guestfs_is_file_opts(g, file_path, GUESTFS_IS_FILE_OPTS_FOLLOWSYMLINKS, 1, -1) > 0) {
    (*file_content) = guestfs_read_file(g, file_path, file_size);
    if ((*file_content) == NULL) {
      exit(EXIT_FAILURE);
    }
  }
  return;
}

static guestfs_h* init_guestfs(const char *disk_path) {
  guestfs_h *g = NULL;
  char **roots, **mountpoints;
  char *root;
  size_t i, j;

  // Create a connection handle
  g = guestfs_create();
  if (g == NULL) {
    exit(EXIT_FAILURE);
  }

  // Adding disk_path to connection handle
  if (guestfs_add_drive_opts(g, disk_path, GUESTFS_ADD_DRIVE_OPTS_READONLY, 1, -1) == -1) {
    exit(EXIT_FAILURE);
  }

  // Launching connection handle
  if (guestfs_launch(g) == -1) {
    exit(EXIT_FAILURE);
  }

  // Pulling rootfs information
  roots = guestfs_inspect_os(g);
  if (roots == NULL) {
    exit(EXIT_FAILURE);
  }

  // Looping through roots to mount mountpoints
  for (j = 0; roots[j] != NULL; j++) {
    root = roots[j];
    mountpoints = guestfs_inspect_get_mountpoints(g, root);
    if (mountpoints == NULL) {
      exit(EXIT_FAILURE);
    }
    // Sorting mountpoints to be in {'${device_path}', '${mount_path}'} format
    qsort(mountpoints, count_mountpoints(mountpoints) / 2, 2 * sizeof (char*), compare_key_len);
      for (i = 0; mountpoints[i] != NULL; i += 2) {
        guestfs_mount_ro(g, mountpoints[i+1], mountpoints[i]);
        free(mountpoints[i]);
        free(mountpoints[i+1]);
      }
      free(mountpoints);
      free(root);
  }
  free(roots);
  return g;
}

static void worker_task(zsock_t *pipe, const char *disk_path) {
  guestfs_h *g;
  g = init_guestfs(disk_path);

  zsock_t *worker = zsock_new_dealer("inproc://workers");

  // ZeroMQ Opens here
  zsock_signal(pipe, 0);

  while (true) {
    zmsg_t *msg = zmsg_recv(pipe);
    if (!msg) {
      break;
    }

    zframe_t *client_id = zmsg_pop(msg);
    // Checking for term message
    if (zframe_streq(client_id, "$TERM")) {
      zframe_destroy(&client_id);
      zmsg_destroy(&msg);
      break;
    }

    struct guestfs_inpsect_command *command = guestfs_inspect_zmsg_to_command(msg);
    zmsg_t *reply = zmsg_new();
    if (!reply) {
      printf("wuddahec\n");
      exit(EXIT_FAILURE);
    }

    switch (command->command) {
      case GUESTFS_COMMAND_LS:;
        zmsg_pushstr(reply, "From worker! -- not implemented");
        break;
      case GUESTFS_COMMAND_CAT:;
        char *res = NULL;
        size_t s = 0;
        cat_file(g, command->args.cat.paths[0], &res, &s);
        zmsg_addmem(reply, res, s);
        free(res);
        break;
      default:;
        zmsg_pushstr(reply, "Unknown command");
        break;
    }

    // Sending reply
    zmsg_pushstr(reply, "");
    zmsg_prepend(reply, &client_id);
    zmsg_send(&reply, worker);

    guestfs_inspect_command_destroy(&command);
    zframe_destroy(&client_id);
    zmsg_destroy(&msg);
    zmsg_destroy(&reply);
  }

  guestfs_shutdown(g);
  guestfs_close(g);
  zsock_destroy(&worker);
  exit(EXIT_SUCCESS);
}

typedef struct {
  char *name;
  zactor_t *worker;
} zactor_worker_map;

char *strdup(const char *);
void mapping_zactor_worker(const char *name, zactor_t *worker, zactor_worker_map **map, size_t mapper_size) {
  map[mapper_size - 1] = malloc(sizeof(zactor_worker_map));
  (*map[mapper_size - 1]).name = strdup(name);
  (*map[mapper_size - 1]).worker = worker;
  return;
}

zactor_t * zactor_worker_lookup(char *name, zactor_worker_map **map, size_t mapper_size) {
  size_t i = 0;
  for (i = 0; i < mapper_size; i++) {
    if (strcmp((*map[i]).name, name) == 0) {
	    return map[i]->worker;
    }
  }
  return NULL;
}

static void log_message(const char *message) {
  time_t ltime;
  struct tm *result;
  char *stime;

  ltime = time(NULL);
  result = localtime(&ltime);
  stime = asctime(result);
  strtok(stime, "\n");
  printf("[%s] %s\n", stime, message);
}

int main(int argc, char **argv) {
  if (argc < 2) {
    printf("Usage: %s <disk-path>:<name> ...\n", argv[0]);
    exit(EXIT_FAILURE);
  }

  print_intro();

  // Zactor setup
  size_t worker_map_size = 0;
  zactor_worker_map **worker_map = NULL;

  log_message("Starting Workers");

  for (int i = 0; i < argc - 1; i++) {
    worker_map_size++;
    worker_map = realloc(worker_map, sizeof(zactor_worker_map *) * worker_map_size);
    char *path = strtok(argv[i+1], ":");
    char *name = strtok(NULL, ":");
    mapping_zactor_worker(name, zactor_new((void *)worker_task, path), worker_map, worker_map_size);
  }

  if (worker_map == NULL) {
    log_message("Worker map was empty");
    exit(EXIT_FAILURE);
  }

  // Setup daemon ipc endpoint
  char *ep = guestfs_inspect_endpoint();
  // Setting up endpoint log message
  char *ep_log = (char *)malloc((strlen(ep) * sizeof(char)) + 32);

  // Sending endpoint log message
  sprintf(ep_log, "Endpoint: %s", ep); // NOLINT
  log_message(ep_log);
  free(ep_log);

  // Setup ZMQ routers
  zsock_t *frontend = zsock_new(ZMQ_ROUTER);
  zsock_bind(frontend, "%s", ep);
  free(ep);

  zsock_t *backend = zsock_new(ZMQ_ROUTER);
  zsock_bind(backend, "inproc://workers");

  // Connecting frontend to backend
  zpoller_t *poller = zpoller_new(frontend, backend, NULL);
  assert(poller);

  log_message("Server listening for requests");

  zframe_t *message_pop = NULL;

  while (true) {
    void *sock = zpoller_wait(poller, -1);

    if (!sock) {
      break;
    }

    if (sock == frontend) {
      zmsg_t *msg = zmsg_recv(frontend);
      // reply id
      zframe_t *identity = zmsg_pop(msg);
      // Null frame
      message_pop = zmsg_pop(msg);
      // Destroy immediately
      zframe_destroy(&message_pop);

      struct guestfs_inpsect_command *c = guestfs_inspect_zmsg_to_command(msg);
      if (c == NULL) {
        log_message("Error creating inspect command");
        zmsg_destroy(&msg);
        zframe_destroy(&identity);
        continue;
      }

      zactor_t *zactor_tmp = zactor_worker_lookup(c->name, worker_map, worker_map_size);
      if (zactor_tmp != NULL) {
        // send request to worker
        zmsg_prepend(msg, &identity);
        zactor_send(zactor_tmp, &msg);
      } else {
        // add failure response
      }

      zactor_tmp = NULL;
      guestfs_inspect_command_destroy(&c);
      zmsg_destroy(&msg);
      zframe_destroy(&identity);
    } else if (sock == backend) {
      zmsg_t *msg = zmsg_recv(backend);
      message_pop = zmsg_pop(msg); // Removing backend id
      // Destroy immediately
      zframe_destroy(&message_pop);
      zmsg_send(&msg, frontend);
      zmsg_destroy(&msg);
    } else {
      log_message("Socket found from unknown zsock");
    }
  }

  log_message("Cleaning up workers");
  // Cleanup
  for (size_t i = 0; i < worker_map_size; i++) {
    zactor_destroy(&worker_map[i]->worker);
    free(worker_map[i]->name);
    free(worker_map[i]);
  }

  free(worker_map);
  zpoller_destroy(&poller);
  zsock_destroy(&frontend);
  zsock_destroy(&backend);
  exit(EXIT_SUCCESS);
}