lib/kunit/string-stream.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /*
   3  * C++ stream style string builder used in KUnit for building messages.
   4  *
   5  * Copyright (C) 2019, Google LLC.
   6  * Author: Brendan Higgins <brendanhiggins@google.com>
   7  */
   8
   9 #include <kunit/test.h>
  10 #include <linux/list.h>
  11 #include <linux/slab.h>
  12
  13 #include "string-stream.h"
  14
  15
  16 static struct string_stream_fragment *alloc_string_stream_fragment(
  17                 struct kunit *test, int len, gfp_t gfp)
  18 {
  19         struct string_stream_fragment *frag;
  20
  21         frag = kunit_kzalloc(test, sizeof(*frag), gfp);
  22         if (!frag)
  23                 return ERR_PTR(-ENOMEM);
  24
  25         frag->fragment = kunit_kmalloc(test, len, gfp);
  26         if (!frag->fragment) {
  27                 kunit_kfree(test, frag);
  28                 return ERR_PTR(-ENOMEM);
  29         }
  30
  31         return frag;
  32 }
  33
  34 static void string_stream_fragment_destroy(struct kunit *test,
  35                                            struct string_stream_fragment *frag)
  36 {
  37         list_del(&frag->node);
  38         kunit_kfree(test, frag->fragment);
  39         kunit_kfree(test, frag);
  40 }
  41
  42 int string_stream_vadd(struct string_stream *stream,
  43                        const char *fmt,
  44                        va_list args)
  45 {
  46         struct string_stream_fragment *frag_container;
  47         int len;
  48         va_list args_for_counting;
  49
  50         /* Make a copy because `vsnprintf` could change it */
  51         va_copy(args_for_counting, args);
  52
  53         /* Need space for null byte. */
  54         len = vsnprintf(NULL, 0, fmt, args_for_counting) + 1;
  55
  56         va_end(args_for_counting);
  57
  58         frag_container = alloc_string_stream_fragment(stream->test,
  59                                                       len,
  60                                                       stream->gfp);
  61         if (IS_ERR(frag_container))
  62                 return PTR_ERR(frag_container);
  63
  64         len = vsnprintf(frag_container->fragment, len, fmt, args);
  65         spin_lock(&stream->lock);
  66         stream->length += len;
  67         list_add_tail(&frag_container->node, &stream->fragments);
  68         spin_unlock(&stream->lock);
  69
  70         return 0;
  71 }
  72
  73 int string_stream_add(struct string_stream *stream, const char *fmt, ...)
  74 {
  75         va_list args;
  76         int result;
  77
  78         va_start(args, fmt);
  79         result = string_stream_vadd(stream, fmt, args);
  80         va_end(args);
  81
  82         return result;
  83 }
  84
  85 static void string_stream_clear(struct string_stream *stream)
  86 {
  87         struct string_stream_fragment *frag_container, *frag_container_safe;
  88
  89         spin_lock(&stream->lock);
  90         list_for_each_entry_safe(frag_container,
  91                                  frag_container_safe,
  92                                  &stream->fragments,
  93                                  node) {
  94                 string_stream_fragment_destroy(stream->test, frag_container);
  95         }
  96         stream->length = 0;
  97         spin_unlock(&stream->lock);
  98 }
  99
 100 char *string_stream_get_string(struct string_stream *stream)
 101 {
 102         struct string_stream_fragment *frag_container;
 103         size_t buf_len = stream->length + 1; /* +1 for null byte. */
 104         char *buf;
 105
 106         buf = kunit_kzalloc(stream->test, buf_len, stream->gfp);
 107         if (!buf)
 108                 return NULL;
 109
 110         spin_lock(&stream->lock);
 111         list_for_each_entry(frag_container, &stream->fragments, node)
 112                 strlcat(buf, frag_container->fragment, buf_len);
 113         spin_unlock(&stream->lock);
 114
 115         return buf;
 116 }
 117
 118 int string_stream_append(struct string_stream *stream,
 119                          struct string_stream *other)
 120 {
 121         const char *other_content;
 122
 123         other_content = string_stream_get_string(other);
 124
 125         if (!other_content)
 126                 return -ENOMEM;
 127
 128         return string_stream_add(stream, other_content);
 129 }
 130
 131 bool string_stream_is_empty(struct string_stream *stream)
 132 {
 133         return list_empty(&stream->fragments);
 134 }
 135
 136 struct string_stream_alloc_context {
 137         struct kunit *test;
 138         gfp_t gfp;
 139 };
 140
 141 struct string_stream *alloc_string_stream(struct kunit *test, gfp_t gfp)
 142 {
 143         struct string_stream *stream;
 144
 145         stream = kunit_kzalloc(test, sizeof(*stream), gfp);
 146         if (!stream)
 147                 return ERR_PTR(-ENOMEM);
 148
 149         stream->gfp = gfp;
 150         stream->test = test;
 151         INIT_LIST_HEAD(&stream->fragments);
 152         spin_lock_init(&stream->lock);
 153
 154         return stream;
 155 }
 156
 157 void string_stream_destroy(struct string_stream *stream)
 158 {
 159         string_stream_clear(stream);
 160 }