warpbox/lib/server/server.go

package server

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/gin-contrib/gzip"
	"github.com/gin-gonic/gin"
)

const uploadRoot = "data/uploads"

func Run(addr string) error {
	router := gin.Default()
	router.LoadHTMLGlob("templates/*.html")

	router.GET("/", func(ctx *gin.Context) {
		ctx.HTML(http.StatusOK, "index.html", gin.H{})
	})

	router.POST("/box", func(ctx *gin.Context) {
		boxID, err := newBoxID()
		if err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not create upload box"})
			return
		}

		if err := os.MkdirAll(boxPath(boxID), 0755); err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not prepare upload box"})
			return
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id":  boxID,
			"box_url": "/box/" + boxID,
		})
	})

	router.POST("/box/:id/upload", func(ctx *gin.Context) {
		boxID := ctx.Param("id")
		if !validBoxID(boxID) {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "Invalid box id"})
			return
		}

		file, err := ctx.FormFile("file")
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "No file received"})
			return
		}

		savedFile, err := saveUploadedFile(ctx, boxID, file)
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
			return
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id":  boxID,
			"box_url": "/box/" + boxID,
			"file":    savedFile,
		})
	})

	router.POST("/upload", func(ctx *gin.Context) {
		form, err := ctx.MultipartForm()
		if err != nil {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "No files received"})
			return
		}

		files := form.File["files"]
		if len(files) == 0 {
			ctx.JSON(http.StatusBadRequest, gin.H{"error": "No files received"})
			return
		}

		boxID, err := newBoxID()
		if err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not create upload box"})
			return
		}

		if err := os.MkdirAll(boxPath(boxID), 0755); err != nil {
			ctx.JSON(http.StatusInternalServerError, gin.H{"error": "Could not prepare upload box"})
			return
		}

		savedFiles := make([]gin.H, 0, len(files))

		for _, file := range files {
			savedFile, err := saveUploadedFile(ctx, boxID, file)
			if err != nil {
				ctx.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
				return
			}

			savedFiles = append(savedFiles, savedFile)
		}

		ctx.JSON(http.StatusOK, gin.H{
			"box_id":  boxID,
			"box_url": "/box/" + boxID,
			"files":   savedFiles,
		})
	})

	compressed := router.Group("/", gzip.Gzip(gzip.DefaultCompression))
	compressed.Static("/static", "./static")

	return router.Run(addr)
}

func newBoxID() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}

	return hex.EncodeToString(bytes), nil
}

func validBoxID(boxID string) bool {
	if len(boxID) != 32 {
		return false
	}

	for _, character := range boxID {
		if !strings.ContainsRune("0123456789abcdef", character) {
			return false
		}
	}

	return true
}

func boxPath(boxID string) string {
	return filepath.Join(uploadRoot, boxID)
}

func saveUploadedFile(ctx *gin.Context, boxID string, file *multipart.FileHeader) (gin.H, error) {
	filename, ok := safeFilename(file.Filename)
	if !ok {
		return nil, fmt.Errorf("Invalid filename")
	}

	boxPath := boxPath(boxID)
	if err := os.MkdirAll(boxPath, 0755); err != nil {
		return nil, fmt.Errorf("Could not prepare upload box")
	}

	filename = uniqueFilename(boxPath, filename)
	destination := filepath.Join(boxPath, filename)
	if err := ctx.SaveUploadedFile(file, destination); err != nil {
		return nil, fmt.Errorf("Could not save uploaded file")
	}

	return gin.H{
		"name": filename,
		"size": file.Size,
	}, nil
}

func safeFilename(name string) (string, bool) {
	filename := filepath.Base(name)
	filename = strings.TrimSpace(filename)
	return filename, filename != "" && filename != "." && filename != string(filepath.Separator)
}

func uniqueFilename(directory string, filename string) string {
	if _, err := os.Stat(filepath.Join(directory, filename)); os.IsNotExist(err) {
		return filename
	}

	extension := filepath.Ext(filename)
	base := strings.TrimSuffix(filename, extension)
	for count := 2; ; count++ {
		candidate := fmt.Sprintf("%s-%d%s", base, count, extension)
		if _, err := os.Stat(filepath.Join(directory, candidate)); os.IsNotExist(err) {
			return candidate
		}
	}
}