/*
 * Copyright (c) 2022 Genome Research Ltd.
 * Author(s): James Bonfield
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *    1. Redistributions of source code must retain the above copyright notice,
 *       this list of conditions and the following disclaimer.
 *
 *    2. Redistributions in binary form must reproduce the above
 *       copyright notice, this list of conditions and the following
 *       disclaimer in the documentation and/or other materials provided
 *       with the distribution.
 *
 *    3. Neither the names Genome Research Ltd and Wellcome Trust Sanger
 *       Institute nor the names of its contributors may be used to endorse
 *       or promote products derived from this software without specific
 *       prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY GENOME RESEARCH LTD AND CONTRIBUTORS "AS
 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GENOME RESEARCH
 * LTD OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include "htslib/config.h"

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <inttypes.h>

#include "utils.h"

#ifndef NO_THREADS
#include <pthread.h>
#endif

//#define TLS_DEBUG

#ifndef NO_THREADS
/*
 * Thread local storage per thread in the pool.
 *
 * We have some large memory blocks for rANS which we cannot store on the
 * stack due to various system limitations.  Allocaitng them can be
 * expensive as some OSes use mmap and will pass the pages back to the OS
 * on each free.  This unfortunately then means zeroing the pages out again
 * on each new malloc, plus additional switching into the kernel.
 *
 * Instead where available, we use pthread_once to allocate a small arena
 * of memory buffers and we continually reuse these same buffers.  We don't
 * need to memset it (calloc equivalent) either as we're sure that any
 * leakage of data is simply an earlier set of precomputed frequency
 * lookups, and not something more sinister such as an encryption key.
 *
 * If we don't have pthreads, then we have to fall back to the slow
 * traditional calloc instead.
 */

#define MAX_TLS_BUFS 10
typedef struct {
    void   *bufs[MAX_TLS_BUFS];
    size_t sizes[MAX_TLS_BUFS];
    int     used[MAX_TLS_BUFS];
} tls_pool;

static pthread_once_t rans_once = PTHREAD_ONCE_INIT;
static pthread_key_t rans_key;

/*
 * Frees all local storage for this thread.
 * Note: this isn't a function to free a specific allocated item.
 */
static void htscodecs_tls_free_all(void *ptr) {
    tls_pool *tls = (tls_pool *)ptr;
    if (!tls)
        return;

    int i;
    for (i = 0; i < MAX_TLS_BUFS; i++) {
#ifdef TLS_DEBUG
        if (tls->bufs[i])
            fprintf(stderr, "Free %ld = %p\n", tls->sizes[i], tls->bufs[i]);
#endif
        if (tls->used[i]) {
            fprintf(stderr, "Closing thread while TLS data is in use\n");
        }
        free(tls->bufs[i]);
    }

    free(tls);
}

static void htscodecs_tls_init(void) {
    pthread_key_create(&rans_key, htscodecs_tls_free_all);
}

/*
 * Allocates size bytes from the global Thread Local Storage pool.
 * This is shared by all subsequent calls within this thread.
 *
 * An simpler alternative could be possible where we have a fixed number
 * of types of alloc, say 5, and specify the correct label when allocating.
 * Eg histogram, name_context, fqzcomp, rans.  We can have multiple types
 * in use in different stack frames (such name_context + hist + rans), but
 * the number is very limited.  That then paves the way to simply check and
 * realloc without needing to keep track of use status or overflowing
 * the maximum number permitted.
 */
void *htscodecs_tls_alloc(size_t size) {
    int i;

    int err = pthread_once(&rans_once, htscodecs_tls_init);
    if (err != 0) {
        fprintf(stderr, "Initialising TLS data failed: pthread_once: %s\n",
                strerror(err));
        return NULL;
    }

    // Initialise tls_pool on first usage
    tls_pool *tls = pthread_getspecific(rans_key);
    if (!tls) {
        if (!(tls = calloc(1, sizeof(*tls))))
            return NULL;
        pthread_setspecific(rans_key, tls);
    }

    // Query pool for size
    int avail = -1;
    for (i = 0; i < MAX_TLS_BUFS; i++) {
        if (!tls->used[i]) {
            if (size <= tls->sizes[i]) {
                tls->used[i] = 1;
#ifdef TLS_DEBUG
                fprintf(stderr, "Reuse %d: %ld/%ld = %p\n",
                        i, size, tls->sizes[i], tls->bufs[i]);
#endif
                return tls->bufs[i];
            } else if (avail == -1) {
                avail = i;
            }
        }
    }

    if (i == MAX_TLS_BUFS && avail == -1) {
        // Shouldn't happen given our very limited use of this function
        fprintf(stderr, "Error: out of rans_tls_alloc slots\n");
        return NULL;
    }

    if (tls->bufs[avail])
        free(tls->bufs[avail]);
    if (!(tls->bufs[avail] = calloc(1, size)))
        return NULL;
#ifdef TLS_DEBUG
    fprintf(stderr, "Alloc %d: %ld = %p\n", avail, size, tls->bufs[avail]);
#endif
    tls->sizes[avail] = size;
    tls->used[avail] = 1;

    return tls->bufs[avail];
}

void *htscodecs_tls_calloc(size_t nmemb, size_t size) {
#ifdef TLS_DEBUG
    fprintf(stderr, "htscodecs_tls_calloc(%ld)\n", nmemb*size);
#endif
    void *ptr = htscodecs_tls_alloc(nmemb * size);
    if (ptr)
        memset(ptr, 0, nmemb * size);
    return ptr;
}

void htscodecs_tls_free(void *ptr) {
    if (!ptr)
        return;

    tls_pool *tls = pthread_getspecific(rans_key);

    int i;
    for (i = 0; i < MAX_TLS_BUFS; i++) {
        if (tls->bufs[i] == ptr)
            break;
    }
#ifdef TLS_DEBUG
    fprintf(stderr, "Fake free %d size %ld ptr %p\n",
            i, tls->sizes[i], tls->bufs[i]);
#endif
    if (i == MAX_TLS_BUFS) {
        fprintf(stderr, "Attempt to htscodecs_tls_free a buffer not allocated"
                " with htscodecs_tls_alloc\n");
        return;
    }
    if (!tls->used[i]) {
        fprintf(stderr, "Attempt to htscodecs_tls_free a buffer twice\n");
        return;
    }
    tls->used[i] = 0;
}

#else
/*
 * Calloc/free equivalents instead.
 *
 * We use calloc instead of malloc as a sufficiently malformed set of input
 * frequencies may not sum to the expected total frequency size, leaving
 * some elements uninitialised.  It's unlikely, but potentially a crafty
 * attacker could somehow exploit this to pull out parts of this allocated
 * buffer and leak them into the decompressed data stream, potentially
 * compromising previous buffers such as encryption keys.  (Although
 * frankly any well-written crypto library should be zeroing such memory
 * before freeing it to ensure it's never visible to a subsequent malloc.)
 */
void *htscodecs_tls_alloc(size_t size) {
    return calloc(1, size);
}

void *htscodecs_tls_calloc(size_t nmemb, size_t size) {
    return calloc(nmemb, size);
}

void htscodecs_tls_free(void *ptr) {
    free(ptr);
}
#endif
