From 3fff3ae6db10e26dc283eb1a3b601322087deb5b Mon Sep 17 00:00:00 2001
From: leitner <leitner>
Date: Fri, 30 Oct 2020 12:13:05 +0000
Subject: [PATCH] add and use bs_capacityassert

---
 buffer/bs_capacityassert.c    | 13 +++++++++++
 buffer/bs_init_bstream_size.c | 42 ++++++++++++++++++++++++++++++++++-
 buffer/prs_asciiz.c           |  8 +++----
 buffer/prs_asciiz_fixedlen.c  |  8 +++----
 parse.h                       |  6 +++++
 5 files changed, 68 insertions(+), 9 deletions(-)
 create mode 100644 buffer/bs_capacityassert.c

diff --git a/buffer/bs_capacityassert.c b/buffer/bs_capacityassert.c
new file mode 100644
index 0000000..fbb0568
--- /dev/null
+++ b/buffer/bs_capacityassert.c
@@ -0,0 +1,13 @@
+#include "parse.h"
+
+// This function asserts that the caller wants to read n bytes from the stream.
+// If there is less space available, set the stream to error state.
+int bs_capacityassert(struct bytestream* bs,size_t capacity) {
+  if (bs->cur>=bs->max || 	// if EOF or error, return 0
+      bs->max - bs->cur < capacity) {	// not EOF but less than that many bytes left
+    bs->cur = 1;	// set stream to error state
+    bs->max = 0;
+    return 0;		// return failure
+  }
+  return 1;
+}
diff --git a/buffer/bs_init_bstream_size.c b/buffer/bs_init_bstream_size.c
index 69ec67e..7033360 100644
--- a/buffer/bs_init_bstream_size.c
+++ b/buffer/bs_init_bstream_size.c
@@ -3,7 +3,7 @@
 void bs_init_bstream_size(struct bytestream* bs,struct bytestream* other,size_t maxlen) {
   bs->type = BSTREAM;
   // check if we have enough capacity in the parent bytestream
-  if (bs_capacitycheck(other, maxlen)) {
+  if (bs_capacityassert(other, maxlen)) {
     bs->cur = 0;
     bs->max = maxlen;
   } else {
@@ -13,3 +13,43 @@ void bs_init_bstream_size(struct bytestream* bs,struct bytestream* other,size_t
   }
   bs->u.bs=other;
 }
+
+#ifdef UNITTEST
+#include <assert.h>
+
+#undef UNITTEST
+#include "buffer/bs_get.c"
+#include "buffer/bs_capacityassert.c"
+#include "buffer/bs_init_membuf.c"
+
+// we use membuf here, mock buffer stuff away
+ssize_t buffer_getc(buffer* b,char* x) { return 0; }
+
+int main() {
+  struct bytestream bs = BS_FROM_MEMBUF("\x12\x34\x56\x78",4);
+  struct bytestream sub;
+  assert(bs_get(&bs) == 0x12);		// eat one byte
+  bs_init_bstream_size(&sub, &bs, 2);	// 2 byte substream
+  assert(bs_get(&sub) == 0x34);		// make sure we got the offset right
+  assert(bs_get(&sub) == 0x56);		// read second (and last) byte
+
+  assert(bs_err(&sub) == 0);		// should have no errors yet
+  assert(bs_err(&bs) == 0);
+
+  assert(bs_get(&sub) == 0);		// next byte should error out
+  assert(bs_err(&sub));			// substream should have error state
+  assert(bs_err(&bs) == 0);		// parent stream not
+  assert(bs_get(&bs) == 0x78);		// and we should not have overstepped
+
+  // now check the range checking in bs_init_bstream_size itself
+  bs_init_membuf(&bs, "\x12\x34\x56\x78", 4);
+
+  bs_init_bstream_size(&sub, &bs, 4);	// tight fit
+  assert(bs_err(&sub) == 0);
+  assert(bs_err(&bs) == 0);
+
+  bs_init_bstream_size(&sub, &bs, 5);
+  assert(bs_err(&sub));
+  assert(bs_err(&bs));
+}
+#endif
diff --git a/buffer/prs_asciiz.c b/buffer/prs_asciiz.c
index 8011671..9dce998 100644
--- a/buffer/prs_asciiz.c
+++ b/buffer/prs_asciiz.c
@@ -45,11 +45,11 @@ ssize_t prs_asciiz(struct bytestream* bs, char* dest, size_t len) {
 #undef UNITTEST
 #include "buffer/bs_init_membuf.c"
 #include "buffer/bs_get.c"
-#include "buffer/buffer_peekc.c"
-#include "buffer/buffer_getc.c"
 #include "buffer/bs_peek.c"
-#include "buffer/buffer_feed.c"
-#include "buffer/buffer_stubborn2.c"
+
+// we use membuf here, mock buffer stuff away
+ssize_t buffer_getc(buffer* b,char* x) { return 0; }
+ssize_t buffer_peekc(buffer* b,char* x) { return 0; }
 
 int main() {
   struct bytestream bs = BS_FROM_MEMBUF("fnord\n\0x",8);
diff --git a/buffer/prs_asciiz_fixedlen.c b/buffer/prs_asciiz_fixedlen.c
index 4b88640..cf69472 100644
--- a/buffer/prs_asciiz_fixedlen.c
+++ b/buffer/prs_asciiz_fixedlen.c
@@ -36,11 +36,11 @@ ssize_t prs_asciiz_fixedlen(struct bytestream* bs, char* dest, size_t len) {
 #undef UNITTEST
 #include "buffer/bs_init_membuf.c"
 #include "buffer/bs_get.c"
-#include "buffer/buffer_peekc.c"
-#include "buffer/buffer_getc.c"
 #include "buffer/bs_peek.c"
-#include "buffer/buffer_feed.c"
-#include "buffer/buffer_stubborn2.c"
+
+// we use membuf here, mock buffer stuff away
+ssize_t buffer_getc(buffer* b,char* x) { return 0; }
+ssize_t buffer_peekc(buffer* b,char* x) { return 0; }
 
 int main() {
   struct bytestream bs = BS_FROM_MEMBUF("fnord\n\0x",8);
diff --git a/parse.h b/parse.h
index 1ff6c94..c4061b0 100644
--- a/parse.h
+++ b/parse.h
@@ -85,6 +85,12 @@ int bs_err(struct bytestream* bs);
 
 /* Can we read this much more bytes from the bytestream? */
 int bs_capacitycheck(struct bytestream* bs,size_t capacity);
+/* Like bs_capacitycheck but will set stream to error state on fail */
+int bs_capacityassert(struct bytestream* bs,size_t capacity);
+
+/* Read n bytes from stream. Return n.
+ * Set stream to error state if not enough space or I/O error. */
+size_t prs_readblob(struct bytestream* bs,unsigned char* dest,size_t destlen);
 
 uint16_t prs_u16(struct bytestream* bs);
 uint16_t prs_u16_big(struct bytestream* bs);