hydra.go

   1 package main
   2
   3 import (
   4         _ "embed"
   5         "github.com/yuin/gopher-lua"
   6         "os"
   7         "os/signal"
   8         "runtime/pprof"
   9         "syscall"
  10         "time"
  11 )
  12
  13 var lastTime = time.Now()
  14 var signalChannel chan os.Signal
  15
  16 var serializeVer uint8 = 29
  17 var protoVer uint16 = 40
  18
  19 //go:embed builtin/luax/init.lua
  20 var builtinLuaX string
  21
  22 //go:embed builtin/vector.lua
  23 var builtinVector string
  24
  25 //go:embed builtin/escapes.lua
  26 var builtinEscapes string
  27
  28 //go:embed builtin/client.lua
  29 var builtinClient string
  30
  31 //go:embed builtin/base64.lua
  32 var builtinBase64 string
  33
  34 var builtinFiles = []string{
  35         builtinLuaX,
  36         builtinVector,
  37         builtinEscapes,
  38         builtinClient,
  39         builtinBase64,
  40 }
  41
  42 var hydraFuncs = map[string]lua.LGFunction{
  43         "client": l_client,
  44         "map":    l_map,
  45         "dtime":  l_dtime,
  46         "poll":   l_poll,
  47         "close":  l_close,
  48 }
  49
  50 func l_dtime(l *lua.LState) int {
  51         l.Push(lua.LNumber(time.Since(lastTime).Seconds()))
  52         lastTime = time.Now()
  53         return 1
  54 }
  55
  56 func l_poll(l *lua.LState) int {
  57         return doPoll(l, getClients(l))
  58 }
  59
  60 func l_close(l *lua.LState) int {
  61         for _, client := range getClients(l) {
  62                 client.closeConn()
  63         }
  64
  65         return 0
  66 }
  67
  68 func main() {
  69         if len(os.Args) < 2 {
  70                 panic("missing filename")
  71         }
  72
  73         signalChannel = make(chan os.Signal, 1)
  74         signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
  75
  76         go func() {
  77                 ch := make(chan os.Signal, 1)
  78                 signal.Notify(ch, syscall.SIGUSR1)
  79                 for {
  80                         <-ch
  81                         pprof.Lookup("goroutine").WriteTo(os.Stdout, 1)
  82                 }
  83         }()
  84
  85         l := lua.NewState()
  86         defer l.Close()
  87
  88         arg := l.NewTable()
  89         for i, a := range os.Args {
  90                 l.RawSetInt(arg, i-1, lua.LString(a))
  91         }
  92         l.SetGlobal("arg", arg)
  93
  94         hydra := l.SetFuncs(l.NewTable(), hydraFuncs)
  95         l.SetField(hydra, "BS", lua.LNumber(10.0))
  96         l.SetField(hydra, "serialize_ver", lua.LNumber(serializeVer))
  97         l.SetField(hydra, "proto_ver", lua.LNumber(protoVer))
  98         l.SetGlobal("hydra", hydra)
  99
 100         l.SetField(l.NewTypeMetatable("hydra.client"), "__index", l.NewFunction(l_client_index))
 101         l.SetField(l.NewTypeMetatable("hydra.map"), "__index", l.SetFuncs(l.NewTable(), mapFuncs))
 102
 103         l.SetField(l.NewTypeMetatable("hydra.comp.auth"), "__index", l.SetFuncs(l.NewTable(), compAuthFuncs))
 104         l.SetField(l.NewTypeMetatable("hydra.comp.map"), "__index", l.SetFuncs(l.NewTable(), compMapFuncs))
 105         l.SetField(l.NewTypeMetatable("hydra.comp.pkts"), "__index", l.SetFuncs(l.NewTable(), compPktsFuncs))
 106
 107         for _, str := range builtinFiles {
 108                 if err := l.DoString(str); err != nil {
 109                         panic(err)
 110                 }
 111         }
 112
 113         if err := l.DoFile(os.Args[1]); err != nil {
 114                 panic(err)
 115         }
 116 }