Fixes elf panic
[akaros.git] / kern / src / kmalloc.c
index 0e1cd0a..3f9ebbf 100644 (file)
@@ -30,10 +30,17 @@ static spinlock_t pages_list_lock = SPINLOCK_INITIALIZER;
 static page_list_t LCKD(&pages_list_lock)pages_list;
 
 struct kmem_cache *kmalloc_caches[NUM_KMALLOC_CACHES];
+
+static void __kfree_release(struct kref *kref);
+
 void kmalloc_init(void)
 {
-       static_assert(sizeof(struct kmalloc_tag) == 16);
-       // build caches of common sizes
+       /* we want at least a 16 byte alignment of the tag so that the bufs kmalloc
+        * returns are 16 byte aligned.  we used to check the actual size == 16,
+        * since we adjusted the KMALLOC_SMALLEST based on that. */
+       static_assert(ALIGNED(sizeof(struct kmalloc_tag), 16));
+       /* build caches of common sizes.  this size will later include the tag and
+        * the actual returned buffer. */
        size_t ksize = KMALLOC_SMALLEST;
        for (int i = 0; i < NUM_KMALLOC_CACHES; i++) {
                kmalloc_caches[i] = kmem_cache_create("kmalloc_cache", ksize,
@@ -65,6 +72,7 @@ void *kmalloc(size_t size, int flags)
                tag->flags = KMALLOC_TAG_PAGES;
                tag->num_pages = num_pgs;
                tag->canary = KMALLOC_CANARY;
+               kref_init(&tag->kref, __kfree_release, 1);
                return buf + sizeof(struct kmalloc_tag);
        }
        // else, alloc from the appropriate cache
@@ -76,13 +84,14 @@ void *kmalloc(size_t size, int flags)
        tag->flags = KMALLOC_TAG_CACHE;
        tag->my_cache = kmalloc_caches[cache_id];
        tag->canary = KMALLOC_CANARY;
+       kref_init(&tag->kref, __kfree_release, 1);
        return buf + sizeof(struct kmalloc_tag);
 }
 
 void *kzmalloc(size_t size, int flags) 
 {
        void *v = kmalloc(size, flags);
-       if (! v)
+       if (!v)
                return v;
        memset(v, 0, size);
        return v;
@@ -122,22 +131,43 @@ void *kzmalloc_align(size_t size, int flags, size_t align)
        return v;
 }
 
+static struct kmalloc_tag *__get_km_tag(void *buf)
+{
+       struct kmalloc_tag *tag = (struct kmalloc_tag*)(buf -
+                                                   sizeof(struct kmalloc_tag));
+       if (tag->canary != KMALLOC_CANARY){
+               printk("__get_km_tag bad canary: %08lx, expected %08lx\n", tag->canary,
+                      KMALLOC_CANARY);
+               hexdump((void *)(buf - sizeof(struct kmalloc_tag)), 256);
+               panic("Bad canary");
+       }
+       return tag;
+}
+
+/* If we kmalloc_aligned, the buf we got back (and are now trying to perform
+ * some operation on) might not be the original, underlying, unaligned buf.
+ *
+ * This returns the underlying, unaligned buf, or 0 if the buf was not realigned
+ * in the first place. */
+static void *__get_unaligned_orig_buf(void *buf)
+{
+       int *tag_flags = (int*)(buf - sizeof(int));
+       if ((*tag_flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_UNALIGN)
+               return (buf - (*tag_flags >> KMALLOC_ALIGN_SHIFT));
+       else
+               return 0;
+}
+
 void *krealloc(void* buf, size_t size, int flags)
 {
        void *nbuf;
        size_t osize = 0;
-       int *tag_flags;
+       struct kmalloc_tag *tag;
+
        if (buf){
-               struct kmalloc_tag *tag = (struct kmalloc_tag*)(buf -
-                                                       sizeof(struct kmalloc_tag));
-               tag_flags = (int*)(buf - sizeof(int));
-               if ((*tag_flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_UNALIGN)
+               if (__get_unaligned_orig_buf(buf))
                        panic("krealloc of a kmalloc_align not supported");
-               if (tag->canary != KMALLOC_CANARY){
-                       printk("krealloc bad canary: %08lx, expected %08lx\n", tag->canary,
-                              KMALLOC_CANARY);
-                       hexdump((void *)(buf - sizeof(struct kmalloc_tag)), 256);
-               }
+               tag = __get_km_tag(buf);
                /* whatever we got from either a slab or the page allocator is meant for
                 * both the buf+size as well as the kmalloc tag */
                if ((tag->flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_CACHE) {
@@ -168,31 +198,47 @@ void *krealloc(void* buf, size_t size, int flags)
        return nbuf;
 }
 
-void kfree(void *addr)
+/* Grabs a reference on a buffer.  Release with kfree().
+ *
+ * Note that a krealloc on a buffer with ref > 1 that needs a new, underlying
+ * buffer will result in two buffers existing.  In this case, the krealloc is a
+ * kmalloc and a kfree, but that kfree does not completely free since the
+ * original ref > 1. */
+void kmalloc_incref(void *buf)
 {
-       struct kmalloc_tag *tag;
-       int *tag_flags;
-       if (addr == NULL)
-               return;
-       tag_flags = (int*)(addr - sizeof(int));
-       if ((*tag_flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_UNALIGN) {
-               kfree(addr - (*tag_flags >> KMALLOC_ALIGN_SHIFT));
-               return;
-       }
-       tag = (struct kmalloc_tag*)(addr - sizeof(struct kmalloc_tag));
-       assert(tag_flags == &tag->flags);
-       if (tag->canary != KMALLOC_CANARY){
-               printk("Canary is bogus: %08lx, expected %08lx\n", tag->canary,
-                      KMALLOC_CANARY);
-               hexdump((void*)(addr - sizeof(struct kmalloc_tag)), 256);
-       }
-       assert(tag->canary == KMALLOC_CANARY);
+       void *orig_buf = __get_unaligned_orig_buf(buf);
+       buf = orig_buf ? orig_buf : buf;
+       /* if we want a smaller tag, we can extract the code from kref and manually
+        * set the release method in kfree. */
+       kref_get(&__get_km_tag(buf)->kref, 1);
+}
+
+int kmalloc_refcnt(void *buf)
+{
+       void *orig_buf = __get_unaligned_orig_buf(buf);
+       buf = orig_buf ? orig_buf : buf;
+       return kref_refcnt(&__get_km_tag(buf)->kref);
+}
+
+static void __kfree_release(struct kref *kref)
+{
+       struct kmalloc_tag *tag = container_of(kref, struct kmalloc_tag, kref);
        if ((tag->flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_CACHE)
                kmem_cache_free(tag->my_cache, tag);
-       else if ((tag->flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_PAGES) {
+       else if ((tag->flags & KMALLOC_FLAG_MASK) == KMALLOC_TAG_PAGES)
                free_cont_pages(tag, LOG2_UP(tag->num_pages));
-       } else 
-               panic("[Italian Accent]: Che Cazzo! BO! Flag in kmalloc!!!");
+       else
+               panic("Bad flag 0x%x in %s", tag->flags, __FUNCTION__);
+}
+
+void kfree(void *buf)
+{
+       void *orig_buf;
+       if (buf == NULL)
+               return;
+       orig_buf = __get_unaligned_orig_buf(buf);
+       buf = orig_buf ? orig_buf : buf;
+       kref_put(&__get_km_tag(buf)->kref);
 }
 
 void kmalloc_canary_check(char *str)