path: root/worker.go

package main

import (
	"errors"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
)

type HTTPRequest struct {
	request *http.Request
	route string
	debug bool
	result chan *HTTPResponse
}

type HTTPResponse struct {
	Code int
	Headers map[string]string
	Body string
}

func HandleHTTPRequest(
	queue chan *HTTPRequest,
	route string,
	req *http.Request,
) chan *HTTPResponse {
	res := make(chan *HTTPResponse)
	queue <- &HTTPRequest{
		request: req,
		route: route,
		result: res,
	}
	return res
}

type Worker struct {
	lua *Lua
	routes map[string]LuaRef
	started bool
	mu sync.Mutex
	evalFn *LuaRef
}

// NewWorker creates a new instance of Worker type.
func NewWorker() *Worker {
	return &Worker {
		routes: make(map[string]LuaRef),
		lua: &Lua{},
	}
}

// Start starts the worker:
// 1) creates a Lua context
// 2) executes the argv in it
// 3) initiates the the "luna" module so it's possible to call Go functions
// from Lua
func (w *Worker) Start(argv []string, module map[string]any) error {
	if len(argv) == 0 {
		return errors.New("argv must at least contain lua file name")
	}
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.started {
		return errors.New("already started")
	}
	w.lua.Start()
	defer w.lua.RestoreStackFunc()()

	// emulate passing arguments to the loaded chunk
	args := []any{}
	if len(argv) > 1 {
		for _, arg := range argv[1:] {
			args = append(args, arg)
		}
	}
	err := w.lua.PushArray(args)
	if err != nil {
		return err
	}
	w.lua.SetGlobal("arg")

	// register the module in the Lua context
	err = w.lua.PushObject(module)
	if err != nil {
		return err
	}
	w.lua.SetGlobal("luna")

	waitCh := make(chan bool)
	w.lua.yield = func () {
		<- waitCh
	}
	w.lua.resume = func () bool {
		waitCh <- true
		return true
	}

	err = w.lua.Require(argv[0])

	if err != nil {
		return err
	}

	w.started = true
	return nil
}

// Listen starts handling HTTP requests from the queue.
func (w *Worker) Listen(queue chan *HTTPRequest) {
	stringListToAny := func(slice []string) []any {
		res := []any{}
		for _, v := range slice {
			res = append(res, v)
		}
		return res
	}
	handle := func(r *HTTPRequest, yield func(), resume func() bool) {
		l := w.lua.NewThread(yield, resume)
		// Save a thread to a reference so it's not garbage collected
		// before we are done with it.
		ref := w.lua.PopToRef()
		defer w.lua.Unref(ref)

		if _, ok := w.routes[r.route]; !ok {
			r.result <- &HTTPResponse{
				Code: 404,
				Headers: make(map[string]string),
				Body: "not found",
			}
			return
		}

		l.PushTracebackHandler()
		l.PushFromRef(w.routes[r.route])
		req := r.request
		res := make(map[string]any)
		res["method"] = req.Method
		res["path"] = req.URL.Path

		fh := make(map[string]any)
		for k := range req.Header {
			fh[k] = req.Header.Get(k)
		}
		res["headers"] = fh

		flatQr := make(map[string]any)
		qr := req.URL.Query()
		for k := range qr {
			flatQr[k] = stringListToAny(qr[k])
		}
		res["query"] = flatQr
		
		// if request body is a multipart form: automatically parse it,
		// save the files and form values and put them in to the
		// request object in the "form" field
		// in all other cases just put the body as a string in the
		// "body" field
		if strings.HasPrefix(
			req.Header.Get("Content-Type"),
			"multipart/form-data",
		) {
			err := req.ParseMultipartForm(0)
			if err != nil {
				r.result <- &HTTPResponse{
					Code: 400,
					Headers: make(map[string]string),
					Body: "bad multipart request",
				}
				return
			}

			form := make(map[string]any)
			for k, v := range req.MultipartForm.File {
				// for now only take the first value
				fh := v[0]
				if fh == nil {
					continue
				}
				fd, err := fh.Open()
				defer fd.Close()
				if err != nil {
					r.result <- &HTTPResponse{
						Code: 500,
						Headers: make(map[string]string),
						Body: "server error",
					}
					log.Println("could not open multipart file:", err)
					return
				}

				// assume fd is a file stored on the filesystem.
				// if it's not the case: write the file from
				// memory to the filesystem.
				f, ok := fd.(*os.File)
				if !ok {
					f, _ = os.CreateTemp(os.TempDir(), "multipart-")
					buf := make([]byte, 8192)
					for {
						nread, _ := fd.Read(buf)
						_, _ = f.Write(buf)
						if nread < 8192 {
							break
						}
					}
					f.Close()
					defer os.Remove(f.Name())
				}
				record := make(map[string]any)
				record["path"] = f.Name()
				record["size"] = fh.Size
				record["name"] = fh.Filename
				form[k] = record
			}
			for k, v := range req.MultipartForm.Value {
				// for now only take the first value
				form[k] = v[0]
			}
			res["form"] = form
		} else {
			body, err := io.ReadAll(req.Body)
			if err != nil {
				r.result <- &HTTPResponse{
					Code: 500,
					Headers: make(map[string]string),
					Body: "server error",
				}
				log.Println("could not read request body:", err)
				return
			}
			res["body"] = string(body)
		}

		err := l.PushObject(res)
		if err != nil {
			r.result <- &HTTPResponse{
				Code: 500,
				Headers: make(map[string]string),
				Body: "server error",
			}
			log.Println("could not form request to lua:", err)
			return
		}

		err = l.PCall(1, 3, -3)
		if err != nil {
			var body string
			if debug {
				body = err.Error()
			} else {
				body = "server error"
			}
			r.result <- &HTTPResponse{
				Code: 500,
				Headers: make(map[string]string),
				Body: body,
			}
			log.Println("could not process request:\n" + err.Error())
			return
		}

		// TODO: probably it would be better to just use l.Scan()
		// here but i'm not really sure if we want to have that
		// overhead here.
		code := l.ToInt(-3)
		rbody := l.ToString(-1)
		// Parse headers.
		headers := make(map[string]string)
		l.Pop(1)
		l.PushNil()
		for l.Next() {
			if !l.IsString(-2) || !l.IsString(-1) {
				l.Pop(1)
				continue
			}
			v := l.ToString(-1)
			l.Pop(1)
			// We must not pop the item key from the stack
			// because otherwise C.lua_next won't work
			// properly.
			k := l.ToString(-1)
			headers[k] = v
		}
		r.result <- &HTTPResponse{
			Code: int(code),
			Headers: headers,
			Body: rbody,
		}
	}

	resCh := make(chan func() bool, 4096)
outer:
	for {
		select {
		case r, ok := <- queue:
			// accept new requests
			if !ok {
				break outer
			}
			resCh <- NewCoroutine(
				func(yield func(), resume func() bool) {
					handle(r, yield, func () bool {
						resCh <- resume
						return true
					})
				},
			)
		case resume, ok := <-resCh:
			// coroutine executor
			if !ok {
				break outer
			}
			resume()
		}
	}
}

// Eval evaluates the code in the Lua context. Not safe for execution when
// there are requests in the processing queue, only meant for development
// purposes.
func (w *Worker) Eval(code string) error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.evalFn != nil {
		w.lua.PushFromRef(*w.evalFn)
		w.lua.PushString(code)
		return w.lua.PCall(1, 0, 0)
	}
	return w.lua.LoadString(code)
}

// SetRoute sets a handler for the route.
func (w *Worker) SetRoute(route string, handler LuaRef) {
	w.routes[route] = handler
}

// Stop stops the worker closing the Lua context. TODO: stop Listen goroutine
// as well.
func (w *Worker) Stop() {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.lua.Close()
}

// HasSameLua checks if the Lua context belongs to the worker.
func (w *Worker) HasSameLua(l *Lua) bool {
	return w.lua == l
}

func NewCoroutine(f func (yield func(), resume func() bool)) (resume func() bool) {
	cin := make(chan bool)
	cout := make(chan bool)
	running := true
	resume = func() bool {
		if !running {
			return false
		}
		cin <- true
		<-cout
		return true
	}
	yield := func() {
		cout <- true
		<-cin
	}
	go func() {
		<-cin
		f(yield, resume)
		running = false
		cout <- true
	}()
	return resume
}