unix/uxagentc.c

   1 /*
   2  * SSH agent client code.
   3  */
   4
   5 #include <stdio.h>
   6 #include <stdlib.h>
   7 #include <assert.h>
   8 #include <unistd.h>
   9 #include <sys/socket.h>
  10 #include <sys/un.h>
  11 #include <fcntl.h>
  12
  13 #include "putty.h"
  14 #include "misc.h"
  15 #include "tree234.h"
  16 #include "puttymem.h"
  17
  18 int agent_exists(void)
  19 {
  20     const char *p = getenv("SSH_AUTH_SOCK");
  21     if (p && *p)
  22         return TRUE;
  23     return FALSE;
  24 }
  25
  26 static tree234 *agent_pending_queries;
  27 struct agent_pending_query {
  28     int fd;
  29     char *retbuf;
  30     char sizebuf[4];
  31     int retsize, retlen;
  32     void (*callback)(void *, void *, int);
  33     void *callback_ctx;
  34 };
  35 static int agent_conncmp(void *av, void *bv)
  36 {
  37     agent_pending_query *a = (agent_pending_query *) av;
  38     agent_pending_query *b = (agent_pending_query *) bv;
  39     if (a->fd < b->fd)
  40         return -1;
  41     if (a->fd > b->fd)
  42         return +1;
  43     return 0;
  44 }
  45 static int agent_connfind(void *av, void *bv)
  46 {
  47     int afd = *(int *) av;
  48     agent_pending_query *b = (agent_pending_query *) bv;
  49     if (afd < b->fd)
  50         return -1;
  51     if (afd > b->fd)
  52         return +1;
  53     return 0;
  54 }
  55
  56 /*
  57  * Attempt to read from an agent socket fd. Returns 0 if the expected
  58  * response is as yet incomplete; returns 1 if it's either complete
  59  * (conn->retbuf non-NULL and filled with something useful) or has
  60  * failed totally (conn->retbuf is NULL).
  61  */
  62 static int agent_try_read(agent_pending_query *conn)
  63 {
  64     int ret;
  65
  66     ret = read(conn->fd, conn->retbuf+conn->retlen,
  67                conn->retsize-conn->retlen);
  68     if (ret <= 0) {
  69         if (conn->retbuf != conn->sizebuf) sfree(conn->retbuf);
  70         conn->retbuf = NULL;
  71         conn->retlen = 0;
  72         return 1;
  73     }
  74     conn->retlen += ret;
  75     if (conn->retsize == 4 && conn->retlen == 4) {
  76         conn->retsize = toint(GET_32BIT(conn->retbuf) + 4);
  77         if (conn->retsize <= 0) {
  78             conn->retbuf = NULL;
  79             conn->retlen = 0;
  80             return -1;                 /* way too large */
  81         }
  82         assert(conn->retbuf == conn->sizebuf);
  83         conn->retbuf = snewn(conn->retsize, char);
  84         memcpy(conn->retbuf, conn->sizebuf, 4);
  85     }
  86
  87     if (conn->retlen < conn->retsize)
  88         return 0;                      /* more data to come */
  89
  90     return 1;
  91 }
  92
  93 void agent_cancel_query(agent_pending_query *conn)
  94 {
  95     uxsel_del(conn->fd);
  96     close(conn->fd);
  97     del234(agent_pending_queries, conn);
  98     sfree(conn);
  99 }
 100
 101 static int agent_select_result(int fd, int event)
 102 {
 103     agent_pending_query *conn;
 104
 105     assert(event == 1);                /* not selecting for anything but R */
 106
 107     conn = find234(agent_pending_queries, &fd, agent_connfind);
 108     if (!conn) {
 109         uxsel_del(fd);
 110         return 1;
 111     }
 112
 113     if (!agent_try_read(conn))
 114         return 0;                      /* more data to come */
 115
 116     /*
 117      * We have now completed the agent query. Do the callback, and
 118      * clean up. (Of course we don't free retbuf, since ownership
 119      * of that passes to the callback.)
 120      */
 121     conn->callback(conn->callback_ctx, conn->retbuf, conn->retlen);
 122     agent_cancel_query(conn);
 123     return 0;
 124 }
 125
 126 agent_pending_query *agent_query(
 127     void *in, int inlen, void **out, int *outlen,
 128     void (*callback)(void *, void *, int), void *callback_ctx)
 129 {
 130     char *name;
 131     int sock;
 132     struct sockaddr_un addr;
 133     int done;
 134     agent_pending_query *conn;
 135
 136     name = getenv("SSH_AUTH_SOCK");
 137     if (!name)
 138         goto failure;
 139
 140     sock = socket(PF_UNIX, SOCK_STREAM, 0);
 141     if (sock < 0) {
 142         perror("socket(PF_UNIX)");
 143         exit(1);
 144     }
 145
 146     cloexec(sock);
 147
 148     addr.sun_family = AF_UNIX;
 149     strncpy(addr.sun_path, name, sizeof(addr.sun_path));
 150     if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
 151         close(sock);
 152         goto failure;
 153     }
 154
 155     for (done = 0; done < inlen ;) {
 156         int ret = write(sock, (char *)in + done, inlen - done);
 157         if (ret <= 0) {
 158             close(sock);
 159             goto failure;
 160         }
 161         done += ret;
 162     }
 163
 164     conn = snew(agent_pending_query);
 165     conn->fd = sock;
 166     conn->retbuf = conn->sizebuf;
 167     conn->retsize = 4;
 168     conn->retlen = 0;
 169     conn->callback = callback;
 170     conn->callback_ctx = callback_ctx;
 171
 172     if (!callback) {
 173         /*
 174          * Bodge to permit making deliberately synchronous agent
 175          * requests. Used by Unix Pageant in command-line client mode,
 176          * which is legit because it really is true that no other part
 177          * of the program is trying to get anything useful done
 178          * simultaneously. But this special case shouldn't be used in
 179          * any more general program.
 180          */
 181         no_nonblock(conn->fd);
 182         while (!agent_try_read(conn))
 183             /* empty loop body */;
 184
 185         *out = conn->retbuf;
 186         *outlen = conn->retlen;
 187         sfree(conn);
 188         return NULL;
 189     }
 190
 191     /*
 192      * Otherwise do it properly: add conn to the tree of agent
 193      * connections currently in flight, return 0 to indicate that the
 194      * response hasn't been received yet, and call the callback when
 195      * select_result comes back to us.
 196      */
 197     if (!agent_pending_queries)
 198         agent_pending_queries = newtree234(agent_conncmp);
 199     add234(agent_pending_queries, conn);
 200
 201     uxsel_set(sock, 1, agent_select_result);
 202     return conn;
 203
 204     failure:
 205     *out = NULL;
 206     *outlen = 0;
 207     return NULL;
 208 }