#include "pub.h"

static const char base64[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

static const char index_64[128] =
{
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 62, -1, -1, -1, 63,
    52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -1, -1, -1, -1, -1, -1,
    -1,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
    15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -1, -1, -1, -1, -1,
    -1, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
    41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -1, -1, -1, -1, -1,
};

int base64_encode(char *dst, int dstlen, const char *src, int srclen)
{
    char *to = dst;
    const unsigned char *from;
    unsigned char c1, c2;
    int dst_needed;

    // 参数检测
    if (dst == NULL || dstlen <= 0 || src == NULL || srclen <= 0)
    {
        return MS_ERROR;
    }

    /* dst 缓冲区的空间不足。
     * Base64 编码是每 3 个原始字符编码成 4 个字符，
     * 如果原始字符串长度不能被 3 整除，使用 0 来补充原始字符串。*/
    dst_needed = (srclen + 2) / 3;
    dst_needed *= 4;
    if (dstlen < dst_needed + 1)
    {
        return MS_ERROR;
    }

    from = (unsigned char *)src;
    while (srclen > 0)
    {
        c1 = *from++;
        srclen--;
        *to++ = base64[c1 >> 2];
        c1 = (c1 & 0x03) << 4;
        if (srclen <= 0)
        {
            *to++ = base64[c1];
            *to++ = '=';
            *to++ = '=';
            break;
        }

        c2 = *from++;
        srclen--;
        c1 |= (c2 >> 4) & 0x0f;
        *to++ = base64[c1];
        c1 = (c2 & 0x0f) << 2;
        if (srclen <= 0)
        {
            *to++ = base64[c1];
            *to++ = '=';
            break;
        }

        c2 = *from++;
        srclen--;
        c1 |= (c2 >> 6) & 0x03;
        *to++ = base64[c1];
        *to++ = base64[c2 & 0x3f];
    }
    *to = '\0';
    return to - dst;
}

int base64_decode(char *dst, int dstlen, const char *src, int srclen)
{
    const unsigned char *p, *q;
    unsigned char *t;
    int c1, c2;

    // 参数检测
    if (dst == NULL || dstlen <= 0 || src == NULL || srclen <= 0)
    {
        return MS_ERROR;
    }

    // 移除首尾的空白字符
    for (p = (const unsigned char *) src; srclen > 0 && isspace(*p); p++, srclen--)
    {
        /* void */ ;
    }
    for (q = p + srclen - 1; q >= p && isspace(*q); q--, srclen--)
    {
        /* void */ ;
    }

    // 长度必须为 4 的倍数
    if (srclen % 4 != 0)
    {
        return MS_ERROR;
    }

    // 目标缓冲区必须足够长
    if (srclen / 4 * 3 + 1 > dstlen)
    {
        return MS_ERROR;
    }

    t = (unsigned char *)dst;
    while (srclen > 0)
    {
        srclen -= 4;
        if (*p >= 128 || (c1 = index_64[*p++]) == -1)
        {
            return MS_ERROR;
        }

        if (*p >= 128 || (c2 = index_64[*p++]) == -1)
        {
            return MS_ERROR;
        }
        *t++ = (c1 << 2) | ((c2 & 0x30) >> 4);

        if (p[0] == '=' && p[1] == '=')
        {
            break;
        }

        if (*p >= 128 || (c1 = index_64[*p++]) == -1)
        {
            return MS_ERROR;
        }
        *t++ = ((c2 & 0x0f) << 4) | ((c1 & 0x3c) >> 2);

        if (p[0] == '=')
        {
            break;
        }

        if (*p >= 128 || (c2 = index_64[*p++]) == -1)
        {
            return MS_ERROR;
        }
        *t++ = ((c1 & 0x03) << 6) | c2;
    }

    return t - (unsigned char *) dst;
}