From d66b097ddc5a7fbaf1998d6516848405d8a6ab7e Mon Sep 17 00:00:00 2001
From: phoebe <phoreverphoebs@proton.me>
Date: Thu, 20 Jul 2023 01:34:08 +0200
Subject: [PATCH] crypto.pem: add decode_only and general improvements to
 decoding (#18908)

---
 vlib/crypto/pem/decode.v   | 31 ++++++++++++++++++++++++++++---
 vlib/crypto/pem/pem_test.v | 12 +++++++++---
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/vlib/crypto/pem/decode.v b/vlib/crypto/pem/decode.v
index 350222e30dd8b6..ba867048865248 100644
--- a/vlib/crypto/pem/decode.v
+++ b/vlib/crypto/pem/decode.v
@@ -2,25 +2,50 @@ module pem
 
 import encoding.base64
 
+// decode_only reads `data` and returns the first parsed PEM Block. `none` is returned
+//  when a header is expected, but not present or when a start of '-----BEGIN' or end of '-----END'
+// can't be found.
+//
+// use decode if you still need the unparsed rest of the string.
+[inline]
+pub fn decode_only(data string) ?Block {
+	block, _ := decode_internal(data)?
+	return block
+}
+
 // decode reads `data` and returns the first parsed PEM Block along with the rest of
 // the string. `none` is returned when a header is expected, but not present
-// or when a start of '-----BEGIN' or end of '-----END' can't be found in `data`
+// or when a start of '-----BEGIN' or end of '-----END' can't be found.
+//
+// use decode_only if you do not need the unparsed rest of the string.
+[direct_array_access; inline]
 pub fn decode(data string) ?(Block, string) {
+	block, rest := decode_internal(data)?
+	return block, rest[rest.index(pem_end)? + pem_end.len..].all_after_first(pem_eol)
+}
+
+// decode_internal allows `decode` variations to deal with the rest of the data as
+// they want to. for example Block.decode could have hindered performance with the final
+// indexing into `rest` that `decode_partial` does.
+[direct_array_access]
+fn decode_internal(data string) ?(Block, string) {
+	// direct_array_access safety: since we use the string.index method here,
+	// we won't get an invalid index since it would otherwise return `none`
 	mut rest := data[data.index(pem_begin)?..]
 	mut block := Block.new(rest[pem_begin.len..].all_before(pem_eol))
 	block.headers, rest = parse_headers(rest[pem_begin.len..].all_after(pem_eol).trim_left(' \n\t\v\f\r'))?
 
 	block_end_index := rest.index(pem_end)?
 	b64_data := rest[..block_end_index].replace_each(['\r', '', '\n', '', '\t', '', ' ', ''])
-
 	block_data_len := block_end_index / 4 * 3
 	block.data = []u8{len: block_data_len, cap: block_data_len + 3, init: 0}
 	decoded_len := base64.decode_in_buffer(&b64_data, &block.data[0])
 	block.data = block.data[..decoded_len]
 
-	return block, rest[rest.index(pem_end)? + pem_end.len..].all_after_first(pem_eol)
+	return block, rest
 }
 
+[direct_array_access]
 fn parse_headers(block string) ?(map[string][]string, string) {
 	headers_str := block.all_before(pem_end).all_before('\n\n')
 
diff --git a/vlib/crypto/pem/pem_test.v b/vlib/crypto/pem/pem_test.v
index 8f02650b150241..fdfbb165d5d878 100644
--- a/vlib/crypto/pem/pem_test.v
+++ b/vlib/crypto/pem/pem_test.v
@@ -5,6 +5,7 @@ fn test_decode_rfc1421() {
 	for i in 0 .. pem.test_data_rfc1421.len {
 		decoded, rest := decode(pem.test_data_rfc1421[i]) or { Block{}, '' }
 		assert decoded == pem.expected_results_rfc1421[i]
+		assert decoded == decode_only(pem.test_data_rfc1421[i]) or { Block{} }
 		assert rest == ''
 	}
 }
@@ -13,6 +14,7 @@ fn test_decode() {
 	for i in 0 .. pem.test_data.len {
 		decoded, rest := decode(pem.test_data[i]) or { Block{}, '' }
 		assert decoded == pem.expected_results[i]
+		assert decoded == decode_only(pem.test_data[i]) or { Block{} }
 		assert rest == pem.expected_rest[i]
 	}
 }
@@ -23,6 +25,7 @@ fn test_encode_rfc1421() {
 		decoded, rest := decode(encoded) or { Block{}, '' }
 		assert rest == ''
 		assert decoded == pem.expected_results_rfc1421[i]
+		assert decoded == decode_only(encoded) or { Block{} }
 	}
 }
 
@@ -32,6 +35,7 @@ fn test_encode() {
 		decoded, rest := decode(encoded) or { Block{}, '' }
 		assert rest == ''
 		assert decoded == pem.expected_results[i]
+		assert decoded == decode_only(encoded) or { Block{} }
 	}
 }
 
@@ -41,15 +45,17 @@ fn test_encode_config() {
 		decoded, rest := decode(encoded) or { Block{}, '' }
 		assert rest == ''
 		assert decoded == pem.expected_results[i]
+		assert decoded == decode_only(encoded) or { Block{} }
 	}
 }
 
 fn test_decode_no_pem() {
 	for test in pem.test_data_no_pem {
 		if _, _ := decode(test) {
-			assert false, 'Block.decode_partial should return `none` on input without PEM data'
-		} else {
-			assert true
+			assert false, 'decode should return `none` on input without PEM data'
+		}
+		if _ := decode_only(test) {
+			assert false, 'decode_only should return `none` on input without PEM data'
 		}
 	}
 }