uclient-fetch: add support for --header cmdline argument

Add --header='Header: value' option to allow adding custom HTTP headers
to a request, compatible with wget. The option can be used multiple times
to add multiple headers and also allows overriding the Content-Type header
for POST requests.

Signed-off-by: Daniel Golle <daniel@makrotopia.org>

tests/cram/test-san_uclient-fetch.t

@@ -12,6 +12,7 @@ check uclient-fetch usage:
   \t-P <dir>\t\t\tSet directory for output files (esc)
   \t--quiet | -q\t\t\tTurn off status messages (esc)
   \t--continue | -c\t\t\tContinue a partially-downloaded file (esc)
+  \t--header='Header: value'\tAdd HTTP header. Multiple allowed (esc)
   \t--user=<user>\t\t\tHTTP authentication username (esc)
   \t--password=<password>\t\tHTTP authentication password (esc)
   \t--user-agent | -U <str>\t\tSet HTTP user agent (esc)

tests/cram/test_uclient-fetch.t

@@ -12,6 +12,7 @@ check uclient-fetch usage:
   \t-P <dir>\t\t\tSet directory for output files (esc)
   \t--quiet | -q\t\t\tTurn off status messages (esc)
   \t--continue | -c\t\t\tContinue a partially-downloaded file (esc)
+  \t--header='Header: value'\tAdd HTTP header. Multiple allowed (esc)
   \t--user=<user>\t\t\tHTTP authentication username (esc)
   \t--password=<password>\t\tHTTP authentication password (esc)
   \t--user-agent | -U <str>\t\tSet HTTP user agent (esc)

uclient-fetch.c

@@ -21,6 +21,7 @@
 #include <sys/socket.h>
 #include <unistd.h>
 #include <stdio.h>
+#include <ctype.h>
 #include <getopt.h>
 #include <fcntl.h>
 #include <glob.h>
@@ -29,6 +30,7 @@
 #include <signal.h>
 
 #include <libubox/blobmsg.h>
+#include <libubox/list.h>
 
 #include "progress.h"
 #include "uclient.h"
@@ -38,6 +40,12 @@
 #define strdupa(x) strcpy(alloca(strlen(x)+1),x)
 #endif
 
+struct header {
+	struct list_head list;
+	char *name;
+	char *value;
+};
+
 static const char *user_agent = "uclient-fetch";
 static const char *post_data;
 static const char *post_file;
@@ -59,6 +67,7 @@ static char **urls;
 static int n_urls;
 static int timeout;
 static bool resume, cur_resume;
+static LIST_HEAD(headers);
 
 static struct progress pmt;
 static struct uloop_timeout pmt_timer;
@@ -317,6 +326,8 @@ static void check_resume_offset(struct uclient *cl)
 
 static int init_request(struct uclient *cl)
 {
+	struct header *h;
+	char *content_type = "application/x-www-form-urlencoded";
 	int rc;
 
 	out_offset = 0;
@@ -338,18 +349,33 @@ static int init_request(struct uclient *cl)
 		return rc;
 
 	uclient_http_reset_headers(cl);
+
+	list_for_each_entry(h, &headers, list) {
+		if (!strcasecmp(h->name, "Content-Type")) {
+			if (!post_data && !post_file)
+				return -EINVAL;
+
+			content_type = h->value;
+		} else if (!strcasecmp(h->name, "User-Agent")) {
+			user_agent = h->value;
+		} else {
+			uclient_http_set_header(cl, h->name, h->value);
+		}
+	}
+
 	uclient_http_set_header(cl, "User-Agent", user_agent);
+
 	if (cur_resume)
 		check_resume_offset(cl);
 
 	if (post_data) {
-		uclient_http_set_header(cl, "Content-Type", "application/x-www-form-urlencoded");
+		uclient_http_set_header(cl, "Content-Type", content_type);
 		uclient_write(cl, post_data, strlen(post_data));
 	}
 	else if(post_file)
 	{
 		FILE *input_file;
-		uclient_http_set_header(cl, "Content-Type", "application/x-www-form-urlencoded");
+		uclient_http_set_header(cl, "Content-Type", content_type);
 
 		input_file = fopen(post_file, "r");
 		if (!input_file)
@@ -494,6 +520,7 @@ static void usage(const char *progname)
 		"	-P <dir>			Set directory for output files\n"
 		"	--quiet | -q			Turn off status messages\n"
 		"	--continue | -c			Continue a partially-downloaded file\n"
+		"	--header='Header: value'	Add HTTP header. Multiple allowed\n"
 		"	--user=<user>			HTTP authentication username\n"
 		"	--password=<password>		HTTP authentication password\n"
 		"	--user-agent | -U <str>		Set HTTP user agent\n"
@@ -539,6 +566,23 @@ static void debug_cb(void *priv, int level, const char *msg)
 	fprintf(stderr, "%s\n", msg);
 }
 
+static bool is_valid_header(char *str)
+{
+	char *tmp = str;
+
+	/* First character must be a letter */
+	if (!isalpha(*tmp))
+		return false;
+
+	/* Subsequent characters must be letters, numbers or dashes */
+	while (*(++tmp) != '\0') {
+		if (!isalnum(*tmp) && *tmp != '-')
+			return false;
+	}
+
+	return true;
+};
+
 enum {
 	L_NO_CHECK_CERTIFICATE,
 	L_CA_CERTIFICATE,
@@ -555,6 +599,7 @@ enum {
 	L_NO_PROXY,
 	L_QUIET,
 	L_VERBOSE,
+	L_HEADER,
 };
 
 static const struct option longopts[] = {
@@ -573,6 +618,7 @@ static const struct option longopts[] = {
 	[L_NO_PROXY] = { "no-proxy", no_argument, NULL, 0 },
 	[L_QUIET] = { "quiet", no_argument, NULL, 0 },
 	[L_VERBOSE] = { "verbose", no_argument, NULL, 0 },
+	[L_HEADER] = { "header", required_argument, NULL, 0 },
 	{}
 };
 
@@ -587,6 +633,8 @@ int main(int argc, char **argv)
 	struct uclient *cl = NULL;
 	int longopt_idx = 0;
 	bool has_cert = false;
+	struct header *h, *th;
+	char *tmp;
 	int i, ch;
 	int rc;
 	int af = -1;
@@ -661,6 +709,30 @@ int main(int argc, char **argv)
 			case L_VERBOSE:
 				debug_level++;
 				break;
+			case L_HEADER:
+				tmp = strchr(optarg, ':');
+				if (!tmp) {
+					usage(progname);
+					goto err_out;
+				}
+				*(tmp++) = '\0';
+				while (isspace(*tmp))
+					++tmp;
+
+				if (*tmp == '\0' || !is_valid_header(optarg) || strchr(tmp, '\n')) {
+					usage(progname);
+					goto err_out;
+				}
+				h = malloc(sizeof(*h));
+				if (!h) {
+					perror("Set HTTP header");
+					error_ret = 1;
+					goto err_out;
+				}
+				h->name = optarg;
+				h->value = tmp;
+				list_add_tail(&h->list, &headers);
+				break;
 			default:
 				usage(progname);
 				goto err_out;
@@ -796,5 +868,10 @@ int main(int argc, char **argv)
 	if (ssl_ctx)
 		ssl_ops->context_free(ssl_ctx);
 
+	list_for_each_entry_safe(h, th, &headers, list) {
+		list_del(&h->list);
+		free(h);
+	}
+
 	return error_ret;
 }