mbedtls_pem_write_buffer: Correctly report needed buffer size for all possible line lengths and counts

Signed-off-by: Jethro Beekman <jethro@fortanix.com>
diff --git a/library/pem.c b/library/pem.c
index 8c8f39d..aed4788 100644
--- a/library/pem.c
+++ b/library/pem.c
@@ -491,7 +491,7 @@
     size_t len = 0, use_len, add_len = 0;
 
     mbedtls_base64_encode(NULL, 0, &use_len, der_data, der_len);
-    add_len = strlen(header) + strlen(footer) + (use_len / 64) + 1;
+    add_len = strlen(header) + strlen(footer) + (((use_len > 2) ? (use_len - 2) : 0) / 64) + 1;
 
     if (use_len + add_len > buf_len) {
         *olen = use_len + add_len;
diff --git a/tests/suites/test_suite_pem.data b/tests/suites/test_suite_pem.data
index 332cb79..a34aa91 100644
--- a/tests/suites/test_suite_pem.data
+++ b/tests/suites/test_suite_pem.data
@@ -16,6 +16,9 @@
 PEM write (exactly two lines + 1)
 mbedtls_pem_write_buffer:"-----START TEST-----\n":"-----END TEST-----\n":"000102030405060708090A0B0C0D0E0F000102030405060708090A0B0C0D0E0F000102030405060708090A0B0C0D0E0F000102030405060708090A0B0C0D0E0F000102030405060708090A0B0C0D0E0F000102030405060708090A0B0C0D0E0F00":"-----START TEST-----\nAAECAwQFBgcICQoLDA0ODwABAgMEBQYHCAkKCwwNDg8AAQIDBAUGBwgJCgsMDQ4P\nAAECAwQFBgcICQoLDA0ODwABAgMEBQYHCAkKCwwNDg8AAQIDBAUGBwgJCgsMDQ4P\nAA==\n-----END TEST-----\n"
 
+PEM write length reporting
+mbedtls_pem_write_buffer_lengths
+
 PEM read (unencrypted, valid)
 mbedtls_pem_read_buffer:"^":"$":"^\nTWJlZCBUTFM=\n$":"":0:"4d62656420544c53"
 
diff --git a/tests/suites/test_suite_pem.function b/tests/suites/test_suite_pem.function
index 918eae5..36801d1 100644
--- a/tests/suites/test_suite_pem.function
+++ b/tests/suites/test_suite_pem.function
@@ -32,6 +32,35 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_PEM_WRITE_C */
+void mbedtls_pem_write_buffer_lengths()
+{
+    unsigned char data[256] = {};
+    unsigned char buf[1024];
+    size_t olen_needed, olen;
+    int ret;
+    for (size_t l = 0; l <= sizeof(data); l++) {
+        ret = mbedtls_pem_write_buffer("\n", "\n", data, l, NULL, 0, &olen_needed);
+        TEST_ASSERT(ret == MBEDTLS_ERR_BASE64_BUFFER_TOO_SMALL);
+
+        /* Test that a bigger buffer still only requires `olen_needed` */
+        ret = mbedtls_pem_write_buffer("\n", "\n", data, l, buf, sizeof(buf), &olen);
+        TEST_ASSERT(ret == 0);
+        TEST_ASSERT(olen_needed == olen);
+
+        /* Test that a buffer of exactly `olen_needed` works */
+        memset(buf, 1, sizeof(buf));
+        ret = mbedtls_pem_write_buffer("\n", "\n", data, l, buf, olen_needed, &olen);
+        TEST_ASSERT(ret == 0);
+        TEST_ASSERT(olen_needed == olen);
+        /* Test the function didn't overflow the given buffer */
+        for (size_t i = olen_needed; i < sizeof(buf); i++) {
+            TEST_ASSERT(buf[i] == 1);
+        }
+    }
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_PEM_PARSE_C */
 void mbedtls_pem_read_buffer(char *header, char *footer, char *data,
                              char *pwd, int res, data_t *out)