Merge pull request #1837 from libgit2/ntk/topic/control_stream_write_size

odb: Error when streaming in less|more bytes than declared

4 files changed, 99 insertions, 3 deletions

diff --git a/include/git2/odb.h b/include/git2/odb.h
index 3e93a932c..e50a2a1c1 100644
--- a/include/git2/odb.h
+++ b/include/git2/odb.h
@@ -238,6 +238,9 @@ GIT_EXTERN(int) git_odb_open_wstream(git_odb_stream **out, git_odb *db, size_t s
 /**
  * Write to an odb stream
  *
+ * This method will fail as soon as the total number of
+ * received bytes exceeds the size declared with `git_odb_open_wstream()`
+ *
  * @param stream the stream
  * @param buffer the data to write
  * @param len the buffer's length
@@ -251,6 +254,9 @@ GIT_EXTERN(int) git_odb_stream_write(git_odb_stream *stream, const char *buffer,
  * The object will take its final name and will be available to the
  * odb.
  *
+ * This method will fail if the total number of received bytes
+ * differs from the size declared with `git_odb_open_wstream()`
+ *
  * @param out pointer to store the resulting object's id
  * @param stream the stream
  * @return 0 on success; an error code otherwise
diff --git a/include/git2/odb_backend.h b/include/git2/odb_backend.h
index bafeec047..e558bbb1c 100644
--- a/include/git2/odb_backend.h
+++ b/include/git2/odb_backend.h
@@ -78,6 +78,9 @@ struct git_odb_stream {
 	unsigned int mode;
 	void *hash_ctx;
 
+	size_t declared_size;
+	size_t received_bytes;
+
 	/**
 	 * Write at most `len` bytes into `buffer` and advance the
 	 * stream.
@@ -93,9 +96,13 @@ struct git_odb_stream {
 	 * Store the contents of the stream as an object with the id
 	 * specified in `oid`.
 	 *
-	 * This method will *not* be invoked by libgit2 if the object pointed at
-	 * by `oid` already exists in any backend. Libgit2 will however take care
-	 * of properly disposing the stream through a call to `free()`.
+	 * This method will *not* be invoked by libgit2 when:
+	 *  - the object pointed at by `oid` already exists in any backend.
+	 * 	- the total number of received bytes differs from the size declared
+	 *    with `git_odb_open_wstream()`
+	 *
+	 * Libgit2 will however take care of properly disposing the stream through
+	 * a call to `free()`.
 	 */
 	int (*finalize_write)(git_odb_stream *stream, const git_oid *oid);
 
diff --git a/src/odb.c b/src/odb.c
index dfb252178..2e6869547 100644
--- a/src/odb.c
+++ b/src/odb.c
@@ -888,17 +888,44 @@ int git_odb_open_wstream(
 	hash_header(ctx, size, type);
 	(*stream)->hash_ctx = ctx;
 
+	(*stream)->declared_size = size;
+	(*stream)->received_bytes = 0;
+
 	return error;
 }
 
+static int git_odb_stream__invalid_length(
+	const git_odb_stream *stream,
+	const char *action)
+{
+	giterr_set(GITERR_ODB,
+		"Cannot %s - "
+		"Invalid length. %"PRIuZ" was expected. The "
+		"total size of the received chunks amounts to %"PRIuZ".",
+		action, stream->declared_size, stream->received_bytes);		
+
+	return -1;
+}
+
 int git_odb_stream_write(git_odb_stream *stream, const char *buffer, size_t len)
 {
 	git_hash_update(stream->hash_ctx, buffer, len);
+
+	stream->received_bytes += len;
+
+	if (stream->received_bytes > stream->declared_size)
+		return git_odb_stream__invalid_length(stream,
+			"stream_write()");
+
 	return stream->write(stream, buffer, len);
 }
 
 int git_odb_stream_finalize_write(git_oid *out, git_odb_stream *stream)
 {
+	if (stream->received_bytes != stream->declared_size)
+		return git_odb_stream__invalid_length(stream,
+			"stream_finalize_write()");
+
 	git_hash_final(out, stream->hash_ctx);
 
 	if (git_odb_exists(stream->backend->odb, out))
diff --git a/tests-clar/odb/streamwrite.c b/tests-clar/odb/streamwrite.c
new file mode 100644
index 000000000..591a20040
--- /dev/null
+++ b/tests-clar/odb/streamwrite.c
@@ -0,0 +1,56 @@
+#include "clar_libgit2.h"
+#include "git2/odb_backend.h"
+
+static git_repository *repo;
+static git_odb *odb;
+static git_odb_stream *stream;
+
+void test_odb_streamwrite__initialize(void)
+{
+	repo = cl_git_sandbox_init("testrepo.git");
+	cl_git_pass(git_repository_odb(&odb, repo));
+
+	cl_git_pass(git_odb_open_wstream(&stream, odb, 14, GIT_OBJ_BLOB));
+	cl_assert_equal_sz(14, stream->declared_size);
+}
+
+void test_odb_streamwrite__cleanup(void)
+{
+	git_odb_stream_free(stream);
+	git_odb_free(odb);
+	cl_git_sandbox_cleanup();
+}
+
+void test_odb_streamwrite__can_accept_chunks(void)
+{
+	git_oid oid;
+
+	cl_git_pass(git_odb_stream_write(stream, "deadbeef", 8));
+	cl_assert_equal_sz(8, stream->received_bytes);
+
+	cl_git_pass(git_odb_stream_write(stream, "deadbeef", 6));
+	cl_assert_equal_sz(8 + 6, stream->received_bytes);
+
+	cl_git_pass(git_odb_stream_finalize_write(&oid, stream));
+}
+
+void test_odb_streamwrite__can_detect_missing_bytes(void)
+{
+	git_oid oid;
+
+	cl_git_pass(git_odb_stream_write(stream, "deadbeef", 8));
+	cl_assert_equal_sz(8, stream->received_bytes);
+
+	cl_git_pass(git_odb_stream_write(stream, "deadbeef", 4));
+	cl_assert_equal_sz(8 + 4, stream->received_bytes);
+
+	cl_git_fail(git_odb_stream_finalize_write(&oid, stream));
+}
+
+void test_odb_streamwrite__can_detect_additional_bytes(void)
+{
+	cl_git_pass(git_odb_stream_write(stream, "deadbeef", 8));
+	cl_assert_equal_sz(8, stream->received_bytes);
+
+	cl_git_fail(git_odb_stream_write(stream, "deadbeef", 7));
+}