merp/server.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
}