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

Do not use argparge for parsing environment
This commit is contained in:
Nikolai Aleksandrovich Pavlov 2014-08-04 03:55:52 +04:00
commit d5964bc25b
4 changed files with 86 additions and 66 deletions

View File

@ -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) \
do { perror(msg); exit(EXIT_FAILURE); } while (0)
#define HANDLE_ERROR(msg) \
do { \
perror(msg); exit(EXIT_FAILURE); \
} while (0)
#ifndef TEMP_FAILURE_RETRY
#define TEMP_FAILURE_RETRY(expression) \
( \
({ long int __result; \
do __result = (long int) (expression); \
while (__result == -1L && errno == EINTR); \
__result; }))
#endif
#define TEMP_FAILURE_RETRY(var, expression) \
do { \
long int __result; \
do { \
__result = (long int) (expression); \
} while (__result == -1L && errno == EINTR); \
var = __result; \
} while (0)
extern char **environ;
void do_write(int sd, const char *raw, int len) {
int written = 0, n = -1;
void do_write(int sd, const char *raw, size_t len) {
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 *newargv[200] = {};
char num_args[NUM_ARGS_SIZE];
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, 50, "/tmp/powerline-ipc-%d", getuid());
#else
snprintf(address, 50, "powerline-ipc-%d", getuid());
#endif
snprintf(address, ADDRESS_SIZE, ADDRESS_TEMPLATE, getuid());
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, address, strlen(address));
#else
strncpy(server.sun_path+1, address, strlen(address));
#endif
strncpy(server.sun_path A, address, strlen(address));
#ifdef __APPLE__
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
if (connect(sd, (struct sockaddr *) &server, (socklen_t) (sizeof(server.sun_family) + strlen(address) A)) < 0) {
close(sd);
// We failed to connect to the daemon, execute powerline instead
argc = (argc < 199) ? argc : 199;
for (i=1; i < argc; i++) newargv[i] = argv[i];
/* We failed to connect to the daemon, execute powerline instead */
argc = (argc < NEW_ARGV_SIZE - 1) ? argc : NEW_ARGV_SIZE - 1;
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);

View File

@ -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))

View File

@ -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 -- '--cwd=%s\0' "$PWD"
printf '%s\0' "$PWD"
env -0
) | socat -lf/dev/null -t 10 - abstract-client:"$ADDRESS"
if test $? -ne 0 ; then

View File

@ -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 = environ.get('PWD', args.cwd or '/')
cwd = cwd or environ.get('PWD', '/')
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]
args = parser.parse_args(args)
return safe_bytes(render(args))
return safe_bytes(render(*parse_args(req)))
except Exception as e:
return safe_bytes(str(e))