Refactor space allocation for base64 encoding/decoding in pgcrypto.

Instead of trying to accurately calculate the space needed, use a StringInfo
that's enlarged as needed. This is just moving things around currently - the
old code was not wrong - but this is in preparation for a patch that adds
support for extra armor headers, and would make the space calculation more
complicated.

Marko Tiikkaja

diff --git a/contrib/pgcrypto/pgp-armor.c b/contrib/pgcrypto/pgp-armor.c
index 40f20550ea1333906a7711b28867eef16e7da473..ec647f0f3f2cb9068fd25dd83854cca457f578bb 100644 (file)
--- a/contrib/pgcrypto/pgp-armor.c
+++ b/contrib/pgcrypto/pgp-armor.c
@@ -203,38 +203,33 @@ crc24(const uint8 *data, unsigned len)
        return crc & 0xffffffL;
 }
 
-int
-pgp_armor_encode(const uint8 *src, unsigned len, uint8 *dst)
+void
+pgp_armor_encode(const uint8 *src, int len, StringInfo dst)
 {
-       int                     n;
-       uint8      *pos = dst;
+       int                     res;
+       unsigned        b64len;
        unsigned        crc = crc24(src, len);
 
-       n = strlen(armor_header);
-       memcpy(pos, armor_header, n);
-       pos += n;
-
-       n = b64_encode(src, len, pos);
-       pos += n;
+       appendStringInfoString(dst, armor_header);
 
-       if (*(pos - 1) != '\n')
-               *pos++ = '\n';
+       /* make sure we have enough room to b64_encode() */
+       b64len = b64_enc_len(len);
+       enlargeStringInfo(dst, (int) b64len);
+       res = b64_encode(src, len, (uint8 *) dst->data + dst->len);
+       if (res > b64len)
+               elog(FATAL, "overflow - encode estimate too small");
+       dst->len += res;
 
-       *pos++ = '=';
-       pos[3] = _base64[crc & 0x3f];
-       crc >>= 6;
-       pos[2] = _base64[crc & 0x3f];
-       crc >>= 6;
-       pos[1] = _base64[crc & 0x3f];
-       crc >>= 6;
-       pos[0] = _base64[crc & 0x3f];
-       pos += 4;
+       if (*(dst->data + dst->len - 1) != '\n')
+               appendStringInfoChar(dst, '\n');
 
-       n = strlen(armor_footer);
-       memcpy(pos, armor_footer, n);
-       pos += n;
+       appendStringInfoChar(dst, '=');
+       appendStringInfoChar(dst, _base64[(crc >> 18) & 0x3f]);
+       appendStringInfoChar(dst, _base64[(crc >> 12) & 0x3f]);
+       appendStringInfoChar(dst, _base64[(crc >> 6) & 0x3f]);
+       appendStringInfoChar(dst, _base64[crc & 0x3f]);
 
-       return pos - dst;
+       appendStringInfoString(dst, armor_footer);
 }
 
 static const uint8 *
@@ -309,7 +304,7 @@ find_header(const uint8 *data, const uint8 *datend,
 }
 
 int
-pgp_armor_decode(const uint8 *src, unsigned len, uint8 *dst)
+pgp_armor_decode(const uint8 *src, int len, StringInfo dst)
 {
        const uint8 *p = src;
        const uint8 *data_end = src + len;
@@ -319,6 +314,7 @@ pgp_armor_decode(const uint8 *src, unsigned len, uint8 *dst)
        const uint8 *base64_end = NULL;
        uint8           buf[4];
        int                     hlen;
+       int                     blen;
        int                     res = PXE_PGP_CORRUPT_ARMOR;
 
        /* armor start */
@@ -360,23 +356,18 @@ pgp_armor_decode(const uint8 *src, unsigned len, uint8 *dst)
        crc = (((long) buf[0]) << 16) + (((long) buf[1]) << 8) + (long) buf[2];
 
        /* decode data */
-       res = b64_decode(base64_start, base64_end - base64_start, dst);
-
-       /* check crc */
-       if (res >= 0 && crc24(dst, res) != crc)
-               res = PXE_PGP_CORRUPT_ARMOR;
+       blen = (int) b64_dec_len(len);
+       enlargeStringInfo(dst, blen);
+       res = b64_decode(base64_start, base64_end - base64_start, (uint8 *) dst->data);
+       if (res > blen)
+               elog(FATAL, "overflow - decode estimate too small");
+       if (res >= 0)
+       {
+               if (crc24((uint8 *) dst->data, res) == crc)
+                       dst->len += res;
+               else
+                       res = PXE_PGP_CORRUPT_ARMOR;
+       }
 out:
        return res;
 }
-
-unsigned
-pgp_armor_enc_len(unsigned len)
-{
-       return b64_enc_len(len) + strlen(armor_header) + strlen(armor_footer) + 16;
-}
-
-unsigned
-pgp_armor_dec_len(unsigned len)
-{
-       return b64_dec_len(len);
-}

diff --git a/contrib/pgcrypto/pgp-pgsql.c b/contrib/pgcrypto/pgp-pgsql.c
index ad1fd084276a356a622f84aea85b5c75789d4a43..5d2d4655d18b0a4549235c3d16d1bbbc02defbf9 100644 (file)
--- a/contrib/pgcrypto/pgp-pgsql.c
+++ b/contrib/pgcrypto/pgp-pgsql.c
@@ -31,6 +31,7 @@
 
 #include "postgres.h"
 
+#include "lib/stringinfo.h"
 #include "mb/pg_wchar.h"
 #include "utils/builtins.h"
 
@@ -820,23 +821,20 @@ pg_armor(PG_FUNCTION_ARGS)
 {
        bytea      *data;
        text       *res;
-       int                     data_len,
-                               res_len,
-                               guess_len;
+       int                     data_len;
+       StringInfoData buf;
 
        data = PG_GETARG_BYTEA_P(0);
        data_len = VARSIZE(data) - VARHDRSZ;
 
-       guess_len = pgp_armor_enc_len(data_len);
-       res = palloc(VARHDRSZ + guess_len);
+       initStringInfo(&buf);
 
-       res_len = pgp_armor_encode((uint8 *) VARDATA(data), data_len,
-                                                          (uint8 *) VARDATA(res));
-       if (res_len > guess_len)
-               ereport(ERROR,
-                               (errcode(ERRCODE_EXTERNAL_ROUTINE_INVOCATION_EXCEPTION),
-                                errmsg("Overflow - encode estimate too small")));
-       SET_VARSIZE(res, VARHDRSZ + res_len);
+       pgp_armor_encode((uint8 *) VARDATA(data), data_len, &buf);
+
+       res = palloc(VARHDRSZ + buf.len);
+       SET_VARSIZE(res, VARHDRSZ + buf.len);
+       memcpy(VARDATA(res), buf.data, buf.len);
+       pfree(buf.data);
 
        PG_FREE_IF_COPY(data, 0);
        PG_RETURN_TEXT_P(res);
@@ -847,27 +845,24 @@ pg_dearmor(PG_FUNCTION_ARGS)
 {
        text       *data;
        bytea      *res;
-       int                     data_len,
-                               res_len,
-                               guess_len;
+       int                     data_len;
+       int                     ret;
+       StringInfoData buf;
 
        data = PG_GETARG_TEXT_P(0);
        data_len = VARSIZE(data) - VARHDRSZ;
 
-       guess_len = pgp_armor_dec_len(data_len);
-       res = palloc(VARHDRSZ + guess_len);
+       initStringInfo(&buf);
 
-       res_len = pgp_armor_decode((uint8 *) VARDATA(data), data_len,
-                                                          (uint8 *) VARDATA(res));
-       if (res_len < 0)
+       ret = pgp_armor_decode((uint8 *) VARDATA(data), data_len, &buf);
+       if (ret < 0)
                ereport(ERROR,
                                (errcode(ERRCODE_EXTERNAL_ROUTINE_INVOCATION_EXCEPTION),
-                                errmsg("%s", px_strerror(res_len))));
-       if (res_len > guess_len)
-               ereport(ERROR,
-                               (errcode(ERRCODE_EXTERNAL_ROUTINE_INVOCATION_EXCEPTION),
-                                errmsg("Overflow - decode estimate too small")));
-       SET_VARSIZE(res, VARHDRSZ + res_len);
+                                errmsg("%s", px_strerror(ret))));
+       res = palloc(VARHDRSZ + buf.len);
+       SET_VARSIZE(res, VARHDRSZ + buf.len);
+       memcpy(VARDATA(res), buf.data, buf.len);
+       pfree(buf.data);
 
        PG_FREE_IF_COPY(data, 0);
        PG_RETURN_TEXT_P(res);

diff --git a/contrib/pgcrypto/pgp.h b/contrib/pgcrypto/pgp.h
index 8d4ab9862dfd5dc3da577f47f4a02ca3cf622874..cecd1814956db4d6db465e0d2d5f6b8db155258c 100644 (file)
--- a/contrib/pgcrypto/pgp.h
+++ b/contrib/pgcrypto/pgp.h
@@ -29,6 +29,8 @@
  * contrib/pgcrypto/pgp.h
  */
 
+#include "lib/stringinfo.h"
+
 #include "mbuf.h"
 #include "px.h"
 
@@ -274,10 +276,8 @@ void               pgp_cfb_free(PGP_CFB *ctx);
 int                    pgp_cfb_encrypt(PGP_CFB *ctx, const uint8 *data, int len, uint8 *dst);
 int                    pgp_cfb_decrypt(PGP_CFB *ctx, const uint8 *data, int len, uint8 *dst);
 
-int                    pgp_armor_encode(const uint8 *src, unsigned len, uint8 *dst);
-int                    pgp_armor_decode(const uint8 *src, unsigned len, uint8 *dst);
-unsigned       pgp_armor_enc_len(unsigned len);
-unsigned       pgp_armor_dec_len(unsigned len);
+void           pgp_armor_encode(const uint8 *src, int len, StringInfo dst);
+int                    pgp_armor_decode(const uint8 *src, int len, StringInfo dst);
 
 int                    pgp_compress_filter(PushFilter **res, PGP_Context *ctx, PushFilter *dst);
 int                    pgp_decompress_filter(PullFilter **res, PGP_Context *ctx, PullFilter *src);