/* Copyright (C) 2002-2005 RealVNC Ltd.  All Rights Reserved.
 * Copyright 2015-2024 m-privacy GmbH, Berlin
 *
 * This is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this software; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307,
 * USA.
 */

#include <zlib.h>

#include <rdr/OutStream.h>
#include <rdr/Exception.h>
#include <rfb/Configuration.h>

using namespace rdr;

rfb::IntParameter overKeepTTL("OverKeepTTL", "Time in seconds to keep old buffers when over KeepOutBuffers", 600, 0, 86400);

OutStream::OutStream() : corked(false), storedKey(0), classLog(&ovlog),
  maxBufferSize(MAXBUFFERSIZE), serial(0), allowMasterKey(true) {
  classLog->debug("Initializing");
  MUTEX_INIT(&serialMutex);
  for (int q = QPRIOMIN; q <= QPRIOMAX; q++) {
    MUTEX_INIT(&bufferPoolMutex[q]);
    MUTEX_INIT(&queueMutex[q]);
    numBuffers[q] = 0;
    spareBuffers[q] = 0;
    maxUsedBuffers[q] = 0;
    maxBuffers[q] = 20;
    keepBuffers[q] = 10;
    flushThreshBuffers[q] = 15;
    overKeepCount[q] = 0;
    submitCount[q] = 0;
    unusedBufferHead[q] = NULL;
    unusedBufferTail[q] = NULL;
    bufferQueueHead[q] = NULL;
    bufferQueueTail[q] = NULL;
  }
}

OutStream::~OutStream() {
  struct queueBuffer * buffer;
  struct queueBuffer * nextBuffer;

  for (int q = QPRIOMIN; q <= QPRIOMAX; q++) {
    MUTEX_LOCK(&queueMutex[q]);
    buffer = bufferQueueHead[q];
    while (buffer) {
      nextBuffer = buffer->next;
      free(buffer->data);
      free(buffer);
      buffer = nextBuffer;
    }
    bufferQueueHead[q] = NULL;
    bufferQueueTail[q] = NULL;
    MUTEX_UNLOCK(&queueMutex[q]);
    MUTEX_DESTROY(&queueMutex[q]);
    MUTEX_LOCK(&bufferPoolMutex[q]);
    buffer = unusedBufferHead[q];
    while (buffer) {
      nextBuffer = buffer->next;
      free(buffer->data);
      free(buffer);
      buffer = nextBuffer;
    }
    unusedBufferHead[q] = NULL;
    unusedBufferTail[q] = NULL;
    if (maxUsedBuffers[q] > 0)
      classLog->debug("Prio %u used maximum of %u buffers, keep %u, submitCount %u, overKeepCount %u", q, maxUsedBuffers[q], keepBuffers[q], submitCount[q], overKeepCount[q]);
    numBuffers[q] = 0;
    spareBuffers[q] = 0;
    MUTEX_UNLOCK(&bufferPoolMutex[q]);
    MUTEX_DESTROY(&bufferPoolMutex[q]);
  }
  MUTEX_DESTROY(&serialMutex);
}

void OutStream::savebacktrace(const char * function) {
#ifdef WIN32
  void *array[20];
  unsigned short size;
  int fd;
  unsigned int i;
  SYMBOL_INFO * symbol;
  HANDLE process;

  classLog->info("OutStream::savebacktrace() called by %s", function);

  fd = open(get_viewer_outstream_trace_file_path(), O_WRONLY | O_APPEND | O_CREAT, S_IRUSR|S_IWUSR);
  if (fd >= 0) {
    FILE * file;

    write(fd, "*********\n", 10);
    file = fdopen(fd, "a");
    process = GetCurrentProcess();
    SymInitialize(process, NULL, TRUE);
    size = CaptureStackBackTrace(0, 20, array, NULL);
    symbol = ( SYMBOL_INFO * )calloc( sizeof( SYMBOL_INFO ) + 256 * sizeof( char ), 1 );
    symbol->MaxNameLen   = 255;
    symbol->SizeOfStruct = sizeof( SYMBOL_INFO );
    for( i = 0; i < size; i++ ) {
      SymFromAddr(process, (DWORD64)(array[i]), 0, symbol);
      fprintf(file, "%i: %s - 0x%0lX\n", size - i - 1, symbol->Name, (long unsigned int) symbol->Address);
    }
    free(symbol);
    close(fd);
  }
#else
  void *array[20];
  size_t size;
  int fd;
  char traceName[100];

  size = backtrace(array, 20);
  snprintf(traceName, 99, "/tmp/Xtightgatevnc-%u-OutStream-trace.log", getpid());
  classLog->info("OutStream::savebacktrace() called by %s, save to %s", function, traceName);
  fd = open(traceName, O_WRONLY | O_APPEND | O_CREAT | O_NOFOLLOW, S_IRUSR|S_IWUSR);
  if (fd >= 0) {
    write(fd, "*********\n", 10);
    backtrace_symbols_fd(array, size, fd);
    close(fd);
  }
#endif
}

// writeU/SN() methods write unsigned and signed N-bit integers.
// inefficient, but working

void OutStream::writeU8(U8 u, int key, QPrio prio) {
  if(!check_key(key, __func__))
    return;
  if(!check_prio(prio, __func__))
    return;

  struct queueBuffer * buffer = getQueueBuffer(1, prio);
  if (buffer) {
    copyU8(buffer->data, u);
    buffer->used = 1;
    if (!submitQueueBuffer(buffer, buffer, key, prio)) {
      classLog->error("writeU8(): failed to submit buffer for prio %u", prio);
      returnQueueBuffer(buffer, prio);
    }
  } else {
    classLog->error("writeU8(): failed to get buffer for prio %u, size 1", prio);
  }
}

void OutStream::writeU16(U16 u, int key, QPrio prio) {
  if(!check_key(key, __func__))
    return;
  if(!check_prio(prio, __func__))
    return;

  struct queueBuffer * buffer = getQueueBuffer(2, prio);
  if (buffer) {
    copyU16(buffer->data, u);
    buffer->used = 2;
    if (!submitQueueBuffer(buffer, buffer, key, prio)) {
      classLog->error("writeU16(): failed to submit buffer for prio %u", prio);
      returnQueueBuffer(buffer, prio);
    }
  } else {
    classLog->error("writeU16(): failed to get buffer for prio %u, size 2", prio);
  }
}

void OutStream::writeU32(U32 u, int key, QPrio prio) {
  if(!check_key(key, __func__))
    return;
  if(!check_prio(prio, __func__))
    return;

  struct queueBuffer * buffer = getQueueBuffer(4, prio);
  if (buffer) {
    copyU32(buffer->data, u);
    buffer->used = 4;
    if (!submitQueueBuffer(buffer, buffer, key, prio)) {
      classLog->error("writeU32(): failed to submit buffer for prio %u", prio);
      returnQueueBuffer(buffer, prio);
    }
  } else {
    classLog->error("writeU32(): failed to get buffer for prio %u, size 4", prio);
  }
}

void OutStream::pad(size_t bytes, int key, QPrio prio) {
  if(!check_key(key, __func__))
    return;
  if(!check_prio(prio, __func__))
    return;
  if (bytes <= 0) {
    classLog->error("pad(): called with bytes %lu, returning", bytes);
    return;
  }

  struct queueBuffer * buffer = getQueueBuffer(bytes, prio);
  if (buffer) {
    memset(buffer->data, 0, bytes);
    buffer->used = bytes;
    if (!submitQueueBuffer(buffer, buffer, key, prio)) {
      classLog->error("pad(): failed to submit buffer for prio %u", prio);
      returnQueueBuffer(buffer, prio);
    }
  } else {
    classLog->error("pad(): failed to get buffer for prio %u, size %lu", prio, bytes);
  }
}

/* writeBytes() writes an exact number of bytes. */
void OutStream::writeBytes(const void* data, size_t numBytes, int key, QPrio prio) {
  U16 chunksize;
  U8 * pos;
  struct queueBuffer * buffer;
  struct queueBuffer * firstBuffer = NULL;
  struct queueBuffer * lastBuffer = NULL;

  if(!check_key(key, __func__))
    return;
  if(!check_prio(prio, __func__))
    return;
  if ((numBytes <= 0) || !data) {
    classLog->debug("writeBytes(): called without data, returning");
    return;
  }

  pos = (U8 *) data;
  while (numBytes > 0) {
    chunksize = (unsigned) numBytes > maxBufferSize ? maxBufferSize : numBytes;
    buffer = getQueueBuffer(chunksize, prio);
    if (!buffer) {
      classLog->error("writeBytes(): failed to get buffer for prio %u, size %u, wait and retry", prio, chunksize);
      Sleep(100);
      continue;
    }
    memcpy(buffer->data, pos, chunksize);
    buffer->used = chunksize;

    if (!firstBuffer) {
      firstBuffer = buffer;
    } else {
      lastBuffer->next = buffer;
    }
    lastBuffer = buffer;

    numBytes -= chunksize;
    pos += chunksize;
  }
  lastBuffer->next = NULL;
  if (!submitQueueBuffer(firstBuffer, lastBuffer, key, prio)) {
    classLog->error("writeBytes(): failed to submit buffer %u for prio %u, size %u", firstBuffer->number, prio, chunksize);
    returnQueueBuffer(firstBuffer, prio);
  }
}

// writeOpaqueN() writes a quantity without byte-swapping.

void OutStream::writeOpaque8(U8 u, int key, QPrio prio) {
  writeU8(u, key, prio);
}

void OutStream::writeOpaque16(U16 u, int key, QPrio prio) {
  if(!check_key(key, __func__))
    return;
  if(!check_prio(prio, __func__))
    return;
  struct queueBuffer * buffer = getQueueBuffer(2, prio);
  if (buffer) {
    *(buffer->data) = ((U8*)&u)[0];
    *(buffer->data + 1) = ((U8*)&u)[1];
    buffer->used = 2;
    if (!submitQueueBuffer(buffer, buffer, key, prio)) {
      classLog->error("writeOpaque16(): failed to submit buffer for prio %u", prio);
      returnQueueBuffer(buffer, prio);
    }
  } else {
    classLog->error("writeOpaque16(): failed to get buffer for prio %u, size 2", prio);
  }
}

void OutStream::writeOpaque32(U32 u, int key, QPrio prio) {
  if(!check_key(key, __func__))
    return;
  if(!check_prio(prio, __func__))
    return;
  struct queueBuffer * buffer = getQueueBuffer(4, prio);
  if (buffer) {
    *(buffer->data) = ((U8*)&u)[0];
    *(buffer->data + 1) = ((U8*)&u)[1];
    *(buffer->data + 2) = ((U8*)&u)[2];
    *(buffer->data + 3) = ((U8*)&u)[3];
    buffer->used = 4;
    if (!submitQueueBuffer(buffer, buffer, key, prio)) {
      classLog->error("writeOpaque32(): failed to submit buffer for prio %u", prio);
      returnQueueBuffer(buffer, prio);
    }
  } else {
    classLog->error("writeOpaque32(): failed to get buffer for prio %u, size 4", prio);
  }
}


// Buffer Queueing
// Buffers are per queue level == prio.
// A buffer got here must be returned to the level's pool with
// returnQueueBuffer() later.
// get unused buffer of given size
struct queueBuffer * OutStream::getQueueBuffer(U32 size, QPrio prio) {
  struct queueBuffer * newBuffer = NULL;
  U32 retry = 0;

  if(!check_prio(prio, __func__))
    return NULL;
  if (!size)
    return NULL;
  if (size > maxBufferSize) {
    classLog->error("getQueueBuffer(): requested buffer size %u for prio %u too big", size, prio);
    return NULL;
  }
  while (!newBuffer) {
    MUTEX_LOCK(&bufferPoolMutex[prio]);
    if (unusedBufferHead[prio]) {
      newBuffer = unusedBufferHead[prio];
      unusedBufferHead[prio] = newBuffer->next;
      if (!unusedBufferHead[prio])
        unusedBufferTail[prio] = NULL;
      spareBuffers[prio]--;
      MUTEX_UNLOCK(&bufferPoolMutex[prio]);
//      classLog->debug("getQueueBuffer(): reuse buffer %u, prio %u for size %u", newBuffer->number, prio, size);
      if(newBuffer->size < size) {
        U8 * newData;
//        classLog->debug("getQueueBuffer(): resize buffer %u, prio %u, from %u to %u", newBuffer->number, prio, newBuffer->size, size);
        newData = (U8 *) realloc(newBuffer->data, size);
        if (newData) {
          newBuffer->data = newData;
          newBuffer->size = size;
        } else {
          classLog->error("getQueueBuffer(): failed to realloc data for buffer %u from %u to %u, destroy buffer and retry", newBuffer->number, newBuffer->size, size);
          MUTEX_LOCK(&bufferPoolMutex[prio]);
          numBuffers[prio]--;
          MUTEX_UNLOCK(&bufferPoolMutex[prio]);
          free(newBuffer->data);
          free(newBuffer);
          newBuffer = NULL;
          flush(storedKey, prio, true);
          continue;
        }
      }
    } else if (numBuffers[prio] < maxBuffers[prio]) {
      U16 number;
      number = numBuffers[prio]++;
      if (numBuffers[prio] > maxUsedBuffers[prio])
         maxUsedBuffers[prio] = numBuffers[prio];
      MUTEX_UNLOCK(&bufferPoolMutex[prio]);
//      classLog->debug("getQueueBuffer(): add new buffer %u of size %u with prio %u", number, size, prio);
      if (numBuffers[prio] > flushThreshBuffers[prio]) {
//        classLog->debug("getQueueBuffer(): buffer %u, prio %u, count %u, going over thresh - flush one", number, prio, numBuffers[prio]);
        flush(storedKey, prio, true);
      }
      newBuffer = (struct queueBuffer *) malloc(sizeof(struct queueBuffer));
      if (!newBuffer) {
        classLog->error("getQueueBuffer(): failed to malloc new buffer, flush and retry");
        MUTEX_LOCK(&bufferPoolMutex[prio]);
        numBuffers[prio]--;
        MUTEX_UNLOCK(&bufferPoolMutex[prio]);
        flush(storedKey, prio, true);
        continue;
      }
      newBuffer->data = (U8 *) malloc(size);
      if (!newBuffer->data) {
        classLog->error("getQueueBuffer(): failed to malloc data of size %u for new buffer, flush and retry", size);
        MUTEX_LOCK(&bufferPoolMutex[prio]);
        numBuffers[prio]--;
        MUTEX_UNLOCK(&bufferPoolMutex[prio]);
        free(newBuffer);
        newBuffer = NULL;
        flush(storedKey, prio, true);
        continue;
      }
      newBuffer->size = size;
      newBuffer->number = number;
    } else {
      MUTEX_UNLOCK(&bufferPoolMutex[prio]);
      retry++;
      if (retry > 100) {
        classLog->error("getQueueBuffer(): all %u buffers for prio %u in use, too many retries, giving up", maxBuffers[prio], prio);
        return NULL;
      } else if (retry > 10) {
        classLog->debug("getQueueBuffer(): all %u buffers for prio %u in use, wait, flush and retry #%u", maxBuffers[prio], prio, retry);
        Sleep(10 * retry);
        flush(storedKey, prio, true);
      } else {
        classLog->debug("getQueueBuffer(): all %u buffers for prio %u in use, flush and retry #%u", maxBuffers[prio], prio, retry);
        flush(storedKey, prio, true);
      }
    }
  }
  newBuffer->used = 0;
  newBuffer->next = NULL;
  return newBuffer;
}

// return a chain of buffers that are no longer needed to buffer pool
// last buffer in chain must have buffer->next == NULL
void OutStream::returnQueueBuffer(struct queueBuffer * buffer, QPrio prio, time_t now) {
  struct queueBuffer * nextBuffer;

  if(!check_prio(prio, __func__))
    return;

  if(now == 0) {
    struct timeval nowTime;

    gettimeofday(&nowTime, NULL);
    now = nowTime.tv_sec;
  }

//  classLog->debug("returnQueueBuffer(): buffer %u, prio %u", buffer->number, prio);
//  if (buffer->next)
//    classLog->debug("returnQueueBuffer(): buffer %u, prio %u, size %u, used %u, is head of a chain", buffer->number, prio, buffer->size, buffer->used);
  while (buffer) {
    nextBuffer = buffer->next;
    MUTEX_LOCK(&bufferPoolMutex[prio]);
    if (numBuffers[prio] > keepBuffers[prio] && unusedBufferTail[prio] && unusedBufferTail[prio]->lastUsed < now - overKeepTTL) {
      numBuffers[prio]--;
      MUTEX_UNLOCK(&bufferPoolMutex[prio]);
      classLog->verbose("returnQueueBuffer(): buffer %u, prio %u, numBuffers %u, over keep and unusedBufferTail over ttl -> delete", buffer->number, prio, numBuffers[prio]);
      free(buffer->data);
      free(buffer);
      overKeepCount[prio]++;
    } else {
      buffer->next = unusedBufferHead[prio];
      buffer->lastUsed = now;
      unusedBufferHead[prio] = buffer;
      if (!unusedBufferTail[prio])
        unusedBufferTail[prio] = buffer;
      spareBuffers[prio]++;
      MUTEX_UNLOCK(&bufferPoolMutex[prio]);
    }
    buffer = nextBuffer;
  }
}

// Submit a chain of buffers to called layer for sending.
// Last buffer in chain must have next == NULL (preset by getQueueBuffer())
// Returns true if success.
// Queues are per layer.
bool OutStream::submitQueueBuffer(struct queueBuffer * firstBuffer, struct queueBuffer * lastBuffer, int key, QPrio prio) {
  if(!firstBuffer)
    return false;
  if(!check_key(key, __func__))
    return false;
  if(!check_prio(prio, __func__))
    return false;
#if 0
  if (firstBuffer->used == 0) {
    classLog->error("submitQueueBuffer(): firstBuffer->used == 0 for key %u, prio %u, firstBuffer %u, rejecting", key, prio, firstBuffer->number);
    return false;
  }
#endif

//  classLog->debug("submitQueueBuffer(): firstBuffer %u, prio %u, size %u, used %u", firstBuffer->number, prio, firstBuffer->size, firstBuffer->used);
//  if (firstBuffer->next)
//    classLog->debug("submitQueueBuffer(): firstBuffer %u, prio %u, size %u, used %u, is head of a chain", firstBuffer->number, prio, firstBuffer->size, firstBuffer->used);
  if (!lastBuffer) {
    lastBuffer = firstBuffer;
    while (lastBuffer->next) {
      lastBuffer = lastBuffer->next;
#if 0
      if (lastBuffer->used == 0) {
        classLog->error("submitQueueBuffer(): buffer->used == 0 for key %u, prio %u, buffer %u, rejecting", key, prio, lastBuffer->number);
        return false;
      }
      if (lastBuffer->used > lastBuffer->size) {
        classLog->error("submitQueueBuffer(): buffer->used > buffer->size for key %u, prio %u, buffer %u, rejecting", key, prio, lastBuffer->number);
        return false;
      }
#endif
    }
  }

  submitCount[prio]++;

  MUTEX_LOCK(&queueMutex[prio]);
  if(bufferQueueTail[prio]) {
    bufferQueueTail[prio]->next = firstBuffer;
  } else {
    bufferQueueHead[prio] = firstBuffer;
  }
  bufferQueueTail[prio] = lastBuffer;
  MUTEX_UNLOCK(&queueMutex[prio]);
  return true;
}

// get next single buffer out of queue, returns NULL if empty
struct queueBuffer * OutStream::popBuffer (QPrio prio) {
  struct queueBuffer * buffer;

  if(!check_prio(prio, __func__))
    return NULL;
  MUTEX_LOCK(&queueMutex[prio]);
  buffer = bufferQueueHead[prio];
  if (buffer) {
    bufferQueueHead[prio] = buffer->next;
    if (!bufferQueueHead[prio])
      bufferQueueTail[prio] = NULL;
    buffer->next = NULL;
  }
  MUTEX_UNLOCK(&queueMutex[prio]);
  return buffer;
}

// put pop'd single buffer back to head of queue in case lower level failed
void OutStream::pushBuffer (struct queueBuffer * buffer, QPrio prio) {
  if(!buffer || !check_prio(prio, __func__))
    return;

  classLog->debug("pushBuffer(): prio %u, buffer %u", prio, buffer->number);
  MUTEX_LOCK(&queueMutex[prio]);
  buffer->next = bufferQueueHead[prio];
  bufferQueueHead[prio] = buffer;
  if (!bufferQueueTail[prio])
    bufferQueueTail[prio] = buffer;
  MUTEX_UNLOCK(&queueMutex[prio]);
}

// queue empty? No need to lock, we do not change the queue
bool OutStream::queueEmpty (QPrio prio) {
  if(!check_prio(prio, __func__))
    return true;
  return (bufferQueueHead[prio] == NULL);
}

void OutStream::printBufferUsage() {
  for (int q = QPRIOMIN; q <= QPRIOMAX; q++) {
    if (numBuffers[q] > 0)
      classLog->debug("Prio %u uses %u buffers, %u spare", q, numBuffers[q], spareBuffers[q]);
    if (maxUsedBuffers[q] > 0)
      classLog->debug("Prio %u used maximum of %u buffers, keep %u", q, maxUsedBuffers[q], keepBuffers[q]);
  }
}

void OutStream::fillCheckHeader(struct queueBuffer * buffer, U8 header[CHECKHEADERSIZE]) {
  U32 crc;

  MUTEX_LOCK(&serialMutex);
//  classLog->debug("send serial number: %llu, size: %u", serial, buffer->used);
  memcpy(header, &serial, sizeof(U32));
  serial++;
  MUTEX_UNLOCK(&serialMutex);
  crc = crc32(0L, Z_NULL, 0);
  crc = crc32(crc, buffer->data, buffer->used);
  memcpy(header + sizeof(U32), &crc, sizeof(U32));
  memcpy(header + sizeof(U32) + sizeof(U32), &(buffer->used), sizeof(U16));
};
