Fix major memory leak in sftp_cmd_ls (thanks to Hans-Juergen Petrich
[u/mdw/putty] / pageant.c
index 20aaaf4..b14b58e 100644 (file)
--- a/pageant.c
+++ b/pageant.c
@@ -9,6 +9,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <ctype.h>
+#include <assert.h>
 #include <tchar.h>
 
 #include "ssh.h"
@@ -43,11 +44,45 @@ static HINSTANCE instance;
 static HWND main_hwnd;
 static HWND keylist;
 static HWND aboutbox;
-static HMENU systray_menu;
+static HMENU systray_menu, session_menu;
 static int already_running;
 static int requested_help;
 
 static char *help_path;
+static char *putty_path;
+
+#define IDM_PUTTY         0x0060
+#define IDM_SESSIONS_BASE 0x1000
+#define IDM_SESSIONS_MAX  0x2000
+#define PUTTY_REGKEY      "Software\\SimonTatham\\PuTTY\\Sessions"
+#define PUTTY_DEFAULT     "Default%20Settings"
+static int initial_menuitems_count;
+
+/* Un-munge session names out of the registry. */
+static void unmungestr(char *in, char *out, int outlen)
+{
+    while (*in) {
+       if (*in == '%' && in[1] && in[2]) {
+           int i, j;
+
+           i = in[1] - '0';
+           i -= (i > 9 ? 7 : 0);
+           j = in[2] - '0';
+           j -= (j > 9 ? 7 : 0);
+
+           *out++ = (i << 4) + j;
+           if (!--outlen)
+               return;
+           in += 3;
+       } else {
+           *out++ = *in++;
+           if (!--outlen)
+               return;
+       }
+    }
+    *out = '\0';
+    return;
+}
 
 static tree234 *rsakeys, *ssh2keys;
 
@@ -66,6 +101,14 @@ void agent_query(void *in, int inlen, void **out, int *outlen);
 int agent_exists(void);
 
 /*
+ * Forward references
+ */
+static void *make_keylist1(int *length);
+static void *make_keylist2(int *length);
+static void *get_keylist1(void);
+static void *get_keylist2(void);
+
+/*
  * We need this to link with the RSA code, because rsaencrypt()
  * pads its data with random bytes. Since we only use rsadecrypt()
  * and the signing functions, which are deterministic, this should
@@ -357,6 +400,65 @@ static void add_keyfile(char *filename)
        return;
     }
 
+    /*
+     * See if the key is already loaded (in the primary Pageant,
+     * which may or may not be us).
+     */
+    {
+       void *blob;
+       unsigned char *keylist, *p;
+       int i, nkeys, bloblen;
+
+       if (ver == 1) {
+           if (!rsakey_pubblob(filename, &blob, &bloblen)) {
+               MessageBox(NULL, "Couldn't load private key.", APPNAME,
+                          MB_OK | MB_ICONERROR);
+               return;
+           }
+           keylist = get_keylist1();
+       } else {
+           unsigned char *blob2;
+           blob = ssh2_userkey_loadpub(filename, NULL, &bloblen);
+           if (!blob) {
+               MessageBox(NULL, "Couldn't load private key.", APPNAME,
+                          MB_OK | MB_ICONERROR);
+               return;
+           }
+           /* For our purposes we want the blob prefixed with its length */
+           blob2 = smalloc(bloblen+4);
+           PUT_32BIT(blob2, bloblen);
+           memcpy(blob2 + 4, blob, bloblen);
+           sfree(blob);
+           blob = blob2;
+
+           keylist = get_keylist2();
+       }
+       if (keylist) {
+           nkeys = GET_32BIT(keylist);
+           p = keylist + 4;
+
+           for (i = 0; i < nkeys; i++) {
+               if (!memcmp(blob, p, bloblen)) {
+                   /* Key is already present; we can now leave. */
+                   sfree(keylist);
+                   sfree(blob);
+                   return;
+               }
+               /* Now skip over public blob */
+               if (ver == 1)
+                   p += rsa_public_blob_len(p);
+               else
+                   p += 4 + GET_32BIT(p);
+               /* Now skip over comment field */
+               p += 4 + GET_32BIT(p);
+           }
+
+           sfree(keylist);
+       }
+
+       sfree(blob);
+    }
+
     if (ver == 1)
        needs_pass = rsakey_encrypted(filename, &comment);
     else
@@ -461,6 +563,9 @@ static void add_keyfile(char *filename)
                MessageBox(NULL, "The already running Pageant "
                           "refused to add the key.", APPNAME,
                           MB_OK | MB_ICONERROR);
+
+           sfree(request);
+           sfree(response);
        } else {
            if (add234(rsakeys, rkey) != rkey)
                sfree(rkey);           /* already present, don't waste RAM */
@@ -503,6 +608,9 @@ static void add_keyfile(char *filename)
                MessageBox(NULL, "The already running Pageant"
                           "refused to add the key.", APPNAME,
                           MB_OK | MB_ICONERROR);
+
+           sfree(request);
+           sfree(response);
        } else {
            if (add234(ssh2keys, skey) != skey) {
                skey->alg->freekey(skey->data);
@@ -513,6 +621,162 @@ static void add_keyfile(char *filename)
 }
 
 /*
+ * Create an SSH1 key list in a malloc'ed buffer; return its
+ * length.
+ */
+static void *make_keylist1(int *length)
+{
+    int i, nkeys, len;
+    struct RSAKey *key;
+    unsigned char *blob, *p, *ret;
+    int bloblen;
+
+    /*
+     * Count up the number and length of keys we hold.
+     */
+    len = 4;
+    nkeys = 0;
+    for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
+       nkeys++;
+       blob = rsa_public_blob(key, &bloblen);
+       len += bloblen;
+       sfree(blob);
+       len += 4 + strlen(key->comment);
+    }
+
+    /* Allocate the buffer. */
+    p = ret = smalloc(len);
+    if (length) *length = len;
+
+    PUT_32BIT(p, nkeys);
+    p += 4;
+    for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
+       blob = rsa_public_blob(key, &bloblen);
+       memcpy(p, blob, bloblen);
+       p += bloblen;
+       sfree(blob);
+       PUT_32BIT(p, strlen(key->comment));
+       memcpy(p + 4, key->comment, strlen(key->comment));
+       p += 4 + strlen(key->comment);
+    }
+
+    assert(p - ret == len);
+    return ret;
+}
+
+/*
+ * Create an SSH2 key list in a malloc'ed buffer; return its
+ * length.
+ */
+static void *make_keylist2(int *length)
+{
+    struct ssh2_userkey *key;
+    int i, len, nkeys;
+    unsigned char *blob, *p, *ret;
+    int bloblen;
+
+    /*
+     * Count up the number and length of keys we hold.
+     */
+    len = 4;
+    nkeys = 0;
+    for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
+       nkeys++;
+       len += 4;              /* length field */
+       blob = key->alg->public_blob(key->data, &bloblen);
+       len += bloblen;
+       sfree(blob);
+       len += 4 + strlen(key->comment);
+    }
+
+    /* Allocate the buffer. */
+    p = ret = smalloc(len);
+    if (length) *length = len;
+
+    /*
+     * Packet header is the obvious five bytes, plus four
+     * bytes for the key count.
+     */
+    PUT_32BIT(p, nkeys);
+    p += 4;
+    for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
+       blob = key->alg->public_blob(key->data, &bloblen);
+       PUT_32BIT(p, bloblen);
+       p += 4;
+       memcpy(p, blob, bloblen);
+       p += bloblen;
+       sfree(blob);
+       PUT_32BIT(p, strlen(key->comment));
+       memcpy(p + 4, key->comment, strlen(key->comment));
+       p += 4 + strlen(key->comment);
+    }
+
+    assert(p - ret == len);
+    return ret;
+}
+
+/*
+ * Acquire a keylist1 from the primary Pageant; this means either
+ * calling make_keylist1 (if that's us) or sending a message to the
+ * primary Pageant (if it's not).
+ */
+static void *get_keylist1(void)
+{
+    void *ret;
+
+    if (already_running) {
+       unsigned char request[5], *response;
+       void *vresponse;
+       int resplen;
+       request[4] = SSH1_AGENTC_REQUEST_RSA_IDENTITIES;
+       PUT_32BIT(request, 4);
+
+       agent_query(request, 5, &vresponse, &resplen);
+       response = vresponse;
+       if (resplen < 5 || response[4] != SSH1_AGENT_RSA_IDENTITIES_ANSWER)
+           return NULL;
+
+       ret = smalloc(resplen-5);
+       memcpy(ret, response+5, resplen-5);
+       sfree(response);
+    } else {
+       ret = make_keylist1(NULL);
+    }
+    return ret;
+}
+
+/*
+ * Acquire a keylist2 from the primary Pageant; this means either
+ * calling make_keylist2 (if that's us) or sending a message to the
+ * primary Pageant (if it's not).
+ */
+static void *get_keylist2(void)
+{
+    void *ret;
+
+    if (already_running) {
+       unsigned char request[5], *response;
+       void *vresponse;
+       int resplen;
+
+       request[4] = SSH2_AGENTC_REQUEST_IDENTITIES;
+       PUT_32BIT(request, 4);
+
+       agent_query(request, 5, &vresponse, &resplen);
+       response = vresponse;
+       if (resplen < 5 || response[4] != SSH2_AGENT_IDENTITIES_ANSWER)
+           return NULL;
+
+       ret = smalloc(resplen-5);
+       memcpy(ret, response+5, resplen-5);
+       sfree(response);
+    } else {
+       ret = make_keylist2(NULL);
+    }
+    return ret;
+}
+
+/*
  * This is the main agent function that answers messages.
  */
 static void answer_msg(void *msg)
@@ -533,42 +797,18 @@ static void answer_msg(void *msg)
         * Reply with SSH1_AGENT_RSA_IDENTITIES_ANSWER.
         */
        {
-           struct RSAKey *key;
-           int len, nkeys;
-           int i;
-
-           /*
-            * Count up the number and length of keys we hold.
-            */
-           len = nkeys = 0;
-           for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
-               nkeys++;
-               len += 4;              /* length field */
-               len += ssh1_bignum_length(key->exponent);
-               len += ssh1_bignum_length(key->modulus);
-               len += 4 + strlen(key->comment);
-           }
+           int len;
+           void *keylist;
 
-           /*
-            * Packet header is the obvious five bytes, plus four
-            * bytes for the key count.
-            */
-           len += 5 + 4;
-           if (len > AGENT_MAX_MSGLEN)
-               goto failure;          /* aaargh! too much stuff! */
-           PUT_32BIT(ret, len - 4);
            ret[4] = SSH1_AGENT_RSA_IDENTITIES_ANSWER;
-           PUT_32BIT(ret + 5, nkeys);
-           p = ret + 5 + 4;
-           for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
-               PUT_32BIT(p, bignum_bitcount(key->modulus));
-               p += 4;
-               p += ssh1_write_bignum(p, key->exponent);
-               p += ssh1_write_bignum(p, key->modulus);
-               PUT_32BIT(p, strlen(key->comment));
-               memcpy(p + 4, key->comment, strlen(key->comment));
-               p += 4 + strlen(key->comment);
+           keylist = make_keylist1(&len);
+           if (len + 5 > AGENT_MAX_MSGLEN) {
+               sfree(keylist);
+               goto failure;
            }
+           PUT_32BIT(ret, len + 1);
+           memcpy(ret + 5, keylist, len);
+           sfree(keylist);
        }
        break;
       case SSH2_AGENTC_REQUEST_IDENTITIES:
@@ -576,47 +816,18 @@ static void answer_msg(void *msg)
         * Reply with SSH2_AGENT_IDENTITIES_ANSWER.
         */
        {
-           struct ssh2_userkey *key;
-           int len, nkeys;
-           unsigned char *blob;
-           int bloblen;
-           int i;
-
-           /*
-            * Count up the number and length of keys we hold.
-            */
-           len = nkeys = 0;
-           for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
-               nkeys++;
-               len += 4;              /* length field */
-               blob = key->alg->public_blob(key->data, &bloblen);
-               len += bloblen;
-               sfree(blob);
-               len += 4 + strlen(key->comment);
-           }
+           int len;
+           void *keylist;
 
-           /*
-            * Packet header is the obvious five bytes, plus four
-            * bytes for the key count.
-            */
-           len += 5 + 4;
-           if (len > AGENT_MAX_MSGLEN)
-               goto failure;          /* aaargh! too much stuff! */
-           PUT_32BIT(ret, len - 4);
            ret[4] = SSH2_AGENT_IDENTITIES_ANSWER;
-           PUT_32BIT(ret + 5, nkeys);
-           p = ret + 5 + 4;
-           for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
-               blob = key->alg->public_blob(key->data, &bloblen);
-               PUT_32BIT(p, bloblen);
-               p += 4;
-               memcpy(p, blob, bloblen);
-               p += bloblen;
-               sfree(blob);
-               PUT_32BIT(p, strlen(key->comment));
-               memcpy(p + 4, key->comment, strlen(key->comment));
-               p += 4 + strlen(key->comment);
+           keylist = make_keylist2(&len);
+           if (len + 5 > AGENT_MAX_MSGLEN) {
+               sfree(keylist);
+               goto failure;
            }
+           PUT_32BIT(ret, len + 1);
+           memcpy(ret + 5, keylist, len);
+           sfree(keylist);
        }
        break;
       case SSH1_AGENTC_RSA_CHALLENGE:
@@ -1267,6 +1478,59 @@ static BOOL AddTrayIcon(HWND hwnd)
     return res;
 }
 
+/* Update the saved-sessions menu. */
+static void update_sessions(void)
+{
+    int num_entries;
+    HKEY hkey;
+    TCHAR buf[MAX_PATH + 1];
+    MENUITEMINFO mii;
+
+    int index_key, index_menu;
+
+    if (!putty_path)
+       return;
+
+    if(ERROR_SUCCESS != RegOpenKey(HKEY_CURRENT_USER, PUTTY_REGKEY, &hkey))
+       return;
+
+    for(num_entries = GetMenuItemCount(session_menu);
+       num_entries > initial_menuitems_count;
+       num_entries--)
+       RemoveMenu(session_menu, 0, MF_BYPOSITION);
+
+    index_key = 0;
+    index_menu = 0;
+
+    while(ERROR_SUCCESS == RegEnumKey(hkey, index_key, buf, MAX_PATH)) {
+       TCHAR session_name[MAX_PATH + 1];
+       unmungestr(buf, session_name, MAX_PATH);
+       if(strcmp(buf, PUTTY_DEFAULT) != 0) {
+           memset(&mii, 0, sizeof(mii));
+           mii.cbSize = sizeof(mii);
+           mii.fMask = MIIM_TYPE | MIIM_STATE | MIIM_ID;
+           mii.fType = MFT_STRING;
+           mii.fState = MFS_ENABLED;
+           mii.wID = (index_menu * 16) + IDM_SESSIONS_BASE;
+           mii.dwTypeData = session_name;
+           InsertMenuItem(session_menu, index_menu, TRUE, &mii);
+           index_menu++;
+       }
+       index_key++;
+    }
+
+    RegCloseKey(hkey);
+
+    if(index_menu == 0) {
+       mii.cbSize = sizeof(mii);
+       mii.fMask = MIIM_TYPE | MIIM_STATE;
+       mii.fType = MFT_STRING;
+       mii.fState = MFS_GRAYED;
+       mii.dwTypeData = _T("(No sessions)");
+       InsertMenuItem(session_menu, index_menu, TRUE, &mii);
+    }
+}
+
 static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
                                WPARAM wParam, LPARAM lParam)
 {
@@ -1301,6 +1565,7 @@ static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
       case WM_SYSTRAY2:
        if (!menuinprogress) {
            menuinprogress = 1;
+           update_sessions();
            SetForegroundWindow(hwnd);
            ret = TrackPopupMenu(systray_menu,
                                 TPM_RIGHTALIGN | TPM_BOTTOMALIGN |
@@ -1312,6 +1577,13 @@ static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
       case WM_COMMAND:
       case WM_SYSCOMMAND:
        switch (wParam & ~0xF) {       /* low 4 bits reserved to Windows */
+         case IDM_PUTTY:
+           if((int)ShellExecute(hwnd, NULL, putty_path, _T(""), _T(""),
+                                SW_SHOW) <= 32) {
+               MessageBox(NULL, "Unable to execute PuTTY!",
+                          "Error", MB_OK | MB_ICONERROR);
+           }
+           break;
          case IDM_CLOSE:
            if (passphrase_box)
                SendMessage(passphrase_box, WM_CLOSE, 0, 0);
@@ -1360,6 +1632,28 @@ static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
                 requested_help = TRUE;
             }
            break;
+         default:
+           {
+               if(wParam >= IDM_SESSIONS_BASE && wParam <= IDM_SESSIONS_MAX) {
+                   MENUITEMINFO mii;
+                   TCHAR buf[MAX_PATH + 1];
+                   TCHAR param[MAX_PATH + 1];
+                   memset(&mii, 0, sizeof(mii));
+                   mii.cbSize = sizeof(mii);
+                   mii.fMask = MIIM_TYPE;
+                   mii.cch = MAX_PATH;
+                   mii.dwTypeData = buf;
+                   GetMenuItemInfo(session_menu, wParam, FALSE, &mii);
+                   strcpy(param, "@");
+                   strcat(param, mii.dwTypeData);
+                   if((int)ShellExecute(hwnd, NULL, putty_path, param,
+                                        _T(""), SW_SHOW) <= 32) {
+                       MessageBox(NULL, "Unable to execute PuTTY!", "Error",
+                                  MB_OK | MB_ICONERROR);
+                   }
+               }
+           }
+           break;
        }
        break;
       case WM_DESTROY:
@@ -1481,6 +1775,8 @@ void spawn_cmd(char *cmdline, char * args, int show)
     }
 }
 
+void cleanup_exit(int code) { exit(code); }
+
 int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
 {
     WNDCLASS wndclass;
@@ -1549,6 +1845,27 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
     }
 
     /*
+     * Look for the PuTTY binary (we will enable the saved session
+     * submenu if we find it).
+     */
+    {
+        char b[2048], *p, *q, *r;
+        FILE *fp;
+        GetModuleFileName(NULL, b, sizeof(b) - 1);
+        r = b;
+        p = strrchr(b, '\\');
+        if (p && p >= r) r = p+1;
+        q = strrchr(b, ':');
+        if (q && q >= r) r = q+1;
+        strcpy(r, "putty.exe");
+        if ( (fp = fopen(b, "r")) != NULL) {
+            putty_path = dupstr(b);
+            fclose(fp);
+        } else
+            putty_path = NULL;
+    }
+
+    /*
      * Find out if Pageant is already running.
      */
     already_running = FALSE;
@@ -1581,8 +1898,15 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
        /* Set up a system tray icon */
        AddTrayIcon(main_hwnd);
 
+        /* Accelerators used: nsvkxa */
         systray_menu = CreatePopupMenu();
-        /* accelerators used: vkxa */
+       if (putty_path) {
+           session_menu = CreateMenu();
+           AppendMenu(systray_menu, MF_ENABLED, IDM_PUTTY, "&New Session");
+           AppendMenu(systray_menu, MF_POPUP | MF_ENABLED,
+                      (UINT) session_menu, "&Saved Sessions");
+           AppendMenu(systray_menu, MF_SEPARATOR, 0, 0);
+       }
         AppendMenu(systray_menu, MF_ENABLED, IDM_VIEWKEYS,
                "&View Keys");
         AppendMenu(systray_menu, MF_ENABLED, IDM_ADDKEY, "Add &Key");
@@ -1592,6 +1916,7 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
         AppendMenu(systray_menu, MF_ENABLED, IDM_ABOUT, "&About");
        AppendMenu(systray_menu, MF_SEPARATOR, 0, 0);
         AppendMenu(systray_menu, MF_ENABLED, IDM_CLOSE, "E&xit");
+       initial_menuitems_count = GetMenuItemCount(session_menu);
 
        ShowWindow(main_hwnd, SW_HIDE);
 
@@ -1710,5 +2035,5 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
 
     if (advapi)
        FreeLibrary(advapi);
-    exit(msg.wParam);
+    return msg.wParam;
 }