ShortDragon/internal/web/routes.go

package web

import (
	"crypto/sha256"
	"crypto/sha512"
	"encoding/hex"
	"fmt"
	"io"
	"math"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/jinzhu/gorm"
	"github.com/spf13/viper"

	"git.1750studios.com/ToddShepard/ShortDragon/internal/database"
)

func redirect(c *gin.Context) {
	var URL database.URL
	short := c.Param("short")
	err := database.Db.First(&URL, "short = ?", short).Error
	if err != nil {
		c.Redirect(http.StatusFound, viper.GetString("DefaultURL"))
		return
	}
	if !URL.Long.Valid && !URL.FileName.Valid {
		c.Redirect(http.StatusFound, viper.GetString("DefaultURL"))
		return
	}
	if URL.Hits.Valid && !(c.Request.Header.Get("DNT") == "1") {
		URL.Hits.Int64 = URL.Hits.Int64 + 1
		database.Db.Save(&URL)
	}
	if URL.Long.Valid {
		c.Redirect(http.StatusFound, URL.Long.String)
	} else {
		c.Redirect(http.StatusFound, "/f/"+URL.Short.String)
	}
}

func encode(c *gin.Context) {
	var URL database.URL
	var count uint
	URL.Long.String = c.PostForm("LongURL")
	if URL.Long.String == "" {
		// No URL, maybe a file
		URL.Long.Valid = false
		file, err := c.FormFile("File")
		if err != nil {
			c.Data(http.StatusFailedDependency, "text/plain", []byte(fmt.Sprintf("Not a URL, nor File; Error: %+v", err)))
			return
		}

		fparts := strings.Split(file.Filename, "/")
		destName := path.Join(viper.GetString("FileFolder"), fparts[len(fparts)-1])
		src, err := file.Open()
		if err != nil {
			c.Data(http.StatusFailedDependency, "text/plain", []byte(fmt.Sprintf("Error: %+v", err)))
			return
		}
		defer src.Close()

		var newHash string
		hasher := sha256.New()
		hasher.Write([]byte(fparts[len(fparts)-1]))
		io.Copy(hasher, src)
		newHash = hex.EncodeToString(hasher.Sum(nil))
		database.Db.Where("file_hash = ?", newHash).FirstOrInit(&URL)

		if !URL.FileName.Valid {
			// No file with given hash in DB
			if _, err := os.Stat(destName); err == nil {
				// File does exist, save it with hash
				destName += newHash
			}
			out, err := os.Create(destName)
			if err != nil {
				c.Data(http.StatusFailedDependency, "text/plain", []byte(fmt.Sprintf("Error: %+v", err)))
				return
			}
			defer out.Close()
			src.Seek(0, io.SeekStart)
			io.Copy(out, src)
			destParts := strings.Split(destName, "/")

			URL.FileName.String = destParts[len(destParts)-1]
			URL.FileName.Valid = true
			URL.FileSize.Int64 = file.Size
			URL.FileSize.Valid = true
			URL.FileHash.String = newHash
			URL.FileHash.Valid = true
			URL.Long.Valid = false

			retention := time.Duration(float64(viper.GetInt("MinAge")) + float64(-viper.GetInt("MaxAge")+viper.GetInt("MinAge"))*math.Pow(float64(URL.FileSize.Int64/viper.GetInt64("MaxSize")-1), 3.0)*24)
			URL.FileDelete.Time = time.Now().Add(retention * time.Hour)
			URL.FileDelete.Valid = true
		} else {
			retention := time.Duration(float64(viper.GetInt("MinAge")) + float64(-viper.GetInt("MaxAge")+viper.GetInt("MinAge"))*math.Pow(float64(URL.FileSize.Int64/viper.GetInt64("MaxSize")-1), 3.0)*24)
			URL.FileDelete.Time = time.Now().Add(retention * time.Hour)
			URL.FileDelete.Valid = true
		}

	} else {
		URL.Long.Valid = true
		database.Db.Where("long = ?", URL.Long.String).FirstOrInit(&URL)
	}

	if URL.Short.Valid {
		c.Data(http.StatusOK, "text/plain", []byte(viper.GetString("ShortURL")+"/r/"+URL.Short.String))
		return
	}
	URL.Short.String = c.DefaultPostForm("ShortURL", "")
	if URL.Short.String == "" {
		var hash string
		if !URL.FileHash.Valid {
			hasher := sha512.New()
			hasher.Write([]byte(URL.Long.String))
			hash = hex.EncodeToString(hasher.Sum(nil))
		} else {
			hash = URL.FileHash.String
		}

		i := 2
		for {
			database.Db.Model(&database.URL{}).Where("short = ?", hash[0:i]).Count(&count)
			if count > 0 && i < len(hash) {
				i = i + 1
			} else if count > 0 {
				c.AbortWithStatus(http.StatusConflict)
				return
			} else {
				URL.Short.String = hash[0:i]
				break
			}
		}
	}
	URL.Short.Valid = true
	URL.Hits.Int64 = 0
	if c.DefaultPostForm("Track", "true") == "true" {
		URL.Hits.Valid = true
	} else {
		URL.Hits.Valid = false
	}
	err := database.Db.Create(&URL).Error
	if err != nil {
		c.AbortWithError(http.StatusBadRequest, err)
		return
	}
	c.Data(http.StatusCreated, "text/plain", []byte(viper.GetString("ShortURL")+"/r/"+URL.Short.String))
	return
}

func decode(c *gin.Context) {
	var URL database.URL
	short := c.Param("short")
	err := database.Db.Where("short = ?", short).Find(&URL).Error
	if err != nil && err == gorm.ErrRecordNotFound {
		c.Data(http.StatusNotFound, "text/plain", []byte("Record not found"))
		return
	} else if err != nil {
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}
	if URL.Long.Valid {
		c.Data(http.StatusOK, "text/plain", []byte(URL.Long.String))
	} else {
		c.Data(http.StatusOK, "text/plain", []byte(URL.FileName.String))
	}
}

func info(c *gin.Context) {
	var URL database.URL
	short := c.Param("short")
	err := database.Db.Where("short = ?", short).Find(&URL).Error
	if err != nil && err == gorm.ErrRecordNotFound {
		c.Data(http.StatusNotFound, "text/plain", []byte("Record not found"))
		return
	} else if err != nil {
		c.AbortWithError(http.StatusInternalServerError, err)
		return
	}
	c.JSON(http.StatusOK, URL)
}

func getFile(c *gin.Context) {
	var URL database.URL
	short := c.Param("short")
	err := database.Db.First(&URL, "short = ?", short).Error
	if err != nil {
		c.Redirect(http.StatusFound, viper.GetString("DefaultURL"))
		return
	}
	file := path.Join(viper.GetString("FileFolder"), URL.FileName.String)
	name := strings.TrimSuffix(file, URL.FileHash.String)
	parts := strings.Split(name, "/")

	if viper.GetBool("XAccel") {
		c.Header("Content-Disposition", "attatchment; filename="+parts[len(parts)-1])
		c.Header("X-Accel-Redirect", viper.GetString("XAccelPrefix")+"/"+URL.FileName.String)
	} else {
		c.FileAttachment(file, parts[len(parts)-1])
	}
}