minimum viable (excluding auth)

api/api.go

@@ -0,0 +1,136 @@
+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
+	}
+}

api/auth.go

@@ -0,0 +1,62 @@
+package api
+
+import (
+	"encoding/json"
+	"github.com/sirupsen/logrus"
+	"github.com/valyala/fasthttp"
+	"math/rand"
+)
+
+func (api *WebApi) CreateClient(ctx *fasthttp.RequestCtx) {
+
+	//TODO: auth
+
+	req := &CreateClientRequest{}
+	err := json.Unmarshal(ctx.Request.Body(), req)
+	if err != nil {
+		ctx.Response.SetStatusCode(400)
+		Json(CreateClientResponse{
+			Ok: false,
+		}, ctx)
+		return
+	}
+
+	if !req.IsValid() {
+		ctx.Response.SetStatusCode(400)
+		Json(CreateClientResponse{
+			Ok: false,
+		}, ctx)
+		return
+	}
+
+	client := api.createClient(req)
+
+	Json(CreateClientResponse{
+		Ok:     true,
+		Secret: client.Secret,
+	}, ctx)
+}
+
+func (api *WebApi) createClient(req *CreateClientRequest) *Client {
+
+	client := &Client{
+		Alias:  req.Alias,
+		Secret: genSecret(),
+	}
+
+	api.db.Create(client)
+
+	logrus.WithFields(logrus.Fields{
+		"client": client,
+	}).Info("Created client")
+
+	return client
+}
+
+func genSecret() string {
+	bytes := make([]byte, 32)
+	for i := 0; i < 32; i++ {
+		bytes[i] = byte(48 + rand.Intn(122-48))
+	}
+	return string(bytes)
+}

api/models.go

@@ -0,0 +1,71 @@
+package api
+
+import (
+	"path/filepath"
+	"strings"
+)
+
+type GenericResponse struct {
+	Ok bool `json:"ok"`
+}
+
+type CreateClientRequest struct {
+	Alias string `json:"alias"`
+}
+
+func (req *CreateClientRequest) IsValid() bool {
+	return len(req.Alias) > 3
+}
+
+type CreateClientResponse struct {
+	Ok     bool   `json:"ok"`
+	Secret string `json:"secret,omitempty"`
+}
+
+type Client struct {
+	ID     int64
+	Alias  string `json:"alias"`
+	Secret string `json:"secret"`
+}
+
+type AllocateUploadSlotRequest struct {
+	Token    string `json:"token"`
+	MaxSize  int64  `json:"max_size"`
+	FileName string `json:"file_name"`
+}
+
+func (req *AllocateUploadSlotRequest) IsValid() bool {
+	if len(req.Token) < 3 {
+		return false
+	}
+
+	if len(req.FileName) <= 0 {
+		return false
+	}
+
+	path := filepath.Join(WorkDir, req.FileName)
+	pathAbs, err := filepath.Abs(path)
+	if err != nil {
+		return false
+	}
+	if !strings.HasPrefix(pathAbs, WorkDir) {
+		return false
+	}
+
+	if req.MaxSize < 0 {
+		return false
+	}
+
+	return true
+}
+
+type UploadSlot struct {
+	MaxSize  int64  `json:"max_size"`
+	Token    string `gorm:"primary_key",json:"token"`
+	FileName string `json:"file_name"`
+}
+
+type WebsocketMotd struct {
+	Info Info   `json:"info"`
+	Motd string `json:"motd"`
+}

api/slot.go

@@ -0,0 +1,180 @@
+package api
+
+import (
+	"encoding/json"
+	"github.com/fasthttp/websocket"
+	"github.com/sirupsen/logrus"
+	"github.com/valyala/fasthttp"
+	"io"
+	"os"
+	"path/filepath"
+	"sync"
+)
+
+const WsBufferSize = 4096
+
+var Mutexes sync.Map
+var upgrader = websocket.FastHTTPUpgrader{
+	ReadBufferSize:    WsBufferSize,
+	WriteBufferSize:   WsBufferSize,
+	EnableCompression: true,
+}
+
+func (api *WebApi) AllocateUploadSlot(ctx *fasthttp.RequestCtx) {
+
+	req := &AllocateUploadSlotRequest{}
+	err := json.Unmarshal(ctx.Request.Body(), req)
+	if err != nil {
+		ctx.Response.SetStatusCode(400)
+		Json(GenericResponse{
+			Ok: false,
+		}, ctx)
+		return
+	}
+
+	if !req.IsValid() {
+		ctx.Response.SetStatusCode(400)
+		Json(CreateClientResponse{
+			Ok: false,
+		}, ctx)
+		return
+	}
+
+	api.allocateUploadSlot(req)
+
+	Json(CreateClientResponse{
+		Ok: true,
+	}, ctx)
+}
+
+func (api *WebApi) Upload(ctx *fasthttp.RequestCtx) {
+
+	token := string(ctx.Request.Header.Peek("X-Upload-Token"))
+	slot := UploadSlot{}
+	err := api.db.Where("token=?", token).First(&slot).Error
+	if err != nil {
+		ctx.Response.Header.SetStatusCode(400)
+		logrus.WithFields(logrus.Fields{
+			"token": token,
+		}).Warning("Upload slot not found")
+		return
+	}
+
+	logrus.WithFields(logrus.Fields{
+		"slot": slot,
+	}).Info("Upgrading connection")
+
+	err = upgrader.Upgrade(ctx, func(ws *websocket.Conn) {
+		defer ws.Close()
+
+		err := ws.WritePreparedMessage(api.MotdMessage)
+		if err != nil {
+			panic(err)
+		}
+
+		mt, reader, err := ws.NextReader()
+		if err != nil {
+			panic(err)
+		}
+		if mt != websocket.BinaryMessage {
+			return
+		}
+
+		mu, _ := Mutexes.LoadOrStore(slot.Token, &sync.RWMutex{})
+		mu.(*sync.RWMutex).Lock()
+		path := filepath.Join(WorkDir, slot.FileName)
+		fp, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+		if err != nil {
+			panic(err)
+		}
+
+		buf := make([]byte, WsBufferSize)
+		totalRead := int64(0)
+		for totalRead < slot.MaxSize {
+			read, err := reader.Read(buf)
+
+			var toWrite int
+			if totalRead+int64(read) > slot.MaxSize {
+				toWrite = int(slot.MaxSize - totalRead)
+			} else {
+				toWrite = read
+			}
+
+			_, _ = fp.Write(buf[:toWrite])
+			if err == io.EOF {
+				break
+			}
+			totalRead += int64(read)
+		}
+
+		logrus.WithFields(logrus.Fields{
+			"totalRead": totalRead,
+		}).Info("Finished reading")
+		err = fp.Close()
+		if err != nil {
+			panic(err)
+		}
+		mu.(*sync.RWMutex).Unlock()
+	})
+	if err != nil {
+		panic(err)
+	}
+}
+
+func (api *WebApi) ReadUploadSlot(ctx *fasthttp.RequestCtx) {
+
+	tokenStr := string(ctx.Request.Header.Peek("X-Upload-Token"))
+
+	slot := UploadSlot{}
+	err := api.db.Where("token=?", tokenStr).First(&slot).Error
+
+	if err != nil {
+		ctx.Response.Header.SetStatusCode(404)
+		logrus.WithFields(logrus.Fields{
+			"token": tokenStr,
+		}).Warning("Upload slot not found")
+		return
+	}
+
+	logrus.WithFields(logrus.Fields{
+		"slot": slot,
+	}).Info("Reading")
+
+	path := filepath.Join(WorkDir, slot.FileName)
+
+	mu, _ := Mutexes.LoadOrStore(slot.Token, &sync.RWMutex{})
+	mu.(*sync.RWMutex).RLock()
+	fp, err := os.OpenFile(path, os.O_RDONLY, 0600)
+	if err != nil {
+		panic(err)
+	}
+
+	buf := make([]byte, WsBufferSize)
+	response := ctx.Response.BodyWriter()
+	for {
+		read, err := fp.Read(buf)
+		_, _ = response.Write(buf[:read])
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			panic(err)
+		}
+	}
+	mu.(*sync.RWMutex).RUnlock()
+}
+
+func (api *WebApi) allocateUploadSlot(req *AllocateUploadSlotRequest) {
+
+	slot := &UploadSlot{
+		MaxSize:  req.MaxSize,
+		FileName: req.FileName,
+		Token:    req.Token,
+	}
+
+	logrus.WithFields(logrus.Fields{
+		"slot": slot,
+	}).Info("Allocated new upload slot")
+
+	api.db.Create(slot)
+}