path: root/main.go

package main

import (
	"bufio"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"runtime"
	"runtime/cgo"
	"strings"
	"sync"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

func main() {
	lAddr := flag.String("l", "127.0.0.1:3000", "Address HTTP-server will listen to")
	wrksNum := flag.Int("n", runtime.NumCPU(), "Number of HTTP-workers to start")
	flag.Parse()

	if flag.NArg() == 0 {
		printUsage()
	}
	luaArgv := flag.Args()
	mustExist(luaArgv[0])

	httpClient := &http.Client{}

	// queue for http messages for workers to handle
	msgs := make(chan any, 4096)
	mux := http.NewServeMux()
	wrks := []*Worker{}
	// track routes for mux to avoid registering the same route twice
	routes := make(map[string]bool)
	// track open dbs to close them on exit
	dbs := []*sql.DB{}
	dbmu := sync.Mutex{}
	defer func() {
		for _, db := range dbs {
			if db == nil {
				continue
			}
			db.Close()
		}
	}()

	// define luna.router module
	routeModule := make(map[string]any)
	routeModule["route"] = func (l *Lua) int {
		var route string
		var fn LuaRef
		err := l.Scan(&route, &fn)
		if err != nil {
			return luaErr(l, err)
		}
		// find corresponding worker for the lua context
		var wrk *Worker
		for _, wrk = range wrks {
			if wrk.HasSameLua(l) {
				wrk.SetRoute(route, fn)
				break
			}
		}
		if wrk == nil {
			return luaOk(l, nil)
		}
		// check if route is already registered. otherwise
		// mux.HandleFunc will panic
		_, ok := routes[route]
		if ok {
			return luaOk(l, nil)
		}
		routes[route] = true
		mux.HandleFunc(
			route,
			func (w http.ResponseWriter, r *http.Request) {
				start := time.Now()
				resCh := HandleHTTPRequest(msgs, route, r)
				res := <- resCh
				for k, h := range res.Headers {
					w.Header().Set(k, h)
				}
				w.WriteHeader(int(res.Code))
				fmt.Fprint(w, string(res.Body))
				log.Printf(
					"%s %s (%s) -> %d\n",
					r.Method,
					r.URL.Path,
					time.Now().Sub(start),
					res.Code,
				)
			},
		)
		return luaOk(l, nil)
	}
	routeModule["static"] = func (l *Lua) int {
		var route, dir string
		err := l.Scan(&route, &dir)
		if err != nil {
			return luaErr(l, err)
		}
		// check if route is already registered. otherwise
		// mux.HandleFunc will panic
		_, ok := routes[route]
		if ok {
			return luaOk(l, nil)
		}
		routes[route] = true
		mux.Handle(
			route,
			http.StripPrefix("/static/", http.FileServer(http.Dir(dir))),
		)
		return luaOk(l, nil)
	}

	// define luna.http module
	httpModule := make(map[string]any)
	httpModule["request"] = func (l *Lua) int {
		var method, url, body string
		headers := make(map[string]string)
		err := l.Scan(&method, &url, &headers, &body)
		if err != nil {
			return luaErr(l, err)
		}

		req, err := http.NewRequest(method, url, strings.NewReader(body))
		if err != nil {
			return luaErr(l, err)
		}
		for k, v := range headers {
			req.Header.Add(k, v)
		}
		resp, err := httpClient.Do(req)
		if err != nil {
			return luaErr(l, err)
		}
		defer resp.Body.Close()
		respBody, err := io.ReadAll(resp.Body)
		if err != nil {
			return luaErr(l, err)
		}
		respHeaders := make(map[string]any)
		for k := range resp.Header {
			respHeaders[k] = resp.Header.Get(k)
		}

		res := make(map[string]any)
		res["status"] = resp.StatusCode
		res["headers"] = respHeaders
		res["body"] = string(respBody)
		return luaOk(l, res)
	}

	// define luna.db module
	dbModule := make(map[string]any)
	dbModule["open"] = func (l *Lua) int {
		var file string
		err := l.Scan(&file)
		if err != nil {
			return luaErr(l, err)
		}
		db, err := sql.Open("sqlite3", file)
		if err != nil {
			return luaErr(l, err)
		}
		dbmu.Lock()
		dbs = append(dbs, db)
		dbmu.Unlock()
		h := cgo.NewHandle(db)
		return luaOk(l, int(h))
	}
	dbModule["begin"] = func (l *Lua) int {
		var handle cgo.Handle
		err := l.Scan(&handle)
		if err != nil {
			return luaErr(l, err)
		}
		db := handle.Value().(*sql.DB)
		tx, err := db.Begin()
		if err != nil {
			return luaErr(l, err)
		}
		txh := cgo.NewHandle(tx)
		return luaOk(l, int(txh))
	}
	dbModule["commit"] = func (l *Lua) int {
		var handle cgo.Handle
		err := l.Scan(&handle)
		if err != nil {
			return luaErr(l, err)
		}
		tx := handle.Value().(*sql.Tx)
		err = tx.Commit()
		if err != nil {
			return luaErr(l, err)
		}
		handle.Delete()
		return luaOk(l, nil)
	}
	dbModule["rollback"] = func (l *Lua) int {
		var handle cgo.Handle
		err := l.Scan(&handle)
		if err != nil {
			return luaErr(l, err)
		}
		tx := handle.Value().(*sql.Tx)
		err = tx.Rollback()
		if err != nil {
			return luaErr(l, err)
		}
		handle.Delete()
		return luaOk(l, nil)
	}
	dbModule["exec-tx"] = func (l *Lua) int {
		var handle cgo.Handle
		var query string
		var params []any
		err := l.Scan(&handle, &query, &params)
		if err != nil {
			return luaErr(l, err)
		}
		tx := handle.Value().(*sql.Tx)
		_, err = tx.Exec(query, params...)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, nil)
	}
	dbModule["exec"] = func (l *Lua) int {
		var handle cgo.Handle
		var query string
		var params []any
		err := l.Scan(&handle, &query, &params)
		if err != nil {
			return luaErr(l, err)
		}
		db := handle.Value().(*sql.DB)
		_, err = db.Exec(query, params...)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, nil)
	}
	dbModule["query"] = func (l *Lua) int {
		var handle cgo.Handle
		var query string
		var params []any
		err := l.Scan(&handle, &query, &params)
		if err != nil {
			return luaErr(l, err)
		}
		db := handle.Value().(*sql.DB)
		rows, err := db.Query(query, params...)
		if err != nil {
			return luaErr(l, err)
		}
		var res [][]any
		cols, _ := rows.Columns()
		for rows.Next() {
			scans := make([]any, len(cols))
			for i, _ := range scans {
				scans[i] = &scans[i]
			}
			var row []any
			rows.Scan(scans...)
			for _, v := range scans {
				row = append(row, v)
			}
			res = append(res, row)
		}
		ares := []any{}
		for _, v := range res {
			ares = append(ares, v)
		}
		err = l.PushArray(ares)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, ares)
	}
	dbModule["query*"] = func (l *Lua) int {
		var handle cgo.Handle
		var query string
		var params []any
		err := l.Scan(&handle, &query, &params)
		if err != nil {
			return luaErr(l, err)
		}
		db := handle.Value().(*sql.DB)
		rows, err := db.Query(query, params...)
		if err != nil {
			return luaErr(l, err)
		}
		var res []map[string]any
		cols, _ := rows.Columns()
		for rows.Next() {
			scans := make([]any, len(cols))
			for i, _ := range scans {
				scans[i] = &scans[i]
			}
			row := make(map[string]any)
			rows.Scan(scans...)
			for i, v := range scans {
				row[cols[i]] = v
			}
			res = append(res, row)
		}
		ares := []any{}
		for _, v := range res {
			ares = append(ares, v)
		}
		return luaOk(l, ares)
	}
	dbModule["close"] = func (l *Lua) int {
		var handle cgo.Handle
		err := l.Scan(&handle)
		if err != nil {
			return luaErr(l, err)
		}
		db := handle.Value().(*sql.DB)
		err = db.Close()
		if err != nil {
			return luaErr(l, err)
		}
		dbmu.Lock()
		for k, _db := range dbs {
			if db == _db {
				dbs[k] = nil
				break
			}
		}
		dbmu.Unlock()
		handle.Delete()
		return luaOk(l, nil)
	}

	// define luna.utf module
	utfModule := make(map[string]any)
	utfModule["len"] = func (l *Lua) int {
		var str string
		err := l.Scan(&str)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, len(str))
	}
	utfModule["lower"] = func (l *Lua) int {
		var str string
		err := l.Scan(&str)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, strings.ToLower(str))
	}
	utfModule["sub"] = func (l *Lua) int {
		var str string
		var start, length int
		err := l.Scan(&str, &start, &length)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, str[start-1:start-1+length])
	}

	module := make(map[string]any)
	module["router"] = routeModule
	module["http"] = httpModule
	module["db"] = dbModule
	module["utf"] = utfModule

	wg := sync.WaitGroup{}
	wg.Add(*wrksNum)
	// start workers
	for i := 0; i < *wrksNum; i++ {
		wrk := NewWorker()
		wrks = append(wrks, wrk)
		go func () {
			err := wrk.Start(luaArgv, module)
			if err != nil {
				log.Fatal(err)
			}
			wrks = append(wrks, wrk)
			log.Printf("worker %d started\n", i)
			wg.Add(-1)
			wrk.Listen(msgs)
		}()
		defer wrk.Stop()
	}
	wg.Wait()

	// listen for evals from stdio
	go func () {
		reader := bufio.NewReader(os.Stdin)
		for {
			code, _ := reader.ReadString('\n')
			if len(code) == 0 {
				continue
			}
			for _, wrk := range wrks {
				if err := wrk.Eval(code); err != nil {
					log.Printf("error: %s\n", err)
					break
				}
			}
		}
	}()

	if len(routes) > 0 {
		log.Printf("luna is running at %s...", *lAddr)
		log.Fatal(http.ListenAndServe(*lAddr, mux))
	}
	log.Printf("no HTTP-routes were registered, exiting...")
}

func printUsage() {
	fmt.Printf("usage: %s [luafile]\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(2)
}

func mustExist(file string) {
	if _, err := os.Stat(file); errors.Is(err, os.ErrNotExist) {
		fmt.Printf("file \"%s\" does not exist\n", file)
		os.Exit(1)
	}
}

func luaErr(l *Lua, err error) int {
	l.PushBoolean(false)
	l.PushString(err.Error())
	return 2
}

func luaOk(l *Lua, res any) int {
	l.PushBoolean(true)
	err := l.PushAny(res)
	if err != nil {
		return luaErr(l, fmt.Errorf("luaOk: %s", err))
	}
	return 2
}