Merge pull request #955 from ZyX-I/faster-powerline_daemon

Do not use argparge for parsing environment

client/powerline.c

@@ -3,111 +3,125 @@
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <stddef.h>
 #include <sys/un.h>
 #include <sys/types.h>
 #include <sys/socket.h>
 #include <unistd.h>
 #include <errno.h>
+#include <string.h>
 
-#define handle_error(msg) \
+#define HANDLE_ERROR(msg) \
-	do { perror(msg); exit(EXIT_FAILURE); } while (0)
+	do { \
+		perror(msg); exit(EXIT_FAILURE); \
+	} while (0)
 
-#ifndef TEMP_FAILURE_RETRY
+#define TEMP_FAILURE_RETRY(var, expression) \
-#define TEMP_FAILURE_RETRY(expression) \
+	do { \
-  (                                                                           \
+		long int __result; \
-	({ long int __result;                                                     \
+		do { \
-	   do __result = (long int) (expression);                                 \
+			__result = (long int) (expression); \
-	   while (__result == -1L && errno == EINTR);                             \
+		} while (__result == -1L && errno == EINTR); \
-	   __result; }))
+		var = __result; \
-#endif
+	} while (0)
 
 extern char **environ;
 
-void do_write(int sd, const char *raw, int len) {
+void do_write(int sd, const char *raw, size_t len) {
-	int written = 0, n = -1;
+	size_t written = 0;
+	ptrdiff_t n = -1;
 
 	while (written < len) {
-		n = TEMP_FAILURE_RETRY(write(sd, raw+written, len-written));
+		TEMP_FAILURE_RETRY(n, write(sd, raw + written, len - written));
 		if (n == -1) {
 			close(sd);
-			handle_error("write() failed");
+			HANDLE_ERROR("write() failed");
 		}
-		written += n;
+		written += (size_t) n;
 	}
 }
 
+#ifdef __APPLE__
+# define ADDRESS_TEMPLATE "/tmp/powerline-ipc-%d"
+# define A
+#else
+# define ADDRESS_TEMPLATE "powerline-ipc-%d"
+# define A +1
+#endif
+
+#define ADDRESS_SIZE sizeof(ADDRESS_TEMPLATE) + (sizeof(uid_t) * 4)
+#define NUM_ARGS_SIZE (sizeof(int) * 2)
+#define BUF_SIZE 4096
+#define NEW_ARGV_SIZE 200
+
 int main(int argc, char *argv[]) {
-	int sd = -1, i;
+	int sd = -1;
+	ptrdiff_t i;
 	struct sockaddr_un server;
-	char address[50] = {};
+	char address[ADDRESS_SIZE];
 	const char eof[2] = "\0\0";
-	char buf[4096] = {};
+	char num_args[NUM_ARGS_SIZE];
-	char *newargv[200] = {};
+	char buf[BUF_SIZE];
+	char *newargv[NEW_ARGV_SIZE];
 	char *wd = NULL;
 	char **envp;
 
-	if (argc < 2) { printf("Must provide at least one argument.\n"); return EXIT_FAILURE; }
+	if (argc < 2) {
+		printf("Must provide at least one argument.\n"); return EXIT_FAILURE;
+	}
 
-#ifdef __APPLE__
+	snprintf(address, ADDRESS_SIZE, ADDRESS_TEMPLATE, getuid());
-	snprintf(address, 50, "/tmp/powerline-ipc-%d", getuid());
-#else
-	snprintf(address, 50, "powerline-ipc-%d", getuid());
-#endif
 
 	sd = socket(AF_UNIX, SOCK_STREAM, 0);
-	if (sd == -1) handle_error("socket() failed");
+	if (sd == -1)
+		HANDLE_ERROR("socket() failed");
 
-	memset(&server, 0, sizeof(struct sockaddr_un)); // Clear 
+	memset(&server, 0, sizeof(struct sockaddr_un));
 	server.sun_family = AF_UNIX;
-#ifdef __APPLE__
+	strncpy(server.sun_path A, address, strlen(address));
-	strncpy(server.sun_path, address, strlen(address));
-#else
-	strncpy(server.sun_path+1, address, strlen(address));
-#endif
 
-#ifdef __APPLE__
+	if (connect(sd, (struct sockaddr *) &server, (socklen_t) (sizeof(server.sun_family) + strlen(address) A)) < 0) {
-	if (connect(sd, (struct sockaddr *) &server, sizeof(server.sun_family) + strlen(address)) < 0) {
-#else
-	if (connect(sd, (struct sockaddr *) &server, sizeof(server.sun_family) + strlen(address)+1) < 0) {
-#endif
 		close(sd);
-		// We failed to connect to the daemon, execute powerline instead
+		/* We failed to connect to the daemon, execute powerline instead */
-		argc = (argc < 199) ? argc : 199;
+		argc = (argc < NEW_ARGV_SIZE - 1) ? argc : NEW_ARGV_SIZE - 1;
-		for (i=1; i < argc; i++) newargv[i] = argv[i];
+		for (i = 1; i < argc; i++)
+			newargv[i] = argv[i];
 		newargv[0] = "powerline-render";
 		newargv[argc] = NULL;
 		execvp("powerline-render", newargv);
 	}
 
+	snprintf(num_args, NUM_ARGS_SIZE, "%x", argc - 1);
+	do_write(sd, num_args, strlen(num_args));
+	do_write(sd, eof, 1);
+
 	for (i = 1; i < argc; i++) {
 		do_write(sd, argv[i], strlen(argv[i]));
 		do_write(sd, eof, 1);
 	}
 
-	for(envp=environ; *envp; envp++) {
-		do_write(sd, "--env=", 6);
-		do_write(sd, *envp, strlen(*envp));
-		do_write(sd, eof, 1);
-	}
-
 	wd = getcwd(NULL, 0);
 	if (wd != NULL) {
-		do_write(sd, "--cwd=", 6);
 		do_write(sd, wd, strlen(wd));
 		free(wd); wd = NULL;
 	}
 
+	for(envp=environ; *envp; envp++) {
+		do_write(sd, *envp, strlen(*envp));
+		do_write(sd, eof, 1);
+	}
+
 	do_write(sd, eof, 2);
 
 	i = -1;
 	while (i != 0) {
-		i = TEMP_FAILURE_RETRY(read(sd, buf, 4096));
+		TEMP_FAILURE_RETRY(i, read(sd, buf, BUF_SIZE));
 		if (i == -1) {
 			close(sd);
-			handle_error("read() failed");
+			HANDLE_ERROR("read() failed");
+		} else if (i > 0) {
+			(void) write(STDOUT_FILENO, buf, (size_t) i);
 		}
-		if (i > 0) 
-			write(STDOUT_FILENO, buf, i) || 0;
 	}
 
 	close(sd);

client/powerline.py

@@ -42,7 +42,8 @@ fenc = sys.getfilesystemencoding() or 'utf-8'
 if fenc == 'ascii':
 	fenc = 'utf-8'
 
-args = [x.encode(fenc) if isinstance(x, type('')) else x for x in sys.argv[1:]]
+args = [bytes('%x' % (len(sys.argv) - 1))]
+args.extend((x.encode(fenc) if isinstance(x, type('')) else x for x in sys.argv[1:]))
 
 try:
 	cwd = os.getcwd()
@@ -51,17 +52,17 @@ except EnvironmentError:
 else:
 	if isinstance(cwd, type('')):
 		cwd = cwd.encode(fenc)
-	args.append(b'--cwd=' + cwd)
+	args.append(cwd)
 
 
 env = (k + b'=' + v for k, v in os.environ.items())
 env = (x if isinstance(x, bytes) else x.encode(fenc, 'replace') for x in env)
-args.extend((b'--env=' + x for x in env))
+args.extend(env)
 
 EOF = b'\0\0'
 
 for a in args:
-	eintr_retry_call(sock.sendall, a + EOF[0])
+	eintr_retry_call(sock.sendall, a + b'\0')
 
 eintr_retry_call(sock.sendall, EOF)
 
@@ -74,4 +75,4 @@ while True:
 
 sock.close()
 
-print (b''.join(received))
+sys.stdout.write(b''.join(received))

client/powerline.sh

@@ -5,11 +5,12 @@ ADDRESS="powerline-ipc-${UID:-`id -u`}"
 # Warning: env -0 does not work in busybox. Consider switching to parsing 
 # `set` output in this case
 (
+	printf '%x\0' "$#"
 	for argv in "$@" ; do
 		printf '%s\0' "$argv"
 	done
-	env -0 | sed 's/\(\x00\)\([^\x00]\)\|^/\1--env=\2/g'
+	printf '%s\0' "$PWD"
-	printf -- '--cwd=%s\0' "$PWD"
+	env -0
 ) | socat -lf/dev/null -t 10 - abstract-client:"$ADDRESS"
 
 if test $? -ne 0 ; then

scripts/powerline-daemon

@@ -48,8 +48,6 @@ class NonInteractiveArgParser(ArgumentParser):
 
 
 parser = get_argparser(parser=NonInteractiveArgParser, description='powerline daemon')
-parser.add_argument('--cwd', metavar='PATH')
-parser.add_argument('--env', action='append')
 
 EOF = b'EOF\0\0'
 
@@ -67,11 +65,10 @@ class PowerlineDaemon(ShellPowerline):
 		return super(PowerlineDaemon, self).get_log_handler()
 
 
-def render(args):
+def render(args, environ, cwd):
 	global logger
 	global config_loader
-	environ = dict(((k, v) for k, v in (x.partition('=')[0::2] for x in args.env)))
+	cwd = cwd or environ.get('PWD', '/')
-	cwd = environ.get('PWD', args.cwd or '/')
 	segment_info = {
 		'getcwd': lambda: cwd,
 		'home': environ.get('HOME', home),
@@ -175,11 +172,18 @@ def safe_bytes(o, encoding=encoding):
 		return safe_bytes(str(e), encoding)
 
 
+def parse_args(req):
+	args = [x.decode(encoding) for x in req.split(b'\0') if x]
+	numargs = int(args[0], 16)
+	shell_args = parser.parse_args(args[1:numargs + 1])
+	cwd = args[numargs + 1]
+	environ = dict(((k, v) for k, v in (x.partition('=')[0::2] for x in args[numargs + 2:])))
+	return shell_args, environ, cwd
+
+
 def do_render(req):
 	try:
-		args = [x.decode(encoding) for x in req.split(b'\0') if x]
+		return safe_bytes(render(*parse_args(req)))
-		args = parser.parse_args(args)
-		return safe_bytes(render(args))
 	except Exception as e:
 		return safe_bytes(str(e))