Merge PR #81 from 'pinheadmz/useragent'

src/daemon.c

@@ -36,6 +36,7 @@ typedef struct hsk_options_s {
   uint8_t *identity_key;
   char *seeds;
   int pool_size;
+  char *user_agent;
 } hsk_options_t;
 
 static void
@@ -52,6 +53,7 @@ hsk_options_init(hsk_options_t *opt) {
   opt->identity_key = NULL;
   opt->seeds = NULL;
   opt->pool_size = HSK_POOL_SIZE;
+  opt->user_agent = NULL;
 }
 
 static void
@@ -141,6 +143,9 @@ help(int r) {
     "  -l, --log-file <filename>\n"
     "    Redirect output to a log file.\n"
     "\n"
+    "  -a, --user-agent <string>\n"
+    "    Add supplemental user agent string in p2p version message.\n"
+    "\n"
 #ifndef _WIN32
     "  -d, --daemon\n"
     "    Fork and background the process.\n"
@@ -156,9 +161,9 @@ help(int r) {
 
 static void
 parse_arg(int argc, char **argv, hsk_options_t *opt) {
-  const static char *optstring = "c:n:r:i:u:p:k:s:l:h"
+  const static char *optstring = "c:n:r:i:u:p:k:s:l:h:a"
 #ifndef _WIN32
-    "d"
+    ":d"
 #endif
     ;
 
@@ -172,6 +177,7 @@ parse_arg(int argc, char **argv, hsk_options_t *opt) {
     { "identity-key", required_argument, NULL, 'k' },
     { "seeds", required_argument, NULL, 's' },
     { "log-file", required_argument, NULL, 'l' },
+    { "user-agent", required_argument, NULL, 'a' },
 #ifndef _WIN32
     { "daemon", no_argument, NULL, 'd' },
 #endif
@@ -307,6 +313,18 @@ parse_arg(int argc, char **argv, hsk_options_t *opt) {
         break;
       }
 
+      case 'a': {
+        if (!optarg || strlen(optarg) == 0)
+          return help(1);
+
+        if (opt->user_agent)
+          free(opt->user_agent);
+
+        opt->user_agent = strdup(optarg);
+
+        break;
+      }
+
 #ifndef _WIN32
       case 'd': {
         background = true;
@@ -428,6 +446,12 @@ hsk_daemon_init(hsk_daemon_t *daemon, uv_loop_t *loop, hsk_options_t *opt) {
     goto fail;
   }
 
+  if (!hsk_pool_set_agent(daemon->pool, opt->user_agent)) {
+    fprintf(stderr, "failed adding user agent\n");
+    rc = HSK_EFAILURE;
+    goto fail;
+  }
+
   daemon->ns = hsk_ns_alloc(loop, daemon->pool);
 
   if (!daemon->ns) {

src/pool.c

@@ -173,6 +173,8 @@ hsk_pool_init(hsk_pool_t *pool, const uv_loop_t *loop) {
   pool->pending_count = 0;
   pool->block_time = 0;
   pool->getheaders_time = 0;
+  pool->user_agent = (char *)malloc(255);
+  strcpy(pool->user_agent, HSK_USER_AGENT);
 
   return HSK_SUCCESS;
 }
@@ -206,6 +208,11 @@ hsk_pool_uninit(hsk_pool_t *pool) {
   hsk_chain_uninit(&pool->chain);
   hsk_addrman_uninit(&pool->am);
   hsk_timedata_uninit(&pool->td);
+
+  if (pool->user_agent) {
+    free(pool->user_agent);
+    pool->user_agent = NULL;
+  }
 }
 
 bool
@@ -288,6 +295,25 @@ hsk_pool_set_seeds(hsk_pool_t *pool, const char *seeds) {
   return true;
 }
 
+bool
+hsk_pool_set_agent(hsk_pool_t *pool, const char *user_agent) {
+  assert(pool);
+
+  if (!user_agent)
+    return true;
+
+  size_t len = strlen(pool->user_agent);  
+  len += strlen(user_agent);
+
+  // Agent size in p2p version message is 1 byte
+  if (len > 0xff)
+    return false;
+
+  pool->user_agent = strcat(pool->user_agent, user_agent);
+
+  return true;
+}
+
 hsk_pool_t *
 hsk_pool_alloc(const uv_loop_t *loop) {
   hsk_pool_t *pool = malloc(sizeof(hsk_pool_t));
@@ -1205,7 +1231,7 @@ hsk_peer_send_version(hsk_peer_t *peer) {
   hsk_addr_copy(&msg.remote.addr, &peer->addr);
 
   msg.nonce = hsk_nonce();
-  strcpy(msg.agent, HSK_USER_AGENT);
+  strcpy(msg.agent, pool->user_agent);
   msg.height = (uint32_t)pool->chain.height;
 
   peer->version_time = hsk_now();

src/pool.h

@@ -104,6 +104,7 @@ typedef struct hsk_pool_s {
   int pending_count;
   int64_t block_time;
   int64_t getheaders_time;
+  char *user_agent;
 } hsk_pool_t;
 
 /*
@@ -125,6 +126,9 @@ hsk_pool_set_size(hsk_pool_t *pool, int max_size);
 bool
 hsk_pool_set_seeds(hsk_pool_t *pool, const char *seeds);
 
+bool
+hsk_pool_set_agent(hsk_pool_t *pool, const char *user_agent);
+
 hsk_pool_t *
 hsk_pool_alloc(const uv_loop_t *loop);