/* ***** BEGIN LICENSE BLOCK *****
 * Version: MPL 1.1/GPL 2.0/LGPL 2.1
 *
 * The contents of this file are subject to the Mozilla Public License Version
 * 1.1 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 * http://www.mozilla.org/MPL/
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 *
 * The Original Code is Ubiquity.
 *
 * The Initial Developer of the Original Code is Mozilla.
 * Portions created by the Initial Developer are Copyright (C) 2007
 * the Initial Developer. All Rights Reserved.
 *
 * Contributor(s):
 *   Atul Varma <atul@mozilla.com>
 *
 * Alternatively, the contents of this file may be used under the terms of
 * either the GNU General Public License Version 2 or later (the "GPL"), or
 * the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
 * in which case the provisions of the GPL or the LGPL are applicable instead
 * of those above. If you wish to allow use of your version of this file only
 * under the terms of either the GPL or the LGPL, and not to allow others to
 * use your version of this file under the terms of the MPL, indicate your
 * decision by deleting the provisions above and replace them with the notice
 * and other provisions required by the GPL or the LGPL. If you do not delete
 * the provisions above, a recipient may use your version of this file under
 * the terms of any one of the MPL, the GPL or the LGPL.
 *
 * ***** END LICENSE BLOCK ***** */

// This file is based on the code in the JSAPI User Guide:
//
// https://developer.mozilla.org/En/SpiderMonkey/JSAPI_User_Guide

#include "string.h"
#include "jsapi.h"
#include "jsdbgapi.h"
#include "jsdhash.h"

#include "wrapper.h"
#include "server_socket.h"

// The name of our JS script containing our Trusted Code Base (TCB).
#define TCB_FILENAME "tcb.js"

// Global references to our TCB.
static JSContext *tcb_cx;
static JSObject  *tcb_global;

// TODO: Make sure we're rooting objects appropriately here so that
// the interpreter doesn't randomly crash or something.

/* The class of the global object. */
static JSClass global_class = {
  "global", JSCLASS_GLOBAL_FLAGS,
  JS_PropertyStub, JS_PropertyStub, JS_PropertyStub, JS_PropertyStub,
  JS_EnumerateStub, JS_ResolveStub, JS_ConvertStub, JS_FinalizeStub,
  JSCLASS_NO_OPTIONAL_MEMBERS
};

// This native JS function prints the given string to the console.

static JSBool print(JSContext *cx, JSObject *obj, uintN argc, jsval *argv,
                    jsval *rval)
{
  char *str;

  if (!JS_ConvertArguments(cx, argc, argv, "s", &str))
    return JS_FALSE;

  printf("%s\n", str);

  return JS_TRUE;
}

// Private structure to track the state of tracing the JS heap.

typedef struct TracingState {
  // Keeps track of what objects we've visited so far.
  JSDHashTable visited;

  // Whether the tracing operation is successful or failed.
  JSBool result;

  // Class to look for.
  JSClass *classp;
};

// Static singleton for tracking the state of tracing the JS heap.
static TracingState tracingState;

// JSTraceCallback for heap dumping and other operations.

static void traceCallback(JSTracer *trc, void *thing, uint32 kind)
{
  switch (kind) {
  case JSTRACE_OBJECT:
    JSDHashEntryStub *entry = (JSDHashEntryStub *)
      JS_DHashTableOperate(&tracingState.visited,
                           thing,
                           JS_DHASH_LOOKUP);
    if (JS_DHASH_ENTRY_IS_FREE((JSDHashEntryHdr *)entry)) {
      entry = (JSDHashEntryStub *) JS_DHashTableOperate(&tracingState.visited,
                                                        thing,
                                                        JS_DHASH_ADD);
      if (entry == NULL) {
        JS_ReportOutOfMemory(trc->context);
        tracingState.result = JS_FALSE;
        return;
      }
      entry->key = thing;
      JSObject *obj = (JSObject *) thing;
      JSClass *classp = JS_GET_CLASS(trc->context, obj);
      if (classp == tracingState.classp) {
        JSObject *constructor = JS_GetConstructor(trc->context, obj);
        if (constructor == NULL) {
          tracingState.result = JS_FALSE;
          return;
        }
        if (JS_ObjectIsFunction(trc->context, constructor)) {
          JSFunction *func = JS_ValueToFunction(
            trc->context,
            OBJECT_TO_JSVAL(constructor)
            );
          if (func == NULL) {
            tracingState.result = JS_FALSE;
            return;
          }
          JSString *name = JS_GetFunctionId(func);
          if (name != NULL) {
            const char *str = JS_GetStringBytes(name);
            printf("function: %s\n", str);
          }
        }
      }
      JS_TraceChildren(trc, thing, kind);
    }
    break;
  case JSTRACE_DOUBLE:
    break;
  case JSTRACE_STRING:
    break;
  }
}

// This native JS function dumps the heap to stdout. It's available even on
// non-debug builds of SpiderMonkey, but it's also not terribly robust.

static JSBool dumpHeap(JSContext *cx, JSObject *obj, uintN argc, jsval *argv,
                       jsval *rval)
{
  JSObject *object;

  if (!JS_ConvertArguments(cx, argc, argv, "o", &object))
    return JS_FALSE;

  tracingState.classp = JS_GET_CLASS(cx, object);

  if (tracingState.classp == NULL) {
    JS_ReportError(cx, "Object has no class.");
    return JS_FALSE;
  }

  JSTracer tracer;

  if (!JS_DHashTableInit(&tracingState.visited, JS_DHashGetStubOps(),
                         NULL, sizeof(JSDHashEntryStub),
                         JS_DHASH_DEFAULT_CAPACITY(100))) {
    JS_ReportOutOfMemory(cx);
    return JS_FALSE;
  }

  tracingState.result = JS_TRUE;
  tracer.context = cx;
  tracer.callback = traceCallback;
  JS_TraceRuntime(&tracer);

  JS_DHashTableFinish(&tracingState.visited);

  if (!tracingState.result)
    return JS_FALSE;
  *rval = JSVAL_VOID;
  return JS_TRUE;
}

// This native JS function forces garbage collection.

static JSBool forceGC(JSContext *cx, JSObject *obj, uintN argc, jsval *argv,
                      jsval *rval)
{
  JS_GC(cx);
  return JS_TRUE;
}

// This native JS function is an implementation of the SecurableModule
// require() function.  For more information, see:
//
// https://wiki.mozilla.org/ServerJS/Modules/SecurableModules
//
// The function takes two parameters: a filename to import, and a JS
// object containing objects to export into the module.  The latter is
// accessible from the loaded module via a global called 'imports'.
//
// This function returns the SecurableModule's 'exports' global.

static JSBool require(JSContext *cx, JSObject *obj, uintN argc, jsval *argv,
                      jsval *rval)
{
  char *filename;
  JSObject *exports;

  if (!JS_ConvertArguments(cx, argc, argv, "so", &filename, &exports))
    return JS_FALSE;

  FILE *f = fopen(filename, "r");
  if (!f) {
    JS_ReportError(cx, "File not found");
    return JS_FALSE;
  }

  fseek(f, 0, SEEK_END);
  long fileSize = ftell(f);
  fseek(f, 0, SEEK_SET);

  char source[fileSize];
  fread(source, fileSize, 1, f);
  fclose(f);

  // TODO: Check for return values here.
  JSContext *module_cx = JS_NewContext(JS_GetRuntime(cx), 8192);
  JSObject *module_global = JS_NewObject(module_cx, &global_class, NULL,
                                         NULL);
  JS_InitStandardClasses(module_cx, module_global);
  JS_DefineProperty(module_cx, module_global, "imports",
                    OBJECT_TO_JSVAL(exports), NULL, NULL, 0);
  JSObject *module_exports = JS_NewObject(module_cx, NULL, NULL,
                                          module_global);
  JS_DefineProperty(module_cx, module_global, "exports",
                    OBJECT_TO_JSVAL(module_exports), NULL, NULL, 0);

  jsval module_rval;
  if (!JS_EvaluateScript(module_cx, module_global, source,
                         strlen(source), filename, 1, &module_rval)) {
    JS_DestroyContext(module_cx);
    return JS_FALSE;
  }

  *rval = OBJECT_TO_JSVAL(module_exports);

  JS_DestroyContext(module_cx);

  return JS_TRUE;
}

// This native JS function is a wrapper for JS_Enumerate().

static JSBool enumerate(JSContext *cx, JSObject *obj, uintN argc,
                        jsval *argv, jsval *rval)
{
  JSObject *target;

  if (!JS_ConvertArguments(cx, argc, argv, "o", &target))
    return JS_FALSE;

  JSIdArray *ids = JS_Enumerate(cx, target);

  if (ids == NULL)
    return JS_FALSE;

  JSObject *array = JS_NewArrayObject(cx, ids->length, ids->vector);
  *rval = OBJECT_TO_JSVAL(array);
  return JS_TRUE;
}

// This native JS function looks up the property of an object, bypassing
// security checks and getters/setters.

static JSBool lookupProperty(JSContext *cx, JSObject *obj, uintN argc,
                             jsval *argv, jsval *rval)
{
  JSObject *target;

  if (argc < 2) {
    JS_ReportError(cx, "Must provide id to lookup.");
    return JS_FALSE;
  }

  if (!JS_ConvertArguments(cx, argc, argv, "o", &target))
    return JS_FALSE;

  return JS_LookupPropertyById(cx, target, argv[1], rval);
}

// This native JS function returns a JS object containing metadata about
// the given function.

static JSBool functionInfo(JSContext *cx, JSObject *obj, uintN argc,
                           jsval *argv, jsval *rval)
{
  JSFunction *func;

  if (!JS_ConvertArguments(cx, argc, argv, "f", &func))
    return JS_FALSE;

  JSScript *script = JS_GetFunctionScript(cx, func);

  if (script == NULL) {
    *rval = JSVAL_NULL;
    return JS_TRUE;
  }
  
  jsval filenameVal = JSVAL_NULL;

  const char *filename = JS_GetScriptFilename(cx, script);
  if (filename) {
    JSString *filenameStr = JS_NewStringCopyZ(cx, filename);
    filenameVal = STRING_TO_JSVAL(filenameStr);
  }

  // TODO: Check for return values here.
  JSObject *funcInfo = JS_NewObject(cx, NULL, NULL, NULL);
  JS_DefineProperty(cx, funcInfo, "filename", filenameVal,
                    NULL, NULL, 0);

  *rval = OBJECT_TO_JSVAL(funcInfo);
  return JS_TRUE;
}

// This native JS function returns a JS representation of the current
// state of the stack, starting with the callee's stack frame and going
// up from there.

static JSBool stack(JSContext *cx, JSObject *obj, uintN argc, jsval *argv,
                    jsval *rval)
{
  JSStackFrame *iterator = NULL;
  JSStackFrame *frame;
  JSObject *prevFrameInfo = NULL;
  JSObject *firstFrameInfo = NULL;
  bool skippedMyFrame = false;

  if (obj == NULL)
    // We're being called from native code, don't skip the topmost frame.
    skippedMyFrame = true;

  while ((frame = JS_FrameIterator(cx, &iterator)) != NULL) {
    if (!skippedMyFrame) {
      skippedMyFrame = true;
      continue;
    }

    jsval functionNameVal = JSVAL_NULL;
    jsval filenameVal = JSVAL_NULL;
    jsval lineNoVal = JSVAL_ZERO;
    jsval functionObjectVal = JSVAL_NULL;
    jsval scopeChainVal = JSVAL_NULL;

    JSFunction *func = JS_GetFrameFunction(cx, frame);
    if (func) {
      JSString *functionName = JS_GetFunctionId(func);
      if (functionName)
        functionNameVal = STRING_TO_JSVAL(functionName);
    }

    if (!JS_IsNativeFrame(cx, frame)) {
      JSScript *script = JS_GetFrameScript(cx, frame);
      jsbytecode *bytecode = JS_GetFramePC(cx, frame);

      const char *filename = JS_GetScriptFilename(cx, script);
      if (filename) {
        JSString *filenameStr = JS_NewStringCopyZ(cx, filename);
        filenameVal = STRING_TO_JSVAL(filenameStr);
      }

      uintN lineNo = JS_PCToLineNumber(cx, script, bytecode);
      lineNoVal = INT_TO_JSVAL(lineNo);

      JSObject *functionObject = JS_GetFrameFunctionObject(cx, frame);
      functionObjectVal = OBJECT_TO_JSVAL(functionObject);

      JSObject *scopeChain = JS_GetFrameScopeChain(cx, frame);
      scopeChainVal = OBJECT_TO_JSVAL(scopeChain);
    }

    // TODO: Check for return values here.
    JSObject *frameInfo = JS_NewObject(cx, NULL, NULL, NULL);
    JS_DefineProperty(cx, frameInfo, "filename", filenameVal,
                      NULL, NULL, 0);
    JS_DefineProperty(cx, frameInfo, "lineNo", lineNoVal,
                      NULL, NULL, 0);
    JS_DefineProperty(cx, frameInfo, "functionName", functionNameVal,
                      NULL, NULL, 0);
    JS_DefineProperty(cx, frameInfo, "functionObject", functionObjectVal,
                      NULL, NULL, 0);
    JS_DefineProperty(cx, frameInfo, "scopeChain", scopeChainVal,
                      NULL, NULL, 0);

    if (prevFrameInfo)
      JS_DefineProperty(cx, prevFrameInfo, "caller",
                        OBJECT_TO_JSVAL(frameInfo), NULL, NULL, 0);
    else
      firstFrameInfo = frameInfo;

    prevFrameInfo = frameInfo;
  }

  if (firstFrameInfo)
    *rval = OBJECT_TO_JSVAL(firstFrameInfo);
  else
    *rval = JSVAL_NULL;

  return JS_TRUE;
}

// Our global hook called whenever an exception is thrown saves the current
// exception and stack information to the 'lastException' and
// 'lastExceptionTraceback' globals of the TCB, respectively, much like
// Python's sys.exc_info().

static JSTrapStatus throwHook(JSContext *cx, JSScript *script, jsbytecode *pc,
                              jsval *rval, void *closure)
{
  JSObject *tcb_global = (JSObject *) closure;
  jsval lastExceptionTraceback;
  jsval lastException;

  jsval exception = *rval;
  if (JS_IsExceptionPending(cx))
    if (!JS_GetPendingException(cx, &exception))
      printf("Getting exception failed.\n");

  if (!JS_GetProperty(cx, tcb_global, "lastException", &lastException))
    printf("Unable to retrieve last exception.");

  if (lastException == exception)
    // The same exception is just propagating through the stack; keep
    // our existing last-exception info.
    return JSTRAP_CONTINUE;

  if (!stack(cx, NULL, 0, NULL, &lastExceptionTraceback)) {
    printf("Generation of exception info failed.");
    lastExceptionTraceback = JSVAL_NULL;
  }

  if (!JS_SetProperty(cx, tcb_global, "lastExceptionTraceback",
                      &lastExceptionTraceback) ||
      !JS_SetProperty(cx, tcb_global, "lastException", &exception))
    printf("Setting of exception info failed.");

  return JSTRAP_CONTINUE;
}

// Our global checkAccess callback just checks to see if a JS function called
// 'checkAccess' has been defined in the TCB, and delegates to that if so. If
// not, though, we do a default checkAccess.

static JSBool checkAccess(JSContext *cx, JSObject *obj, jsval id,
                          JSAccessMode mode, jsval *vp)
{
  jsval checkAccess;
  if (tcb_global && JS_GetProperty(tcb_cx, tcb_global, "checkAccess",
                                   &checkAccess) &&
      JSVAL_IS_OBJECT(checkAccess) &&
      JS_ObjectIsFunction(tcb_cx, JSVAL_TO_OBJECT(checkAccess))) {
    jsval argv[2];
    argv[0] = OBJECT_TO_JSVAL(obj);
    argv[1] = id;
    return JS_CallFunctionValue(tcb_cx, tcb_global, checkAccess, 2, argv,
                                vp);
  }

  // TODO: This doesn't work for the 'caller' attribute, and possibly other
  // things too.
  return JS_LookupPropertyById(cx, obj, id, vp);
}

static JSFunctionSpec tcb_global_functions[] = {
  JS_FS("print",          print,              1, 0, 0),
  JS_FS("getWrapper",     getWrapper,         1, 0, 0),
  JS_FS("unwrap",         unwrapObject,       1, 0, 0),
  JS_FS("wrap",           wrapObject,         2, 0, 0),
  JS_FS("stack",          stack,              0, 0, 0),
  JS_FS("require",        require,            2, 0, 0),
  JS_FS("lookupProperty", lookupProperty,     2, 0, 0),
  JS_FS("functionInfo",   functionInfo,       1, 0, 0),
  JS_FS("enumerate",      enumerate,          1, 0, 0),
  JS_FS("forceGC",        forceGC,            0, 0, 0),
  JS_FS("dumpHeap",       dumpHeap,           0, 0, 0),
  JS_FS("Socket",         createServerSocket, 0, 0, 0),
  JS_FS_END
};

static JSSecurityCallbacks securityCallbacks = {
  checkAccess,
  NULL,
  NULL
};

int main(int argc, const char *argv[])
{
  /* JS variables. */
  JSRuntime *rt;

  /* Create a JS runtime. */
  rt = JS_NewRuntime(8L * 1024L * 1024L);
  if (rt == NULL)
    return 1;

  JS_SetRuntimeSecurityCallbacks(rt, &securityCallbacks);

  /* Create a context. */
  tcb_cx = JS_NewContext(rt, 8192);
  if (tcb_cx == NULL)
    return 1;
  JS_SetOptions(tcb_cx, JSOPTION_VAROBJFIX);
  JS_SetVersion(tcb_cx, JSVERSION_LATEST);

  /* Create the TCB's global object. */
  tcb_global = JS_NewObject(tcb_cx, &global_class, NULL, NULL);
  if (tcb_global == NULL)
    return 1;

  /* Populate the tcb_global object with the standard globals,
     like Object and Array. */
  if (!JS_InitStandardClasses(tcb_cx, tcb_global))
    return 1;

  if (!JS_DefineFunctions(tcb_cx, tcb_global, tcb_global_functions))
    return 1;

  // TODO: Check for return values here.
  JS_SetThrowHook(rt, throwHook, tcb_global);
  JS_DefineProperty(tcb_cx, tcb_global, "lastExceptionTraceback", JSVAL_NULL,
                    NULL, NULL, 0);
  JS_DefineProperty(tcb_cx, tcb_global, "lastException", JSVAL_NULL,
                    NULL, NULL, 0);

  /* Your application code here. This may include JSAPI calls
     to create your own custom JS objects and run scripts. */
  FILE *f = fopen(TCB_FILENAME, "r");
  if (!f)
    return 1;

  fseek(f, 0, SEEK_END);
  long fileSize = ftell(f);
  fseek(f, 0, SEEK_SET);
  
  char source[fileSize];
  fread(source, fileSize, 1, f);
  fclose(f);

  jsval rval;
  if (!JS_EvaluateScript(tcb_cx, tcb_global, source,
                         strlen(source), TCB_FILENAME, 1,
                         &rval)) {
    jsval handleError;
    if (JS_GetProperty(tcb_cx, tcb_global, "handleError", &handleError) &&
        JSVAL_IS_OBJECT(handleError) &&
        JS_ObjectIsFunction(tcb_cx, JSVAL_TO_OBJECT(handleError))) {
      JS_CallFunctionValue(tcb_cx, tcb_global, handleError, 0, NULL, &rval);
    }
    return 1;
  }

  /* Cleanup. */
  JS_DestroyContext(tcb_cx);
  JS_DestroyRuntime(rt);
  JS_ShutDown();

  printf("Farewell.\n");
  return 0;
}