Mercurial > public > mercurial-scm > hg-stable
diff contrib/chg/hgclient.c @ 28060:726f8d6cc324
chg: import frontend sources
These files are copied from
https://bitbucket.org/yuja/chg/ -r f897faa79687
author | Yuya Nishihara <yuya@tcha.org> |
---|---|
date | Sun, 03 Jan 2016 12:39:27 +0900 |
parents | |
children | 098cb7bd46a7 |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/contrib/chg/hgclient.c Sun Jan 03 12:39:27 2016 +0900 @@ -0,0 +1,527 @@ +/* + * A command server client that uses Unix domain socket + * + * Copyright (c) 2011 Yuya Nishihara <yuya@tcha.org> + * + * This software may be used and distributed according to the terms of the + * GNU General Public License version 2 or any later version. + */ + +#include <arpa/inet.h> /* for ntohl(), htonl() */ +#include <assert.h> +#include <ctype.h> +#include <errno.h> +#include <fcntl.h> +#include <signal.h> +#include <stdint.h> +#include <stdio.h> +#include <stdlib.h> +#include <string.h> +#include <sys/socket.h> +#include <sys/stat.h> +#include <sys/un.h> +#include <unistd.h> + +#include "hgclient.h" +#include "util.h" + +enum { + CAP_GETENCODING = 0x0001, + CAP_RUNCOMMAND = 0x0002, + /* cHg extension: */ + CAP_ATTACHIO = 0x0100, + CAP_CHDIR = 0x0200, + CAP_GETPAGER = 0x0400, + CAP_SETENV = 0x0800, +}; + +typedef struct { + const char *name; + unsigned int flag; +} cappair_t; + +static const cappair_t captable[] = { + {"getencoding", CAP_GETENCODING}, + {"runcommand", CAP_RUNCOMMAND}, + {"attachio", CAP_ATTACHIO}, + {"chdir", CAP_CHDIR}, + {"getpager", CAP_GETPAGER}, + {"setenv", CAP_SETENV}, + {NULL, 0}, /* terminator */ +}; + +typedef struct { + char ch; + char *data; + size_t maxdatasize; + size_t datasize; +} context_t; + +struct hgclient_tag_ { + int sockfd; + pid_t pid; + context_t ctx; + unsigned int capflags; +}; + +static const size_t defaultdatasize = 4096; + +static void initcontext(context_t *ctx) +{ + ctx->ch = '\0'; + ctx->data = malloc(defaultdatasize); + ctx->maxdatasize = (ctx->data) ? defaultdatasize : 0; + ctx->datasize = 0; + debugmsg("initialize context buffer with size %zu", ctx->maxdatasize); +} + +static void enlargecontext(context_t *ctx, size_t newsize) +{ + if (newsize <= ctx->maxdatasize) + return; + + newsize = defaultdatasize + * ((newsize + defaultdatasize - 1) / defaultdatasize); + char *p = realloc(ctx->data, newsize); + if (!p) + abortmsg("failed to allocate buffer"); + ctx->data = p; + ctx->maxdatasize = newsize; + debugmsg("enlarge context buffer to %zu", ctx->maxdatasize); +} + +static void freecontext(context_t *ctx) +{ + debugmsg("free context buffer"); + free(ctx->data); + ctx->data = NULL; + ctx->maxdatasize = 0; + ctx->datasize = 0; +} + +/* Read channeled response from cmdserver */ +static void readchannel(hgclient_t *hgc) +{ + assert(hgc); + + ssize_t rsize = recv(hgc->sockfd, &hgc->ctx.ch, sizeof(hgc->ctx.ch), 0); + if (rsize != sizeof(hgc->ctx.ch)) + abortmsg("failed to read channel"); + + uint32_t datasize_n; + rsize = recv(hgc->sockfd, &datasize_n, sizeof(datasize_n), 0); + if (rsize != sizeof(datasize_n)) + abortmsg("failed to read data size"); + + /* datasize denotes the maximum size to write if input request */ + hgc->ctx.datasize = ntohl(datasize_n); + enlargecontext(&hgc->ctx, hgc->ctx.datasize); + + if (isupper(hgc->ctx.ch) && hgc->ctx.ch != 'S') + return; /* assumes input request */ + + size_t cursize = 0; + while (cursize < hgc->ctx.datasize) { + rsize = recv(hgc->sockfd, hgc->ctx.data + cursize, + hgc->ctx.datasize - cursize, 0); + if (rsize < 0) + abortmsg("failed to read data block"); + cursize += rsize; + } +} + +static void sendall(int sockfd, const void *data, size_t datasize) +{ + const char *p = data; + const char *const endp = p + datasize; + while (p < endp) { + ssize_t r = send(sockfd, p, endp - p, 0); + if (r < 0) + abortmsg("cannot communicate (errno = %d)", errno); + p += r; + } +} + +/* Write lengh-data block to cmdserver */ +static void writeblock(const hgclient_t *hgc) +{ + assert(hgc); + + const uint32_t datasize_n = htonl(hgc->ctx.datasize); + sendall(hgc->sockfd, &datasize_n, sizeof(datasize_n)); + + sendall(hgc->sockfd, hgc->ctx.data, hgc->ctx.datasize); +} + +static void writeblockrequest(const hgclient_t *hgc, const char *chcmd) +{ + debugmsg("request %s, block size %zu", chcmd, hgc->ctx.datasize); + + char buf[strlen(chcmd) + 1]; + memcpy(buf, chcmd, sizeof(buf) - 1); + buf[sizeof(buf) - 1] = '\n'; + sendall(hgc->sockfd, buf, sizeof(buf)); + + writeblock(hgc); +} + +/* Build '\0'-separated list of args. argsize < 0 denotes that args are + * terminated by NULL. */ +static void packcmdargs(context_t *ctx, const char *const args[], + ssize_t argsize) +{ + ctx->datasize = 0; + const char *const *const end = (argsize >= 0) ? args + argsize : NULL; + for (const char *const *it = args; it != end && *it; ++it) { + const size_t n = strlen(*it) + 1; /* include '\0' */ + enlargecontext(ctx, ctx->datasize + n); + memcpy(ctx->data + ctx->datasize, *it, n); + ctx->datasize += n; + } + + if (ctx->datasize > 0) + --ctx->datasize; /* strip last '\0' */ +} + +/* Extract '\0'-separated list of args to new buffer, terminated by NULL */ +static const char **unpackcmdargsnul(const context_t *ctx) +{ + const char **args = NULL; + size_t nargs = 0, maxnargs = 0; + const char *s = ctx->data; + const char *e = ctx->data + ctx->datasize; + for (;;) { + if (nargs + 1 >= maxnargs) { /* including last NULL */ + maxnargs += 256; + args = realloc(args, maxnargs * sizeof(args[0])); + if (!args) + abortmsg("failed to allocate args buffer"); + } + args[nargs] = s; + nargs++; + s = memchr(s, '\0', e - s); + if (!s) + break; + s++; + } + args[nargs] = NULL; + return args; +} + +static void handlereadrequest(hgclient_t *hgc) +{ + context_t *ctx = &hgc->ctx; + size_t r = fread(ctx->data, sizeof(ctx->data[0]), ctx->datasize, stdin); + ctx->datasize = r; + writeblock(hgc); +} + +/* Read single-line */ +static void handlereadlinerequest(hgclient_t *hgc) +{ + context_t *ctx = &hgc->ctx; + if (!fgets(ctx->data, ctx->datasize, stdin)) + ctx->data[0] = '\0'; + ctx->datasize = strlen(ctx->data); + writeblock(hgc); +} + +/* Execute the requested command and write exit code */ +static void handlesystemrequest(hgclient_t *hgc) +{ + context_t *ctx = &hgc->ctx; + enlargecontext(ctx, ctx->datasize + 1); + ctx->data[ctx->datasize] = '\0'; /* terminate last string */ + + const char **args = unpackcmdargsnul(ctx); + if (!args[0] || !args[1]) + abortmsg("missing command or cwd in system request"); + debugmsg("run '%s' at '%s'", args[0], args[1]); + int32_t r = runshellcmd(args[0], args + 2, args[1]); + free(args); + + uint32_t r_n = htonl(r); + memcpy(ctx->data, &r_n, sizeof(r_n)); + ctx->datasize = sizeof(r_n); + writeblock(hgc); +} + +/* Read response of command execution until receiving 'r'-esult */ +static void handleresponse(hgclient_t *hgc) +{ + for (;;) { + readchannel(hgc); + context_t *ctx = &hgc->ctx; + debugmsg("response read from channel %c, size %zu", + ctx->ch, ctx->datasize); + switch (ctx->ch) { + case 'o': + fwrite(ctx->data, sizeof(ctx->data[0]), ctx->datasize, + stdout); + break; + case 'e': + fwrite(ctx->data, sizeof(ctx->data[0]), ctx->datasize, + stderr); + break; + case 'd': + /* assumes last char is '\n' */ + ctx->data[ctx->datasize - 1] = '\0'; + debugmsg("server: %s", ctx->data); + break; + case 'r': + return; + case 'I': + handlereadrequest(hgc); + break; + case 'L': + handlereadlinerequest(hgc); + break; + case 'S': + handlesystemrequest(hgc); + break; + default: + if (isupper(ctx->ch)) + abortmsg("cannot handle response (ch = %c)", + ctx->ch); + } + } +} + +static unsigned int parsecapabilities(const char *s, const char *e) +{ + unsigned int flags = 0; + while (s < e) { + const char *t = strchr(s, ' '); + if (!t || t > e) + t = e; + const cappair_t *cap; + for (cap = captable; cap->flag; ++cap) { + size_t n = t - s; + if (strncmp(s, cap->name, n) == 0 && + strlen(cap->name) == n) { + flags |= cap->flag; + break; + } + } + s = t + 1; + } + return flags; +} + +static void readhello(hgclient_t *hgc) +{ + readchannel(hgc); + context_t *ctx = &hgc->ctx; + if (ctx->ch != 'o') + abortmsg("unexpected channel of hello message (ch = %c)", + ctx->ch); + enlargecontext(ctx, ctx->datasize + 1); + ctx->data[ctx->datasize] = '\0'; + debugmsg("hello received: %s (size = %zu)", ctx->data, ctx->datasize); + + const char *s = ctx->data; + const char *const dataend = ctx->data + ctx->datasize; + while (s < dataend) { + const char *t = strchr(s, ':'); + if (!t || t[1] != ' ') + break; + const char *u = strchr(t + 2, '\n'); + if (!u) + u = dataend; + if (strncmp(s, "capabilities:", t - s + 1) == 0) { + hgc->capflags = parsecapabilities(t + 2, u); + } else if (strncmp(s, "pid:", t - s + 1) == 0) { + hgc->pid = strtol(t + 2, NULL, 10); + } + s = u + 1; + } + debugmsg("capflags=0x%04x, pid=%d", hgc->capflags, hgc->pid); +} + +static void attachio(hgclient_t *hgc) +{ + debugmsg("request attachio"); + static const char chcmd[] = "attachio\n"; + sendall(hgc->sockfd, chcmd, sizeof(chcmd) - 1); + readchannel(hgc); + context_t *ctx = &hgc->ctx; + if (ctx->ch != 'I') + abortmsg("unexpected response for attachio (ch = %c)", ctx->ch); + + static const int fds[3] = {STDIN_FILENO, STDOUT_FILENO, STDERR_FILENO}; + struct msghdr msgh; + memset(&msgh, 0, sizeof(msgh)); + struct iovec iov = {ctx->data, ctx->datasize}; /* dummy payload */ + msgh.msg_iov = &iov; + msgh.msg_iovlen = 1; + char fdbuf[CMSG_SPACE(sizeof(fds))]; + msgh.msg_control = fdbuf; + msgh.msg_controllen = sizeof(fdbuf); + struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msgh); + cmsg->cmsg_level = SOL_SOCKET; + cmsg->cmsg_type = SCM_RIGHTS; + cmsg->cmsg_len = CMSG_LEN(sizeof(fds)); + memcpy(CMSG_DATA(cmsg), fds, sizeof(fds)); + msgh.msg_controllen = cmsg->cmsg_len; + ssize_t r = sendmsg(hgc->sockfd, &msgh, 0); + if (r < 0) + abortmsg("sendmsg failed (errno = %d)", errno); + + handleresponse(hgc); + int32_t n; + if (ctx->datasize != sizeof(n)) + abortmsg("unexpected size of attachio result"); + memcpy(&n, ctx->data, sizeof(n)); + n = ntohl(n); + if (n != sizeof(fds) / sizeof(fds[0])) + abortmsg("failed to send fds (n = %d)", n); +} + +static void chdirtocwd(hgclient_t *hgc) +{ + if (!getcwd(hgc->ctx.data, hgc->ctx.maxdatasize)) + abortmsg("failed to getcwd (errno = %d)", errno); + hgc->ctx.datasize = strlen(hgc->ctx.data); + writeblockrequest(hgc, "chdir"); +} + +/*! + * Open connection to per-user cmdserver + * + * If no background server running, returns NULL. + */ +hgclient_t *hgc_open(const char *sockname) +{ + int fd = socket(AF_UNIX, SOCK_STREAM, 0); + if (fd < 0) + abortmsg("cannot create socket (errno = %d)", errno); + + /* don't keep fd on fork(), so that it can be closed when the parent + * process get terminated. */ + int flags = fcntl(fd, F_GETFD); + if (flags < 0) + abortmsg("cannot get flags of socket (errno = %d)", errno); + if (fcntl(fd, F_SETFD, flags | FD_CLOEXEC) < 0) + abortmsg("cannot set flags of socket (errno = %d)", errno); + + struct sockaddr_un addr; + addr.sun_family = AF_UNIX; + strncpy(addr.sun_path, sockname, sizeof(addr.sun_path)); + addr.sun_path[sizeof(addr.sun_path) - 1] = '\0'; + + debugmsg("connect to %s", addr.sun_path); + int r = connect(fd, (struct sockaddr *)&addr, sizeof(addr)); + if (r < 0) { + close(fd); + if (errno == ENOENT || errno == ECONNREFUSED) + return NULL; + abortmsg("cannot connect to %s (errno = %d)", + addr.sun_path, errno); + } + + hgclient_t *hgc = malloc(sizeof(hgclient_t)); + if (!hgc) + abortmsg("failed to allocate hgclient object"); + memset(hgc, 0, sizeof(*hgc)); + hgc->sockfd = fd; + initcontext(&hgc->ctx); + + readhello(hgc); + if (!(hgc->capflags & CAP_RUNCOMMAND)) + abortmsg("insufficient capability: runcommand"); + if (hgc->capflags & CAP_ATTACHIO) + attachio(hgc); + if (hgc->capflags & CAP_CHDIR) + chdirtocwd(hgc); + + return hgc; +} + +/*! + * Close connection and free allocated memory + */ +void hgc_close(hgclient_t *hgc) +{ + assert(hgc); + freecontext(&hgc->ctx); + close(hgc->sockfd); + free(hgc); +} + +pid_t hgc_peerpid(const hgclient_t *hgc) +{ + assert(hgc); + return hgc->pid; +} + +/*! + * Execute the specified Mercurial command + * + * @return result code + */ +int hgc_runcommand(hgclient_t *hgc, const char *const args[], size_t argsize) +{ + assert(hgc); + + packcmdargs(&hgc->ctx, args, argsize); + writeblockrequest(hgc, "runcommand"); + handleresponse(hgc); + + int32_t exitcode_n; + if (hgc->ctx.datasize != sizeof(exitcode_n)) { + abortmsg("unexpected size of exitcode"); + } + memcpy(&exitcode_n, hgc->ctx.data, sizeof(exitcode_n)); + return ntohl(exitcode_n); +} + +/*! + * (Re-)send client's stdio channels so that the server can access to tty + */ +void hgc_attachio(hgclient_t *hgc) +{ + assert(hgc); + if (!(hgc->capflags & CAP_ATTACHIO)) + return; + attachio(hgc); +} + +/*! + * Get pager command for the given Mercurial command args + * + * If no pager enabled, returns NULL. The return value becomes invalid + * once you run another request to hgc. + */ +const char *hgc_getpager(hgclient_t *hgc, const char *const args[], + size_t argsize) +{ + assert(hgc); + + if (!(hgc->capflags & CAP_GETPAGER)) + return NULL; + + packcmdargs(&hgc->ctx, args, argsize); + writeblockrequest(hgc, "getpager"); + handleresponse(hgc); + + if (hgc->ctx.datasize < 1 || hgc->ctx.data[0] == '\0') + return NULL; + enlargecontext(&hgc->ctx, hgc->ctx.datasize + 1); + hgc->ctx.data[hgc->ctx.datasize] = '\0'; + return hgc->ctx.data; +} + +/*! + * Update server's environment variables + * + * @param envp list of environment variables in "NAME=VALUE" format, + * terminated by NULL. + */ +void hgc_setenv(hgclient_t *hgc, const char *const envp[]) +{ + assert(hgc && envp); + if (!(hgc->capflags & CAP_SETENV)) + return; + packcmdargs(&hgc->ctx, envp, /*argsize*/ -1); + writeblockrequest(hgc, "setenv"); +}