From 6f579386f6ecc7ddf0af86577eb243f243dcb5f8 Mon Sep 17 00:00:00 2001
From: Quentin Dufour <quentin@deuxfleurs.fr>
Date: Mon, 20 Jan 2020 16:35:52 +0100
Subject: [PATCH] Add a stream repair function

---
 src/algo_dup2.c      |  7 ++++++-
 src/algo_lightning.c |  8 +++++++-
 src/donar_client.c   | 45 +++++++++++++++++++++++++++++++++++++++++++-
 src/donar_server.c   |  8 +++++++-
 src/proxy.h          |  3 +++
 5 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/src/algo_dup2.c b/src/algo_dup2.c
index 6045c8d..b1631fe 100644
--- a/src/algo_dup2.c
+++ b/src/algo_dup2.c
@@ -114,6 +114,11 @@ int algo_dup2_on_datagram(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdin
 }
 
 int algo_dup2_on_err(struct evt_core_ctx *ctx, struct evt_core_fdinfo *fdinfo) {
-  // We do nothing
+  struct algo_ctx* app_ctx = fdinfo->cat->app_ctx;
+  if (strcmp("tcp-read", fdinfo->cat->name) == 0 || strcmp("tcp-write", fdinfo->cat->name) == 0)
+    return app_ctx->ap.sr(ctx, fdinfo);
+
+  fprintf(stderr, "%s is not eligible for a reconnect\n", fdinfo->url);
+ // We do nothing
   return 1;
 }
diff --git a/src/algo_lightning.c b/src/algo_lightning.c
index 657b7b6..e87146a 100644
--- a/src/algo_lightning.c
+++ b/src/algo_lightning.c
@@ -562,5 +562,11 @@ int algo_lightning_on_datagram(struct evt_core_ctx* ctx, struct evt_core_fdinfo*
 }
 
 int algo_lightning_on_err(struct evt_core_ctx *ctx, struct evt_core_fdinfo *fdinfo) {
-  return 0;
+  struct algo_ctx* app_ctx = fdinfo->cat->app_ctx;
+  if (strcmp("tcp-read", fdinfo->cat->name) == 0 || strcmp("tcp-write", fdinfo->cat->name) == 0)
+    return app_ctx->ap.sr(ctx, fdinfo);
+
+  fprintf(stderr, "%s is not eligible for a reconnect\n", fdinfo->url);
+ // We do nothing
+  return 1;
 }
diff --git a/src/donar_client.c b/src/donar_client.c
index a175467..23ec18b 100644
--- a/src/donar_client.c
+++ b/src/donar_client.c
@@ -69,6 +69,48 @@ void init_socks5_sinks(struct donar_client_ctx* app_ctx) {
   evt_core_add_cat(&app_ctx->evts, &template);
 }
 
+int donar_client_stream_repair(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
+  // @FIXME: Ugly way to get donar_client_ctx. Shame on me :/
+  fprintf(stdout, "[donar-client] %s broke\n", fdinfo->url);
+  struct evt_core_cat* cat = evt_core_get_from_cat (ctx, "socks5-failed");
+  if (cat == NULL) {
+    fprintf(stderr, "Unable to reconnect stream as socks5-failed cat is not available...\n");
+    exit(EXIT_FAILURE);
+  }
+  struct donar_client_ctx* app_ctx = cat->app_ctx;
+
+  struct evt_core_fdinfo* fdtarget = NULL;
+  int port = url_get_port_int (fdinfo->url);
+  int pos = port - 7500, removed = 0;
+  char buffer[256];
+
+  sprintf(buffer, "tcp:read:127.0.0.1:%d", port);
+  fdtarget = evt_core_get_from_url (ctx, buffer);
+  if (fdtarget != NULL) {
+    evt_core_rm_fd(ctx, fdtarget->fd);
+    removed++;
+  }
+
+  sprintf(buffer, "tcp:write:127.0.0.1:%d", port);
+  fdtarget = evt_core_get_from_url (ctx, buffer);
+  if (fdtarget != NULL) {
+    evt_core_rm_fd(ctx, fdtarget->fd);
+    removed++;
+  }
+
+  if (removed == 2) {
+    fprintf(stdout, "[donar-client] Retriggering socks5 for port %d\n", port);
+    init_socks5_client (app_ctx, pos);
+    return 0;
+  } else if (removed == 0) {
+    fprintf(stdout, "[donar-client] Socks5 has already been retriggered for port %d\n", port);
+    return 0;
+  } else {
+    fprintf(stderr, "[donar-client] We only removed 1 link and not 2 for port %d, strange behaviour, exiting...\n", port);
+    exit(EXIT_FAILURE);
+  }
+}
+
 void donar_client(struct donar_client_ctx* ctx, struct donar_params* dp) {
   struct algo_params ap = {
     .is_waiting_bootstrap = dp->is_waiting_bootstrap,
@@ -77,7 +119,8 @@ void donar_client(struct donar_client_ctx* ctx, struct donar_params* dp) {
     .links = dp->links,
     .fresh_data = dp->fresh_data,
     .redundant_data = dp->redundant_data,
-    .capture_file = dp->capture_file
+    .capture_file = dp->capture_file,
+    .sr = donar_client_stream_repair
   };
 
   evt_core_init (&(ctx->evts), dp->verbose);
diff --git a/src/donar_server.c b/src/donar_server.c
index bfb1729..6839aa4 100644
--- a/src/donar_server.c
+++ b/src/donar_server.c
@@ -51,6 +51,11 @@ socket_create_err:
   exit(EXIT_FAILURE);
 }
 
+int donar_server_stream_repair(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fd) {
+  fprintf(stderr, "I am a server, I do nothing with broken streams...\n");
+  return 1;
+}
+
 void donar_server(struct donar_server_ctx* ctx,  struct donar_params* dp) {
   struct algo_params ap = {
     .is_waiting_bootstrap = dp->is_waiting_bootstrap,
@@ -59,7 +64,8 @@ void donar_server(struct donar_server_ctx* ctx,  struct donar_params* dp) {
     .links = dp->links,
     .fresh_data = dp->fresh_data,
     .redundant_data = dp->redundant_data,
-    .capture_file = dp->capture_file
+    .capture_file = dp->capture_file,
+    .sr = donar_server_stream_repair
   };
 
   evt_core_init (&(ctx->evts), dp->verbose);
diff --git a/src/proxy.h b/src/proxy.h
index 0217535..4a0fd1d 100644
--- a/src/proxy.h
+++ b/src/proxy.h
@@ -11,10 +11,13 @@
 #include "utils.h"
 #include "packet.h"
 
+typedef int (*stream_repair)(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo);
+
 struct algo_params {
   uint8_t is_waiting_bootstrap;
   char *algo_name, *capture_file, *algo_specific_params;
   int links, fresh_data, redundant_data;
+  stream_repair sr;
 };
 
 struct algo_ctx;