ws_bucket/api/api.go

package api

import (
	"encoding/json"
	"github.com/buaazp/fasthttprouter"
	"github.com/fasthttp/websocket"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/postgres"
	"github.com/sirupsen/logrus"
	"github.com/valyala/fasthttp"
	"os"
	"path/filepath"
)

var WorkDir, _ = filepath.Abs("./data/")

type Info struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

var info = Info{
	Name:    "ws_bucket",
	Version: "1.0",
}

var motd = WebsocketMotd{
	Info: info,
	Motd: "Hello, world",
}

type WebApi struct {
	server      fasthttp.Server
	db          *gorm.DB
	MotdMessage *websocket.PreparedMessage
}

func Index(ctx *fasthttp.RequestCtx) {
	Json(info, ctx)
}

func Json(object interface{}, ctx *fasthttp.RequestCtx) {

	resp, err := json.Marshal(object)
	if err != nil {
		panic(err)
	}

	ctx.Response.Header.Set("Content-Type", "application/json")
	_, err = ctx.Write(resp)
	if err != nil {
		panic(err)
	}
}

func LogRequestMiddleware(h fasthttp.RequestHandler) fasthttp.RequestHandler {
	return fasthttp.RequestHandler(func(ctx *fasthttp.RequestCtx) {

		logrus.WithFields(logrus.Fields{
			"path":   string(ctx.Path()),
			"header": ctx.Request.Header.String(),
		}).Trace(string(ctx.Method()))

		h(ctx)
	})
}

func New(db *gorm.DB) *WebApi {

	api := &WebApi{}

	logrus.SetLevel(getLogLevel())

	router := fasthttprouter.New()
	router.GET("/", LogRequestMiddleware(Index))

	router.POST("/client", LogRequestMiddleware(api.CreateClient))

	router.POST("/slot", LogRequestMiddleware(api.AllocateUploadSlot))
	router.GET("/slot", LogRequestMiddleware(api.ReadUploadSlot))
	router.GET("/upload", LogRequestMiddleware(api.Upload))

	api.server = fasthttp.Server{
		Handler: router.Handler,
		Name:    "ws_bucket",
	}

	api.db = db
	db.AutoMigrate(&Client{})
	db.AutoMigrate(&UploadSlot{})

	api.setupMotd()

	return api
}

func (api *WebApi) setupMotd() {
	var data []byte
	data, _ = json.Marshal(motd)
	motdMsg, _ := websocket.NewPreparedMessage(websocket.TextMessage, data)
	api.MotdMessage = motdMsg
}

func (api *WebApi) Run() {
	address := GetServerAddress()

	logrus.WithFields(logrus.Fields{
		"addr": address,
	}).Info("Starting web server")

	err := api.server.ListenAndServe(address)
	if err != nil {
		logrus.Fatalf("Error in ListenAndServe: %s", err)
	}
}

func GetServerAddress() string {
	serverAddress := os.Getenv("WS_BUCKET_ADDR")
	if serverAddress == "" {
		serverAddress = "0.0.0.0:3020"
	}
	return serverAddress
}

func getLogLevel() logrus.Level {
	levelStr := os.Getenv("WS_BUCKET_LOGLEVEL")
	if levelStr == "" {
		return logrus.TraceLevel
	} else {
		level, err := logrus.ParseLevel(levelStr)
		if err != nil {
			panic(err)
		}
		return level
	}
}