First commit

8 files changed, 515 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..5b156a7
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+.*.swp
+*.o
+ref
diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..7559941
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,16 @@
+
+SRC = $(wildcard *.c)
+OBJECT = $(SRC:.c=.o)
+TARGET = nfcollect
+CC = gcc
+CFLAGS = -g -Wall -Wextra -DDEBUG
+LDFLAGS = -lnetfilter_log -lpthread
+
+$(TARGET): $(OBJECT)
+	$(CC) -o $(TARGET) $^ $(LDFLAGS)
+
+%.o: %.c
+	$(CC) $(CFLAGS) -c -o $@ $<
+
+clean:
+	rm -f $(TARGET) $(OBJECT)
diff --git a/commit.c b/commit.c
new file mode 100644
index 0000000..90cecfd
--- /dev/null
+++ b/commit.c
@@ -0,0 +1,43 @@
+#include <errno.h>
+#include <string.h>
+#include "commit.h"
+
+extern char *storage_dir;
+extern char *storage_prefix;
+const uint32_t write_blk_size = 8196;
+const uint32_t commit_file_open_flag = O_RDWR | O_CREAT | O_APPEND | O_DIRECT;
+
+void nfl_commit_init() {
+
+}
+
+void nfl_commit_worker(nflog_header_t* header, nflog_entry_t* store) {
+    FILE* f;
+    char filename[1024];
+    uint32_t id = header->id;
+
+    sprintf(filename, "%s/%s_%d", storage_dir, storage_prefix, id);
+    debug("Comm worker #%u: commit to file %s\n", header->id, filename);
+    fd = open
+    ERR((f = fopen(filename, "wb")) == NULL, strerror(errno));
+    fwrite(header, sizeof(nflog_header_t), 1, f);
+
+    uint32_t total_size = sizeof(nflog_entry_t) * header->max_n_entries;
+    uint32_t total_blk = total_size / write_blk_size;
+    uint32_t i, written = 0;
+    for(i = 0; i < total_blk; ++i) {
+        written = fwrite(store, 1, write_blk_size, f);
+
+        while(written < write_blk_size) {
+            written += fwrite(store, 1, write_blk_size - written, f);
+        }
+    }
+    
+    int remain = total_size - total_blk*write_blk_size;
+    while(remain > 0) {
+        remain -= fwrite(store, 1, remain, f);
+    }
+
+    fclose(f);
+}
+
diff --git a/commit.h b/commit.h
new file mode 100644
index 0000000..d424867
--- /dev/null
+++ b/commit.h
@@ -0,0 +1,5 @@
+#pragma once
+#include "main.h"
+
+void nfl_commit_init();
+void nfl_commit_worker(nflog_header_t* header, nflog_entry_t* store);
diff --git a/main.c b/main.c
new file mode 100644
index 0000000..0606e0f
--- /dev/null
+++ b/main.c
@@ -0,0 +1,186 @@
+
+// The MIT License (MIT)
+
+// Copyright (c) 2017 Yun-Chih Chen
+// Copyright (c) 2013 Florian Richter (nflogtable)
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in
+// all
+// copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#include "commit.h"
+#include "main.h"
+#include "nflog.h"
+#include <fcntl.h>
+#include <getopt.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+sem_t nfl_commit_queue;
+uint16_t nfl_group_id;
+char *storage_dir = NULL;
+const char *storage_prefix = "nflog_storage";
+
+const char *version_text = "nfcollect Version 0.1\n";
+const char *help_text =
+    "Usage: nfcollect [OPTION]\n"
+    "Foo bar\n"
+    "\n"
+    "Options:\n"
+    "  -h --help                 print this help\n"
+    "  -v --version              print version information\n"
+    "     --nflog-group=<id>     nflog group\n"
+    "\n";
+
+void sig_handler(int signo) {
+    if (signo == SIGHUP) {
+        /* TODO */
+    }
+}
+/*
+void *open_trunk_storage(uint32_t size){
+    int fd;
+
+    // open file and map to memory
+    ERR((fd = open(mapfile, O_RDWR | O_CREAT | O_TRUNC, (mode_t)0600)) == -1,
+                 "Could not open file");
+
+    ERR(lseek(fd, size - 1, SEEK_SET) == -1,
+                 "Could not seek to end of file");
+
+    // ERR(fallocate(fd, 0, 0, storage_size_byte) == -1,
+    //             "Could not seek to end of file");
+
+    ERR(write(fd, "", 1) == -1, "Could not write to end of file");
+
+    return mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
+}
+*/
+
+nflog_state_t *get_nflog_state(uint32_t id, uint32_t entries_max) {
+    nflog_state_t *state =
+        (nflog_state_t *)malloc(sizeof(nflog_state_t));
+    pthread_mutex_init(&(state->lock), NULL);
+    state->store = (nflog_entry_t *)malloc(sizeof(nflog_entry_t) *
+                                                  entries_max);
+    state->header.id = id;
+    state->header.max_n_entries = entries_max;
+    state->header.n_entries = 0;
+    return state;
+}
+
+void free_nflog_state(nflog_state_t **state) { *state = NULL; }
+
+int main(int argc, char *argv[]) {
+
+    uint32_t i, max_commit_worker = 0, storage_size = 0;
+    int nflog_group_id;
+
+    struct option longopts[] = {/* name, has_args, flag, val */
+                                {"nflog-group", required_argument, NULL, 'g'},
+                                {"storage_dir", required_argument, NULL, 'd'},
+                                {"storage_size", required_argument, NULL, 's'},
+                                {"help", no_argument, NULL, 'h'},
+                                {"version", no_argument, NULL, 'v'},
+                                {0, 0, 0, 0}};
+
+    int opt;
+    while ((opt = getopt_long(argc, argv, "g:d:hv", longopts, NULL)) != -1) {
+        switch (opt) {
+        case 'h':
+            printf("%s", help_text);
+            exit(0);
+            break;
+        case 'v':
+            printf("%s", version_text);
+            exit(0);
+            break;
+        case 'f':
+            storage_dir = optarg;
+            break;
+        case 'g':
+            nflog_group_id = atoi(optarg);
+            break;
+        case 's':
+            storage_size = atoi(optarg);
+            break;
+        case '?':
+            fprintf(stderr, "Unknown argument, see --help\n");
+            exit(1);
+        }
+    }
+
+    // verify arguments
+    ASSERT(nflog_group_id != -1,
+           "You must provide a nflog group (see --help)!\n");
+    ASSERT(storage_dir != NULL,
+           "You must provide a storage directory (see --help)\n");
+    ASSERT(storage_size == 0, "You must provide the desired size of log file "
+                              "(in MiB) (see --help)\n");
+    struct stat _d;
+    if(stat(storage_dir, &_d) != 0 || !S_ISDIR(_d.st_mode)){
+        fprintf(stderr, "storage directory '%s' not exist", storage_dir);
+    }
+
+    // max number of commit worker defaults to #processor - 1
+    if (max_commit_worker == 0) {
+        max_commit_worker = sysconf(_SC_NPROCESSORS_ONLN) - 1;
+        max_commit_worker = max_commit_worker > 0 ? max_commit_worker : 1;
+    }
+
+    nfl_group_id = nflog_group_id;
+
+    // register signal handler
+    ERR(signal(SIGHUP, sig_handler) == SIG_ERR, "Could not set SIGHUP handler");
+
+    uint32_t pgsize = getpagesize();
+    uint32_t trunk_size_byte = storage_size / TRUNK_SIZE * 1024 * 1024; // MiB
+    trunk_size_byte = (trunk_size_byte / pgsize) * pgsize; // align with pagesize
+
+    uint32_t trunk_cnt = CEILING(storage_size, trunk_size_byte);
+    uint32_t entries_max = (trunk_size_byte - sizeof(nflog_header_t)) /
+                           sizeof(nflog_entry_t);
+
+    // Set up commit worker
+    sem_init(&nfl_commit_queue, 0, max_commit_worker);
+
+    // Set up nflog receiver worker
+    nflog_state_t **trunks = (nflog_state_t **)malloc(
+        sizeof(nflog_state_t *) * trunk_cnt);
+    for (i = 0; i < trunk_cnt; ++i) {
+        trunks[i] = NULL;
+    }
+    
+    nfl_commit_init(trunk_cnt);
+
+    for (i = 0;; i = (i + 1) % trunk_cnt) {
+        trunks[i] =
+            trunks[i] != NULL ? trunks[i] : get_nflog_state(i, entries_max);
+        pthread_mutex_lock(&(trunks[i]->lock));
+        pthread_create(&(trunks[i]->thread), NULL, nflog_worker,
+                       (void *)trunks[i]);
+        pthread_join(trunks[i]->thread, NULL);
+    }
+    
+    sem_destroy(&nfl_commit_queue);
+    exit(0);
+}
diff --git a/main.h b/main.h
new file mode 100644
index 0000000..96bc65a
--- /dev/null
+++ b/main.h
@@ -0,0 +1,108 @@
+
+// The MIT License (MIT)
+
+// Copyright (c) 2017 Yun-Chih Chen
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in
+// all
+// copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#pragma once
+#include <semaphore.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <netinet/in.h>
+#include <netinet/ip.h>
+#include <netinet/tcp.h>
+#include <netinet/udp.h>
+
+#ifdef DEBUG
+#define DEBUG_ON 1
+#else
+#define DEBUG_ON 0
+#endif
+
+#define ASSERT(condition, error_msg)                                           \
+  if (!(condition)) {                                                          \
+    fputs((error_msg), stderr);                                                \
+    exit(1);                                                                   \
+  }
+#define ERR(command, error_msg)                                                \
+  if (command) {                                                               \
+    perror((error_msg));                                                       \
+    exit(1);                                                                   \
+  }
+#define debug(format, ...)                                                     \
+  if (DEBUG_ON) {                                                              \
+    fprintf(stdout, format, ##__VA_ARGS__);                                    \
+  }
+
+#define CEILING(a,b) ((a)%(b) == 0 ? ((a)/(b)) : ((a)/(b)+1))
+#define TRUNK_SIZE (4096 * 150)
+
+typedef struct __attribute__((packed)) _nflog_header_t {
+	uint32_t                   id;                   /*     0     4 */
+	uint32_t                   n_entries;            /*     4     4 */
+	uint32_t                   max_n_entries;        /*     8     4 */
+	uint32_t                   _unused;              /*    12     4 */
+	uint64_t                   start_time;           /*    16     8 */
+	uint64_t                   end_time;             /*    24     8 */
+
+	/* size: 32, cachelines: 1, members: 6 */
+} nflog_header_t;
+
+
+typedef struct __attribute__((packed)) _nflog_entry_t {
+	// current timestamp since UNIX epoch
+	time_t                     timestamp;            /*     0     8 */
+
+ 	// dest address
+	struct in_addr             daddr;                /*     8     4 */
+
+	// uid
+	uint32_t                   uid;                  /*    12     4 */
+
+	// unused space, just for padding
+	uint8_t                    __unused1;            /*    16     1 */
+
+ 	// IP protocol (UDP or TCP)
+	uint8_t                    protocol;             /*    17     1 */
+
+	// unused space, just for padding
+	uint16_t                   __unused2;            /*    18     2 */
+
+	// source port
+	uint16_t                   sport;                /*    20     2 */
+
+	// destination port
+	uint16_t                   dport;                /*    22     2 */
+
+	/* size: 24, cachelines: 1, members: 8 */
+} nflog_entry_t;
+
+
+typedef struct _nflog_state_t {
+    nflog_header_t header;
+    nflog_entry_t* store;
+
+	struct nflog_handle *nfl_fd;
+	struct nflog_g_handle *nfl_group_fd;
+
+	pthread_mutex_t lock;
+	pthread_t thread;
+} nflog_state_t;
diff --git a/nflog.c b/nflog.c
new file mode 100644
index 0000000..bd32b83
--- /dev/null
+++ b/nflog.c
@@ -0,0 +1,151 @@
+
+// The MIT License (MIT)
+
+// Copyright (c) 2017 Yun-Chih Chen
+
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+
+// The above copyright notice and this permission notice shall be included in
+// all
+// copies or substantial portions of the Software.
+
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#include "commit.h"
+#include "nflog.h"
+#include "main.h"
+#include <stddef.h>    // size_t for libnetfilter_log
+#include <sys/types.h> // u_int32_t for libnetfilter_log
+#include <libnetfilter_log/libnetfilter_log.h>
+#include <pthread.h>
+#include <time.h>
+
+extern sem_t nfl_commit_queue;
+extern uint16_t nfl_group_id;
+
+static void nfl_cleanup(nflog_state_t *nf);
+static void nfl_init(nflog_state_t *nf);
+static void *_nfl_commit_worker(void *targs);
+static void nfl_commit(nflog_state_t *nf);
+
+static int handle_packet(struct nflog_g_handle *gh, struct nfgenmsg *nfmsg,
+                         struct nflog_data *nfa, void *_nf) {
+    char *payload;
+    int payload_len = nflog_get_payload(nfa, &payload);
+    nflog_state_t *nf = (nflog_state_t *)_nf;
+
+    // only process ipv4 packet
+    if (payload_len >= 0 && ((payload[0] & 0xf0) == 0x40)) {
+        struct iphdr *iph = (struct iphdr *)payload;
+        nflog_entry_t *entry = &(nf->store[nf->header.n_entries]);
+
+        void *inner_hdr = iph + iph->ihl;
+        // Only accept TCP / UDP packets
+        if (iph->protocol == IPPROTO_TCP) {
+            struct tcphdr *tcph = (struct tcphdr *)inner_hdr;
+            entry->sport = ntohs(tcph->source);
+            entry->dport = ntohs(tcph->dest);
+        } else if (iph->protocol == IPPROTO_UDP) {
+            struct udphdr *tcph = (struct udphdr *)inner_hdr;
+            entry->sport = ntohs(tcph->source);
+            entry->dport = ntohs(tcph->dest);
+        } else
+            return 1; // Ignore other types of packet
+
+        entry->daddr.s_addr = iph->daddr;
+        entry->protocol = iph->protocol;
+
+        // get sender uid
+        uint32_t uid;
+        if (nflog_get_uid(nfa, &uid) == 0)
+            entry->uid = uid;
+        else
+            entry->uid = (uint32_t)~0;
+
+        // get current timestamp
+        time(&entry->timestamp);
+        nf->header.n_entries++;
+    }
+
+    // Ignore IPv6 packet for now Q_Q
+    return 0;
+}
+
+static void nfl_init(nflog_state_t *nf) {
+    // open nflog
+    ERR((nf->nfl_fd = nflog_open()) == NULL, "error during nflog_open()")
+
+    // monitor IPv4 packets only
+    ERR(nflog_bind_pf(nf->nfl_fd, AF_INET) < 0, "error during nflog_bind_pf()");
+
+    // bind to group
+    nf->nfl_group_fd = nflog_bind_group(nf->nfl_fd, nfl_group_id);
+
+    // only copy size of ipv4 header + tcp/udp src/dest port (first 4 bytes of their headers
+    ERR(nflog_set_mode(nf->nfl_group_fd, NFULNL_COPY_PACKET, sizeof(struct iphdr) + 4) < 0,
+        "Could not set copy mode");
+
+    nflog_callback_register(nf->nfl_group_fd, &handle_packet, NULL);
+}
+
+static void nfl_cleanup(nflog_state_t *nf) {
+    nflog_unbind_group(nf->nfl_group_fd);
+    nflog_close(nf->nfl_fd);
+}
+
+void *nflog_worker(void *targs) {
+    nflog_state_t *nf = (nflog_state_t *)targs;
+    nfl_init(nf);
+
+    int fd = nflog_fd(nf->nfl_fd);
+    uint32_t *p_cnt_now = &(nf->header.n_entries);
+    uint32_t cnt_max = nf->header.max_n_entries;
+
+    debug("Recv worker #%u: main loop starts\n", nf->header.id);
+    nf->header.start_time = time(NULL);
+    while (*p_cnt_now < cnt_max) {
+        int rv; char buf[4096];
+        if ((rv = recv(fd, buf, sizeof(buf), 0)) && rv >= 0) {
+            debug("Recv worker #%u: nflog packet received (len=%u)\n", nf->header.id,
+                  rv);
+            nflog_handle_packet(nf->nfl_fd, buf, rv);
+        }
+    }
+
+    nf->header.end_time = time(NULL);
+    nfl_cleanup(nf);
+    nfl_commit(nf);
+
+    /* TODO: can return exit status */
+    pthread_exit(NULL);
+}
+
+void nfl_commit(nflog_state_t *nf) {
+    pthread_t tid;
+    pthread_create(&tid, NULL, _nfl_commit_worker, (void *)nf);
+    pthread_detach(tid);
+}
+
+void *_nfl_commit_worker(void *targs) {
+    nflog_state_t* nf = (nflog_state_t*) targs;
+    debug("Comm worker #%u: thread started\n", nf->header.id);
+
+    sem_wait(&nfl_commit_queue);
+    debug("Comm worker #%u: commit started\n", nf->header.id);
+    nfl_commit_worker(&(nf->header), nf->store);
+    debug("Comm worker #%u: commit done\n", nf->header.id);
+    sem_post(&nfl_commit_queue);
+
+    pthread_mutex_unlock(&(nf->lock));
+}
diff --git a/nflog.h b/nflog.h
new file mode 100644
index 0000000..a677cc0
--- /dev/null
+++ b/nflog.h
@@ -0,0 +1,3 @@
+#pragma once
+
+void* nflog_worker(void *targs);