base64_encode/base64_decode should have a similar interface

Prepare functions so they can be exported in the public interface
See #1038

src/civetweb.c

@@ -7216,14 +7216,27 @@ mg_get_cookie(const char *cookie_header,
 }
 
 
-#if defined(USE_WEBSOCKET) || defined(USE_LUA)
-static void
-base64_encode(const unsigned char *src, int src_len, char *dst)
+static int
+base64_encode(const unsigned char *src, int src_len, char *dst, size_t *dst_len)
 {
 	static const char *b64 =
 	    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
 	int i, j, a, b, c;
 
+	if (dst_len != NULL) {
+		/* Expected length including 0 termination: */
+		/* IN 1 -> OUT 5, IN 2 -> OUT 5, IN 3 -> OUT 5, IN 4 -> OUT 9,
+		 * IN 5 -> OUT 9, IN 6 -> OUT 9, IN 7 -> OUT 13, etc. */
+		size_t expected_len = ((src_len + 2) / 3) * 4 + 1;
+		if (*dst_len < expected_len) {
+			if (*dst_len > 0) {
+				dst[0] = '\0';
+			}
+			*dst_len = expected_len;
+			return 0;
+		}
+	}
+
 	for (i = j = 0; i < src_len; i += 3) {
 		a = src[i];
 		b = ((i + 1) >= src_len) ? 0 : src[i + 1];
@@ -7242,11 +7255,16 @@ base64_encode(const unsigned char *src, int src_len, char *dst)
 		dst[j++] = '=';
 	}
 	dst[j++] = '\0';
+
+	if (dst_len != NULL) {
+		*dst_len = (size_t)j;
+	}
+
+	/* Return -1 for "OK" */
+	return -1;
 }
-#endif
 
 
-#if defined(USE_LUA)
 static unsigned char
 b64reverse(char letter)
 {
@@ -7277,10 +7295,16 @@ base64_decode(const unsigned char *src, int src_len, char *dst, size_t *dst_len)
 {
 	int i;
 	unsigned char a, b, c, d;
+	size_t dst_len_limit = (size_t)-1;
+	size_t dst_len_used = 0;
 
-	*dst_len = 0;
+	if (dst_len != NULL) {
+		dst_len_limit = *dst_len;
+		*dst_len = 0;
+	}
 
 	for (i = 0; i < src_len; i += 4) {
+		/* Read 4 characters from BASE64 string */
 		a = b64reverse(src[i]);
 		if (a >= 254) {
 			return i;
@@ -7301,17 +7325,44 @@ base64_decode(const unsigned char *src, int src_len, char *dst, size_t *dst_len)
 			return i + 3;
 		}
 
-		dst[(*dst_len)++] = (a << 2) + (b >> 4);
+		/* Add first (of 3) decoded character */
+		if (dst_len_used < dst_len_limit) {
+			dst[dst_len_used] = (a << 2) + (b >> 4);
+		}
+		dst_len_used++;
+
 		if (c != 255) {
-			dst[(*dst_len)++] = (b << 4) + (c >> 2);
+			if (dst_len_used < dst_len_limit) {
+
+				dst[dst_len_used] = (b << 4) + (c >> 2);
+			}
+			dst_len_used++;
 			if (d != 255) {
-				dst[(*dst_len)++] = (c << 6) + d;
+				if (dst_len_used < dst_len_limit) {
+					dst[dst_len_used] = (c << 6) + d;
+				}
+				dst_len_used++;
 			}
 		}
 	}
+
+	/* Add terminating zero */
+	if (dst_len_used < dst_len_limit) {
+		dst[dst_len_used] = '\0';
+	}
+	dst_len_used++;
+	if (dst_len != NULL) {
+		*dst_len = dst_len_used;
+	}
+
+	if (dst_len_used > dst_len_limit) {
+		/* Out of memory */
+		return 0;
+	}
+
+	/* Return -1 for "OK" */
 	return -1;
 }
-#endif
 
 
 static int
@@ -12379,6 +12430,7 @@ send_websocket_handshake(struct mg_connection *conn, const char *websock_key)
 {
 	static const char *magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
 	char buf[100], sha[20], b64_sha[sizeof(sha) * 2];
+	size_t dst_len = sizeof(b64_sha);
 	SHA_CTX sha_ctx;
 	int truncated;
 
@@ -12394,7 +12446,7 @@ send_websocket_handshake(struct mg_connection *conn, const char *websock_key)
 	SHA1_Init(&sha_ctx);
 	SHA1_Update(&sha_ctx, (unsigned char *)buf, (uint32_t)strlen(buf));
 	SHA1_Final((unsigned char *)sha, &sha_ctx);
-	base64_encode((unsigned char *)sha, sizeof(sha), b64_sha);
+	base64_encode((unsigned char *)sha, sizeof(sha), b64_sha, &dst_len);
 	mg_printf(conn,
 	          "HTTP/1.1 101 Switching Protocols\r\n"
 	          "Upgrade: websocket\r\n"

src/mod_lua.inl

@@ -1506,7 +1506,6 @@ lsp_base64_encode(lua_State *L)
 	int num_args = lua_gettop(L);
 	const char *text;
 	size_t text_len;
-	char *dst;
 	struct mg_context *ctx;
 
 	lua_pushlightuserdata(L, (void *)&lua_regkey_ctx);
@@ -1516,10 +1515,15 @@ lsp_base64_encode(lua_State *L)
 	if (num_args == 1) {
 		text = lua_tolstring(L, 1, &text_len);
 		if (text) {
-			dst = (char *)mg_malloc_ctx(text_len * 8 / 6 + 4, ctx);
+			/* Base 64 encodes 8 bits into 6 */
+			size_t dst_len = text_len * 8 / 6 + 4;
+			char *dst = (char *)mg_malloc_ctx(dst_len, ctx);
 			if (dst) {
-				base64_encode((const unsigned char *)text, (int)text_len, dst);
-				lua_pushstring(L, dst);
+				base64_encode((const unsigned char *)text,
+				              (int)text_len,
+				              dst,
+				              &dst_len);
+				lua_pushlstring(L, dst, dst_len);
 				mg_free(dst);
 			} else {
 				return luaL_error(L, "out of memory in base64_encode() call");

unittest/private.c

@@ -964,7 +964,6 @@ START_TEST(test_encode_decode)
 	int ret;
 	size_t len;
 
-#if defined(USE_WEBSOCKET) || defined(USE_LUA)
 	const char *alpha_b64_enc = "YWJjZGVmZ2hpamtsbW5vcHFyc3R1dnd4eXo=";
 	const char *nonalpha_b64_enc =
 	    "ICEiIyQlJicoKSorLC0uLzAxMjM0NTY3ODk6Ozw9Pj9A";
@@ -972,43 +971,57 @@ START_TEST(test_encode_decode)
 	mark_point();
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)"a", 1, buf);
+	base64_encode((unsigned char *)"a", 1, buf, NULL);
 	ck_assert_str_eq(buf, "YQ==");
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)"ab", 1, buf);
+	len = 9999;
+	base64_encode((unsigned char *)"ab", 1, buf, &len);
 	ck_assert_str_eq(buf, "YQ==");
+	ck_assert_int_eq((int)len, 5);
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)"ab", 2, buf);
+	len = 9999;
+	base64_encode((unsigned char *)"ab", 2, buf, &len);
 	ck_assert_str_eq(buf, "YWI=");
+	ck_assert_int_eq((int)len, 5);
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)alpha, 3, buf);
+	len = 9999;
+	base64_encode((unsigned char *)alpha, 3, buf, &len);
 	ck_assert_str_eq(buf, "YWJj");
+	ck_assert_int_eq((int)len, 5);
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)alpha, 4, buf);
+	len = 9999;
+	base64_encode((unsigned char *)alpha, 4, buf, &len);
 	ck_assert_str_eq(buf, "YWJjZA==");
+	ck_assert_int_eq((int)len, 9);
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)alpha, 5, buf);
+	len = 9999;
+	base64_encode((unsigned char *)alpha, 5, buf, &len);
 	ck_assert_str_eq(buf, "YWJjZGU=");
+	ck_assert_int_eq((int)len, 9);
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)alpha, 6, buf);
+	len = 9999;
+	base64_encode((unsigned char *)alpha, 6, buf, &len);
 	ck_assert_str_eq(buf, "YWJjZGVm");
+	ck_assert_int_eq((int)len, 9);
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)alpha, (int)strlen(alpha), buf);
+	len = 9999;
+	base64_encode((unsigned char *)alpha, (int)strlen(alpha), buf, &len);
 	ck_assert_str_eq(buf, alpha_b64_enc);
+	ck_assert_int_eq((int)len, (int)strlen(alpha_b64_enc) + 1);
 
 	memset(buf, 77, sizeof(buf));
-	base64_encode((unsigned char *)nonalpha, (int)strlen(nonalpha), buf);
+	len = 9999;
+	base64_encode((unsigned char *)nonalpha, (int)strlen(nonalpha), buf, &len);
 	ck_assert_str_eq(buf, nonalpha_b64_enc);
-#endif
+	ck_assert_int_eq((int)len, (int)strlen(nonalpha_b64_enc) + 1);
 
-#if defined(USE_LUA)
 	memset(buf, 77, sizeof(buf));
 	len = 9999;
 	ret = base64_decode((unsigned char *)alpha_b64_enc,
@@ -1016,14 +1029,13 @@ START_TEST(test_encode_decode)
 	                    buf,
 	                    &len);
 	ck_assert_int_eq(ret, -1);
-	ck_assert_uint_eq((unsigned int)len, (unsigned int)strlen(alpha));
+	ck_assert_int_eq((int)len, (int)strlen(alpha) + 1);
 	ck_assert_str_eq(buf, alpha);
 
 	memset(buf, 77, sizeof(buf));
 	len = 9999;
 	ret = base64_decode((unsigned char *)"AAA*AAA", 7, buf, &len);
 	ck_assert_int_eq(ret, 3);
-#endif
 
 	memset(buf, 77, sizeof(buf));
 	ret = mg_url_encode(alpha, buf, sizeof(buf));