kunit: add locking to string-stream operations

Change-Id: I6b21aceab9fe64267928aa88af709911f2724197
Signed-off-by: Felix Guo <felixguo@google.com>
diff --git a/include/test/string-stream.h b/include/test/string-stream.h
index 8e8b531..db46e33 100644
--- a/include/test/string-stream.h
+++ b/include/test/string-stream.h
@@ -2,6 +2,7 @@
 #define _TEST_STRING_STREAM_H
 
 #include <linux/types.h>
+#include <linux/spinlock.h>
 #include <stdarg.h>
 
 struct string_stream_fragment {
@@ -13,6 +14,8 @@
 	size_t length;
 	struct list_head fragments;
 
+	/* length and fragments are protected by this lock */
+	struct spinlock lock;
 	int (*add)(struct string_stream *this, const char *fmt, ...);
 	int (*vadd)(struct string_stream *this, const char *fmt, va_list args);
 	char *(*get_string)(struct string_stream *this);
diff --git a/test/string-stream.c b/test/string-stream.c
index f6049d6..993d95d 100644
--- a/test/string-stream.c
+++ b/test/string-stream.c
@@ -9,6 +9,7 @@
 	struct string_stream_fragment *fragment;
 	int len;
 	va_list args_for_counting;
+	unsigned long flags;
 
 	/* Make a copy because `vsnprintf` could change it */
 	va_copy(args_for_counting, args);
@@ -22,15 +23,17 @@
 	if (!fragment)
 		return -ENOMEM;
 
-	fragment->fragment = kmalloc(len + 1, GFP_KERNEL);
+	fragment->fragment = kmalloc(len, GFP_KERNEL);
 	if (!fragment->fragment) {
 		kfree(fragment);
 		return -ENOMEM;
 	}
 
 	len = vsnprintf(fragment->fragment, len, fmt, args);
+	spin_lock_irqsave(&this->lock, flags);
 	this->length += len;
 	list_add_tail(&fragment->node, &this->fragments);
+	spin_unlock_irqrestore(&this->lock, flags);
 	return 0;
 }
 
@@ -48,7 +51,9 @@
 static void string_stream_clear(struct string_stream *this)
 {
 	struct string_stream_fragment *fragment, *fragment_safe;
+	unsigned long flags;
 
+	spin_lock_irqsave(&this->lock, flags);
 	list_for_each_entry_safe(fragment,
 				 fragment_safe,
 				 &this->fragments,
@@ -57,28 +62,38 @@
 		kfree(fragment->fragment);
 		kfree(fragment);
 	}
-
 	this->length = 0;
+	spin_unlock_irqrestore(&this->lock, flags);
 }
 
 static char *string_stream_get_string(struct string_stream *this)
 {
 	struct string_stream_fragment *fragment;
 	char *buf;
+	unsigned long flags;
 
 	buf = kzalloc(this->length + 1, GFP_KERNEL);
 	if (!buf)
 		return NULL;
 
-	list_for_each_entry(fragment, &this->fragments, node) {
+	spin_lock_irqsave(&this->lock, flags);
+	list_for_each_entry(fragment, &this->fragments, node)
 		strcat(buf, fragment->fragment);
-	}
+	spin_unlock_irqrestore(&this->lock, flags);
+
 	return buf;
 }
 
 static bool string_stream_is_empty(struct string_stream *this)
 {
-	return list_empty(&this->fragments);
+	bool is_empty;
+	unsigned long flags;
+
+	spin_lock_irqsave(&this->lock, flags);
+	is_empty = list_empty(&this->fragments);
+	spin_unlock_irqrestore(&this->lock, flags);
+
+	return is_empty;
 }
 
 void destroy_string_stream(struct string_stream *stream)
@@ -95,7 +110,7 @@
 		return NULL;
 
 	INIT_LIST_HEAD(&stream->fragments);
-	stream->length = 0;
+	spin_lock_init(&stream->lock);
 	stream->add = string_stream_add;
 	stream->vadd = string_stream_vadd;
 	stream->get_string = string_stream_get_string;