Merge pull request #955 from ZyX-I/faster-powerline_daemon
Do not use argparge for parsing environment
This commit is contained in:
commit
d5964bc25b
|
@ -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);
|
||||
|
|
|
@ -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))
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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))
|
||||
|
||||
|
|
Loading…
Reference in New Issue