path: root/main.go

package main

import (
	"bufio"
	"database/sql"
	"crypto/sha1"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"image"
	"image/jpeg"
	"image/png"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"unicode/utf8"
	"runtime"
	"runtime/cgo"
	"strings"
	"sync"
	"time"

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

var debug bool

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")
	dbg := flag.Bool("D", false, "Debug/development mode")
	flag.Parse()
	debug = *dbg

	if flag.NArg() == 0 {
		printUsage()
	}
	luaArgv := flag.Args()
	mustExist(luaArgv[0])
	httpClient := &http.Client{
		// do not follow redirects
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	// queue for http messages for workers to handle
	msgs := make(chan *HTTPRequest, 4096)
	mux := http.NewServeMux()
	wrks := []*Worker{}
	// track routes for mux to avoid registering the same route twice
	routes := make(map[string]bool)
	// since mattn/sqlite3 library doesn't do connection pooling we have
	// to do it ourselves. otherwise sqlite becomes the bottleneck when
	// there are many parallel requests to the same DB.
	dbs := make(map[string]*Pool[sql.DB])
	mu := sync.Mutex{}

	// 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)
		}
		mu.Lock()
		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,
				)
			},
		)
		mu.Unlock()
		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
		mu.Lock()
		mux.Handle(
			route,
			http.StripPrefix(
				strings.TrimPrefix(route, "GET "),
				http.FileServer(http.Dir(dir)),
			),
		)
		mu.Unlock()
		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)
	}
	httpModule["encode-url"] = func (l *Lua) int {
		var str string
		err := l.Scan(&str)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, url.QueryEscape(str))
	}
	httpModule["encodeURL"] = httpModule["encode-url"]
	httpModule["decode-url"] = func (l *Lua) int {
		var str string
		err := l.Scan(&str)
		if err != nil {
			return luaErr(l, err)
		}
		res, err := url.QueryUnescape(str)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, res)
	}
	httpModule["decodeURL"] = httpModule["decode-url"]

	// 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)
		}
		mu.Lock()
		pool, ok := dbs[file]
		if !ok {
			pool = &Pool[sql.DB]{}
			pool.New = func () (*sql.DB, error) {
				return sql.Open("sqlite3", file)
			}
			dbs[file] = pool
		}
		// prepopulate pool with at least one connection so we know
		// if there is an error
		db, err := dbs[file].New()
		if err != nil {
			return luaErr(l, err)
		}
		dbs[file].Put(db)
		mu.Unlock()
		return luaOk(l, file)
	}
	dbModule["begin"] = func (l *Lua) int {
		var file string
		err := l.Scan(&file)
		if err != nil {
			return luaErr(l, err)
		}
		db, err := dbs[file].Get()
		if err != nil {
			return luaErr(l, err)
		}
		defer dbs[file].Put(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["execTx"] = dbModule["exec-tx"]
	dbModule["exec"] = func (l *Lua) int {
		var file, query string
		var params []any
		err := l.Scan(&file, &query, &params)
		if err != nil {
			return luaErr(l, err)
		}
		db, err := dbs[file].Get()
		if err != nil {
			return luaErr(l, err)
		}
		defer dbs[file].Put(db)
		_, err = db.Exec(query, params...)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, nil)
	}
	dbModule["query"] = func (l *Lua) int {
		var file, query string
		var params []any
		err := l.Scan(&file, &query, &params)
		if err != nil {
			return luaErr(l, err)
		}
		db, err := dbs[file].Get()
		if err != nil {
			return luaErr(l, err)
		}
		defer dbs[file].Put(db)
		ares := []any{}
		rows, err := db.Query(query, params...)
		if err != nil {
			return luaErr(l, err)
		}
		go func() {
			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)
			}
			for _, v := range res {
				ares = append(ares, v)
			}
			l.resume()
		}()
		l.yield()
		return luaOk(l, ares)
	}
	dbModule["query-assoc"] = func (l *Lua) int {
		var file, query string
		var params []any
		err := l.Scan(&file, &query, &params)
		if err != nil {
			return luaErr(l, err)
		}
		db, err := dbs[file].Get()
		if err != nil {
			return luaErr(l, err)
		}
		defer dbs[file].Put(db)
		rows, err := db.Query(query, params...)
		if err != nil {
			return luaErr(l, err)
		}
		ares := []any{}
		go func() {
			cols, _ := rows.Columns()
			var res []map[string]any
			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)
			}
			for _, v := range res {
				ares = append(ares, v)
			}
			l.resume()
		}()
		l.yield()
		return luaOk(l, ares)
	}
	dbModule["queryAssoc"] = dbModule["query-assoc"]
	dbModule["close"] = func (l *Lua) int {
		var file string
		err := l.Scan(&file)
		if err != nil {
			return luaErr(l, err)
		}
		// FIXME: noop for now
		return luaOk(l, nil)
	}

	// define luna.utf module
	utf8Module := make(map[string]any)
	utf8Module["len"] = func (l *Lua) int {
		var str string
		err := l.Scan(&str)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, utf8.RuneCountInString(str))
	}
	utf8Module["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))
	}
	utf8Module["upper"] = func (l *Lua) int {
		var str string
		err := l.Scan(&str)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, strings.ToUpper(str))
	}
	utf8Module["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)
		}
		runes := []rune(str)
		cut := runes[start-1:start-1+length]
		return luaOk(l, string(cut))
	}

	fsModule := make(map[string]any)
	fsModule["files"] = func (l *Lua) int {
		var dir string
		err := l.Scan(&dir)
		if err != nil {
			return luaErr(l, err)
		}
		files, err := os.ReadDir(dir)
		if err != nil {
			return luaErr(l, err)
		}
		var res []any
		for _, file := range files {
			entry := make(map[string]any)
			entry["name"] = file.Name()
			entry["dir"] = file.IsDir()
			res = append(res, entry)
		}
		return luaOk(l, res)
	}
	fsModule["size"] = func (l *Lua) int {
		var path string
		err := l.Scan(&path)
		if err != nil {
			return luaErr(l, err)
		}
		stat, err := os.Stat(path)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, stat.Size())
	}
	fsModule["mimetype"] = func (l *Lua) int {
		var path string
		err := l.Scan(&path)
		if err != nil {
			return luaErr(l, err)
		}
		file, err := os.Open(path)
		if err != nil {
			return luaErr(l, err)
		}
		defer file.Close()
		res, err := mimetype(file)
		if err != nil {
			return luaErr(l, err)
		}
		return luaOk(l, res)
	}

	imageModule := make(map[string]any)
	imageModule["create-thumbnail"] = func (l *Lua) int {
		var in, out string
		var width int
		err := l.Scan(&in, &out, &width)
		if err != nil {
			return luaErr(l, err)
		}
		file, err := os.Open(in)
		if err != nil {
			return luaErr(l, err)
		}
		defer file.Close()
		mimetype, err := mimetype(file)
		if err != nil {
			return luaErr(l, err)
		}
		var img image.Image
		if mimetype == "image/jpeg" {
			img, err = jpeg.Decode(file)
		} else if mimetype == "image/png" {
			img, err = png.Decode(file)
		} else {
			return luaErr(l, errors.New("unknown image type: only jpeg and png are supported"))
		}
		if err != nil {
			return luaErr(l, err)
		}
		resizedImg := resize.Resize(uint(width), 0, img, resize.Lanczos3)
		outFile, err := os.Create(out)
		if err != nil {
			return luaErr(l, err)
		}
		defer outFile.Close()
		jpeg.Encode(outFile, resizedImg, nil)
		return luaOk(l, nil)
	}
	imageModule["createThumbnail"] = imageModule["create-thumbnail"]

	cryptoModule := make(map[string]any)
	cryptoModule["sha1"] = func (l *Lua) int {
		var str string
		err := l.Scan(&str)
		if err != nil {
			return luaErr(l, err)
		}
		h := sha1.New()
		io.WriteString(h, str)
		return luaOk(l, hex.EncodeToString(h.Sum(nil)))
	}

	module := make(map[string]any)
	module["router"] = routeModule
	module["http"] = httpModule
	module["db"] = dbModule
	module["utf8"] = utf8Module
	module["fs"] = fsModule
	module["image"] = imageModule
	module["crypto"] = cryptoModule
	module["on-eval"] = func (l *Lua) int {
		var evalFn LuaRef
		err := l.Scan(&evalFn)
		if err != nil {
			return luaErr(l, err)
		}
		for _, wrk := range wrks {
			if wrk.HasSameLua(l) {
				wrk.evalFn = &evalFn
				break
			}
		}
		return luaOk(l, nil)
	}
	module["onEval"] = module["on-eval"]
	module["debug"] = debug

	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 commands from stdio
	if debug {
		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 [options] 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
}

func mimetype(f *os.File) (string, error) {
	firstBytes := make([]byte, 512)
	_, err := f.ReadAt(firstBytes, 0)
	if err != nil && err != io.EOF {
		return "", err
	}
	return http.DetectContentType(firstBytes), nil
}