libndr: Return error code from ndr_token_peek()

This makes it safer to change our code to remove tokens after use
if failing to obtain a token would result in an error.

This means changing ndr_get_array_size() and ndr_get_array_length()
to also return an error code.

BUG: https://bugzilla.samba.org/show_bug.cgi?id=14710

Signed-off-by: Andrew Bartlett <abartlet@samba.org>
Reviewed-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>

4 files changed, 23 insertions, 28 deletions

diff --git a/librpc/ABI/ndr-2.0.0.sigs b/librpc/ABI/ndr-2.0.0.sigs
index b06431b3c2c..5089be76f30 100644
--- a/librpc/ABI/ndr-2.0.0.sigs
+++ b/librpc/ABI/ndr-2.0.0.sigs
@@ -21,8 +21,8 @@ ndr_check_array_size: enum ndr_err_code (struct ndr_pull *, void *, uint32_t)
 ndr_check_padding: void (struct ndr_pull *, size_t)
 ndr_check_pipe_chunk_trailer: enum ndr_err_code (struct ndr_pull *, int, uint32_t)
 ndr_check_string_terminator: enum ndr_err_code (struct ndr_pull *, uint32_t, uint32_t)
-ndr_get_array_length: uint32_t (struct ndr_pull *, const void *)
-ndr_get_array_size: uint32_t (struct ndr_pull *, const void *)
+ndr_get_array_length: enum ndr_err_code (struct ndr_pull *, const void *, uint32_t *)
+ndr_get_array_size: enum ndr_err_code (struct ndr_pull *, const void *, uint32_t *)
 ndr_map_error2errno: int (enum ndr_err_code)
 ndr_map_error2ntstatus: NTSTATUS (enum ndr_err_code)
 ndr_map_error2string: const char *(enum ndr_err_code)
@@ -257,7 +257,7 @@ ndr_syntax_id_from_string: bool (const char *, struct ndr_syntax_id *)
 ndr_syntax_id_null: uuid = {time_low = 0, time_mid = 0, time_hi_and_version = 0, clock_seq = "\000", node = "\000\000\000\000\000"}, if_version = 0
 ndr_syntax_id_to_string: char *(TALLOC_CTX *, const struct ndr_syntax_id *)
 ndr_token_max_list_size: size_t (void)
-ndr_token_peek: uint32_t (struct ndr_token_list *, const void *)
+ndr_token_peek: enum ndr_err_code (struct ndr_token_list *, const void *, uint32_t *)
 ndr_token_retrieve: enum ndr_err_code (struct ndr_token_list *, const void *, uint32_t *)
 ndr_token_retrieve_cmp_fn: enum ndr_err_code (struct ndr_token_list *, const void *, uint32_t *, comparison_fn_t, bool)
 ndr_token_store: enum ndr_err_code (TALLOC_CTX *, struct ndr_token_list *, const void *, uint32_t)
diff --git a/librpc/ndr/libndr.h b/librpc/ndr/libndr.h
index 25b68db3466..4a13a16167c 100644
--- a/librpc/ndr/libndr.h
+++ b/librpc/ndr/libndr.h
@@ -653,12 +653,12 @@ enum ndr_err_code ndr_token_store(TALLOC_CTX *mem_ctx,
 enum ndr_err_code ndr_token_retrieve_cmp_fn(struct ndr_token_list *list, const void *key, uint32_t *v,
 					    int(*_cmp_fn)(const void*,const void*), bool erase);
 enum ndr_err_code ndr_token_retrieve(struct ndr_token_list *list, const void *key, uint32_t *v);
-uint32_t ndr_token_peek(struct ndr_token_list *list, const void *key);
+enum ndr_err_code ndr_token_peek(struct ndr_token_list *list, const void *key, uint32_t *v);
 enum ndr_err_code ndr_pull_array_size(struct ndr_pull *ndr, const void *p);
-uint32_t ndr_get_array_size(struct ndr_pull *ndr, const void *p);
+enum ndr_err_code ndr_get_array_size(struct ndr_pull *ndr, const void *p, uint32_t *size);
 enum ndr_err_code ndr_check_array_size(struct ndr_pull *ndr, void *p, uint32_t size);
 enum ndr_err_code ndr_pull_array_length(struct ndr_pull *ndr, const void *p);
-uint32_t ndr_get_array_length(struct ndr_pull *ndr, const void *p);
+enum ndr_err_code ndr_get_array_length(struct ndr_pull *ndr, const void *p, uint32_t *length);
 enum ndr_err_code ndr_check_array_length(struct ndr_pull *ndr, void *p, uint32_t length);
 enum ndr_err_code ndr_push_pipe_chunk_trailer(struct ndr_push *ndr, int ndr_flags, uint32_t count);
 enum ndr_err_code ndr_check_pipe_chunk_trailer(struct ndr_pull *ndr, int ndr_flags, uint32_t count);
diff --git a/librpc/ndr/ndr.c b/librpc/ndr/ndr.c
index c384e01e31f..115e617da4a 100644
--- a/librpc/ndr/ndr.c
+++ b/librpc/ndr/ndr.c
@@ -148,6 +148,7 @@ _PUBLIC_ enum ndr_err_code ndr_pull_pop(struct ndr_pull *ndr)
 {
 	uint32_t skip = 0;
 	uint32_t append = 0;
+	enum ndr_err_code ndr_err;
 
 	if (ndr->relative_base_offset != 0) {
 		return ndr_pull_error(ndr, NDR_ERR_RELATIVE,
@@ -179,8 +180,8 @@ _PUBLIC_ enum ndr_err_code ndr_pull_pop(struct ndr_pull *ndr)
 	ndr->offset -= skip;
 	ndr->data_size -= skip;
 
-	append = ndr_token_peek(&ndr->array_size_list, ndr);
-	if (append != UINT32_MAX) {
+	ndr_err = ndr_token_peek(&ndr->array_size_list, ndr, &append);
+	if (ndr_err == NDR_ERR_TOKEN) {
 		/*
 		 * here we assume, that ndr->data is not a
 		 * talloc child of ndr.
@@ -1065,18 +1066,10 @@ _PUBLIC_ enum ndr_err_code ndr_token_retrieve(struct ndr_token_list *list,
 /*
   peek at but don't removed a token from a ndr context
 */
-_PUBLIC_ uint32_t ndr_token_peek(struct ndr_token_list *list, const void *key)
+_PUBLIC_ enum ndr_err_code ndr_token_peek(struct ndr_token_list *list,
+					  const void *key, uint32_t *v)
 {
-	unsigned i;
-	struct ndr_token *tokens = list->tokens;
-
-	for (i = list->count - 1; i < list->count; i--) {
-		if (tokens[i].key == key) {
-			return tokens[i].value;
-		}
-	}
-
-	return 0;
+	return ndr_token_retrieve_cmp_fn(list, key, v, NULL, false);
 }
 
 /*
@@ -1099,9 +1092,9 @@ _PUBLIC_ enum ndr_err_code ndr_pull_array_size(struct ndr_pull *ndr, const void
 /*
   get the stored array size field
 */
-_PUBLIC_ uint32_t ndr_get_array_size(struct ndr_pull *ndr, const void *p)
+_PUBLIC_ enum ndr_err_code ndr_get_array_size(struct ndr_pull *ndr, const void *p, uint32_t *size)
 {
-	return ndr_token_peek(&ndr->array_size_list, p);
+	return ndr_token_peek(&ndr->array_size_list, p, size);
 }
 
 /*
@@ -1110,7 +1103,7 @@ _PUBLIC_ uint32_t ndr_get_array_size(struct ndr_pull *ndr, const void *p)
 _PUBLIC_ enum ndr_err_code ndr_check_array_size(struct ndr_pull *ndr, void *p, uint32_t size)
 {
 	uint32_t stored;
-	stored = ndr_token_peek(&ndr->array_size_list, p);
+	NDR_CHECK(ndr_token_peek(&ndr->array_size_list, p, &stored));
 	if (stored != size) {
 		return ndr_pull_error(ndr, NDR_ERR_ARRAY_SIZE,
 				      "Bad array size - got %u expected %u\n",
@@ -1144,9 +1137,9 @@ _PUBLIC_ enum ndr_err_code ndr_pull_array_length(struct ndr_pull *ndr, const voi
 /*
   get the stored array length field
 */
-_PUBLIC_ uint32_t ndr_get_array_length(struct ndr_pull *ndr, const void *p)
+_PUBLIC_ enum ndr_err_code ndr_get_array_length(struct ndr_pull *ndr, const void *p, uint32_t *length)
 {
-	return ndr_token_peek(&ndr->array_length_list, p);
+	return ndr_token_peek(&ndr->array_length_list, p, length);
 }
 
 /*
@@ -1155,7 +1148,7 @@ _PUBLIC_ uint32_t ndr_get_array_length(struct ndr_pull *ndr, const void *p)
 _PUBLIC_ enum ndr_err_code ndr_check_array_length(struct ndr_pull *ndr, void *p, uint32_t length)
 {
 	uint32_t stored;
-	stored = ndr_token_peek(&ndr->array_length_list, p);
+	NDR_CHECK(ndr_token_peek(&ndr->array_length_list, p, &stored));
 	if (stored != length) {
 		return ndr_pull_error(ndr, NDR_ERR_ARRAY_SIZE,
 				      "Bad array length - got %u expected %u\n",
diff --git a/librpc/ndr/ndr_basic.c b/librpc/ndr/ndr_basic.c
index 82d2f3cfae6..e239cfb27d9 100644
--- a/librpc/ndr/ndr_basic.c
+++ b/librpc/ndr/ndr_basic.c
@@ -805,18 +805,20 @@ _PUBLIC_ enum ndr_err_code ndr_push_unique_ptr(struct ndr_push *ndr, const void
 */
 _PUBLIC_ enum ndr_err_code ndr_push_full_ptr(struct ndr_push *ndr, const void *p)
 {
+	enum ndr_err_code ret = NDR_ERR_SUCCESS;
 	uint32_t ptr = 0;
 	if (p) {
 		/* Check if the pointer already exists and has an id */
-		ptr = ndr_token_peek(&ndr->full_ptr_list, p);
-		if (ptr == 0) {
-			enum ndr_err_code ret = NDR_ERR_SUCCESS;
+		ret = ndr_token_peek(&ndr->full_ptr_list, p, &ptr);
+		if (ret == NDR_ERR_TOKEN) {
 			ndr->ptr_count++;
 			ptr = ndr->ptr_count;
 			ret = ndr_token_store(ndr, &ndr->full_ptr_list, p, ptr);
 			if (ret != NDR_ERR_SUCCESS) {
 				return ret;
 			}
+		} else if (ret != NDR_ERR_SUCCESS) {
+			return ret;
 		}
 	}
 	return ndr_push_uint3264(ndr, NDR_SCALARS, ptr);