The code that powers https://tmg.pw

tmgapi.go

package main

import (
    "github.com/jinzhu/gorm"
    _ "github.com/lib/pq"
    _ "github.com/go-sql-driver/mysql"
    "github.com/spf13/viper"
    "fmt"
    "github.com/gin-gonic/gin"
    "net/http"
    "io/ioutil"
    "math/rand"
    "os"
    "io"
    "time"
)

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func main() {
  rand.Seed(time.Now().UnixNano())

  //configuration options
  viper.SetDefault("database", map[string]string{"type": "postgres", "connectionString": "user=tmg password=tmg dbname=tmg host=localhost sslmode=disable"})
  viper.SetDefault("port", 8080)
  viper.SetDefault("host", "0.0.0.0")
  viper.SetDefault("siteURL", "https://tmg.pw")
  viper.SetDefault("storageDir", "uploads")
  viper.SetDefault("defaultMessage", "Hey, this is tmg.pw, my personal screenshot server. It's written using golang and is open source on my github.");

  //config.yml, config.json
  viper.SetConfigName("config")
  err := viper.ReadInConfig()
  if err != nil {
    panic(fmt.Errorf("Could not read config through viper %s \n", err))
  }

  //two variables used throughout the routes
  siteURL := viper.GetString("siteURL");
  storageDir := viper.GetString("storageDir")
  defaultMessage := viper.GetString("defaultMessage")

  databaseConfig := viper.GetStringMapString("database")

  //make directories
  err = os.MkdirAll(storageDir, 0777)
  if err != nil {
    panic(fmt.Errorf("Could not create storage directory! %s \n", err));
  }

  //open a database connection
  db, err := gorm.Open(databaseConfig["type"], databaseConfig["connectionString"])
  if err != nil {
    panic(fmt.Errorf("Could not connect to database %s \n", err))
  }

  //beautiful auto migrations
  db.AutoMigrate(&APIKey{}, &Upload{})
  fmt.Printf("Migrated tables!\n")

  if len(os.Args) > 1 && os.Args[1] == "genkey" {
    apiKey := APIKey{Key: randSeq(64), Description: "Newly created key"}
    db.Create(&apiKey)
    fmt.Printf("Created a new API key\n %s\t%s", apiKey.Key, apiKey.Description)
    os.Exit(0)
    return
  }

  //setup the http server
  router := gin.Default()
  //upload route allows you to upload files to the server
  router.POST("/up", func(c *gin.Context) {
    //get the API Key the user is trying to use, if it is not in the DB then report an error
    var apiKey APIKey;
    providedKey := c.PostForm("key")
    if providedKey == "" {
      c.String(http.StatusUnauthorized, "You did not supply any API Key!")
      return
    }
    db.Where(&APIKey{Key: providedKey}).First(&apiKey)
    if apiKey.ID == 0 {
      c.String(http.StatusUnauthorized, "You did not supply a valid API Key")
      return
    }

    //next, grab the form file
    file, header, err := c.Request.FormFile("file")
    if err != nil {
      c.String(http.StatusBadRequest, "There was an error encountered " + err.Error() + "!")
      return
    }

    //create the instance of the model, and push it to the DB
    headers := header.Header
    upload := Upload{Filename: randSeq(36), URL: randSeq(12), Type: headers.Get("Content-Type"), OriginalFilename : header.Filename, ApiKey: apiKey}
    if upload.Type == "" {
      c.String(http.StatusBadRequest, "The server did not find a Content-Type header on your file!");
      return
    }

    //perform saving
    defer file.Close()
    out, err := os.Create(storageDir + "/" + upload.Filename)
    if err != nil {
      c.String(http.StatusInternalServerError, "The server could not save this file!")
      return
    }
    defer out.Close()
    _, err = io.Copy(out, file)
    if err != nil {
      c.String(http.StatusInternalServerError, "The server could not save this file!")
      return
    }

    db.Create(&upload)
    c.Redirect(http.StatusFound, siteURL + "/" + upload.URL)
  })

  router.GET("/:url", func(c *gin.Context) {
    upload := &Upload{}
    db.Where("url = ?", c.Param("url")).First(upload)
    if upload.ID == 0 {
      c.String(http.StatusNotFound, "Could not find this upload!")
      return
    }

    file, err := ioutil.ReadFile(storageDir + "/" + upload.Filename)

    if err != nil {
      c.String(http.StatusGone, "Could not find the file backing this upload. Strange...")
      return
    }

    var filename string
    if upload.OriginalFilename == "" {
      filename = upload.URL
    } else {
      filename = upload.OriginalFilename
    }

    outputHeaders := c.Writer.Header()
    outputHeaders.Add("Content-Type", upload.Type)
    outputHeaders.Add("Content-Disposition", fmt.Sprintf("inline;filename=\"%s\"", filename))
    c.Writer.WriteHeader(http.StatusOK)
    c.Writer.Write(file)
  })

  router.GET("/", func (c *gin.Context) {
    c.String(http.StatusOK, defaultMessage)
  })

  router.Run(fmt.Sprintf("%s:%d", viper.GetString("host"), viper.GetInt("port")))
}

func randSeq(n int) string {
    b := make([]rune, n)
    for i := range b {
        b[i] = letters[rand.Intn(len(letters))]
    }
    return string(b)
}

type APIKey struct {
  gorm.Model
  Key string
  Description string
}

type Upload struct {
  gorm.Model
  Filename string
  URL string
  Type string
  OriginalFilename string
  ApiKey APIKey
  ApiKeyId int
}