merp/server.go

223 lines
5.1 KiB
Go

package merp
import "context"
import "encoding/json"
import "fmt"
import "net/http"
import "os"
import "os/signal"
import "regexp"
import "strconv"
import "sync"
import "syscall"
import "time"
import "github.com/didip/tollbooth"
import "github.com/didip/tollbooth_gin"
import "github.com/gin-gonic/gin"
import "github.com/rs/zerolog/log"
import "github.com/dn365/gin-zerolog"
import "github.com/thoas/stats"
import "github.com/astaxie/beego/orm"
import _ "github.com/lib/pq" //revive:disable-line
var thingRegex = regexp.MustCompile(`^[a-zA-Z0-9\_\-\.]+$`)
type MerpTopic string
const (
// Shifts
KiB = 10
MiB = 20
GiB = 30
)
// Server is the central structure of the HTTP API server.
type Server struct {
db orm.Ormer
debug bool
gin *gin.Engine
port uint
server *http.Server
stats *stats.Stats
ll *sync.Mutex // protects listeners below
listeners map[MerpTopic][]*EventListener
}
// NewServer returns a Server, so that you can run the API.
func NewServer() *Server {
ms := new(Server)
ms.ll = new(sync.Mutex)
ms.listeners = make(map[MerpTopic][]*EventListener)
ms.init()
return ms
}
type EventListener struct {
Topic MerpTopic
NewMerpJSONChannel chan string
}
func NewEventListener(t MerpTopic) *EventListener {
el := new(EventListener)
el.init()
el.Topic = t
return el
}
func (el *EventListener) init() {
el.NewMerpJSONChannel = make(chan string)
}
func (ms *Server) AddListener(l *EventListener) {
ms.ll.Lock()
defer ms.ll.Unlock()
if ms.listeners[l.Topic] == nil {
// is this an allocation DoS even with rate-limiting middleware?
ms.listeners[l.Topic] = make([]*EventListener, 0)
}
ms.listeners[l.Topic] = append(ms.listeners[l.Topic], l)
}
func (ms *Server) DelListener(l *EventListener) {
// FIXME(sneak)
}
func (ms *Server) init() {
if os.Getenv("DEBUG") != "" {
ms.debug = true
}
ms.port = 8080
var s uint64
var err error
if os.Getenv("PORT") != "" {
if s, err = strconv.ParseUint(os.Getenv("PORT"), 10, 64); err != nil {
panic("invalid PORT in environment")
}
ms.port = uint(s)
}
ms.connectDB()
gin.DefaultWriter = log.With().Str("component", "gin").Logger()
gin.DefaultErrorWriter = log.With().Str("component", "gin").Logger()
ms.setupRoutes()
ms.server = &http.Server{
Addr: fmt.Sprintf(":%d", ms.port),
Handler: ms.gin,
ReadTimeout: 10 * time.Second,
WriteTimeout: 60 * time.Second,
MaxHeaderBytes: 5 << KiB,
}
}
func (ms *Server) connectDB() {
ms.db = GetDB()
}
// ServeForever causes merp to serve http forever
func (ms *Server) ServeForever() {
// start server
go func() {
err := ms.server.ListenAndServe()
if err != nil {
panic(err)
}
}()
// listen for signals
quit := make(chan os.Signal)
// kill -9 is syscall.SIGKILL but can't be caught, so don't need to add it
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
// blocking wait for signal
<-quit
log.Info().Msg("shutting down server")
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := ms.server.Shutdown(ctx); err != nil {
log.Error().Err(err)
}
// catching ctx.Done(). timeout of 5 seconds.
select {
case <-ctx.Done():
log.Info().Msg("server shutdown")
}
}
func (ms *Server) healthCheckHandler() http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
result := gin.H{
"status": "ok",
"now": time.Now().UTC().Format(time.RFC3339),
}
json, err := json.Marshal(result)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
w.Header().Set("Content-Type", "application/json")
w.Write(json)
}
}
func (ms *Server) statsHandler() http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
stats := ms.stats.Data()
b, _ := json.Marshal(stats)
w.Write(b)
}
}
func (ms *Server) setupRoutes() {
if !ms.debug {
gin.SetMode(gin.ReleaseMode)
}
limiter := tollbooth.NewLimiter(5, nil)
ms.stats = stats.New()
// empty router
r := gin.New()
// wrap panics:
r.Use(gin.Recovery())
// attach logger middleware
r.Use(ginzerolog.Logger("gin"))
r.Use(func(c *gin.Context) {
beginning, recorder := ms.stats.Begin(c.Writer)
c.Next()
ms.stats.End(beginning, stats.WithRecorder(recorder))
})
//FIXME(sneak) use a http.MaxBytesReader middleware to limit request size
r.GET("/.well-known/healthcheck.json", gin.WrapF(ms.healthCheckHandler()))
r.GET("/admin/healthcheck.json", gin.WrapF(ms.healthCheckHandler()))
r.GET("/admin/stats.json", gin.WrapF(ms.statsHandler()))
r.GET("/admin/other.json", gin.WrapF(ms.statsHandler()))
r.GET("/merp/for/:thing", tollbooth_gin.LimitHandler(limiter), ms.handleNewMerp())
r.GET("/listen/for/merps/from/:thing", tollbooth_gin.LimitHandler(limiter), ms.listenForMerps())
r.GET("/get/latest/merp/for/:thing", tollbooth_gin.LimitHandler(limiter), ms.getLatestMerps())
r.GET("/get/latest/merps", tollbooth_gin.LimitHandler(limiter), ms.getLatestMerps())
r.GET("/get/merps/for/:thing", tollbooth_gin.LimitHandler(limiter), ms.getLatestMerps())
r.NoRoute(func(c *gin.Context) {
c.JSON(404, gin.H{"message": "not found"})
})
ms.gin = r
}