version 0.7.0: add thumbnail support

README.md

@@ -3,6 +3,12 @@
 Simple but quick Golang webserver that serves requests to get files and
 redirects from [PostgreSQL](https://www.postgresql.org).
 
+### Features
+- Serves files, short URLs and "tombstones" (deleted file markers)
+- Allows for URL previewing on short URLs (add `?preview`)
+- Allows for thumbnail generation on images (if enabled, add `?thumbnail`)
+- Can be configured to store generalized metrics
+
 ### Requirements
 
 - PostgreSQL server with `objects` table
@@ -40,11 +46,10 @@ means using the `.json` files in [lib/metrics/](lib/metrics).
 
 ### TODO
 
-- [ ] Write tests
-- [ ] Add thumbnail functionality
+- `OPTIONS`/`HEAD` support
+- Write tests
 
 ### License
 
 `cdn-origin` is licensed under the MIT license. A copy of the MIT license can be
 found in [LICENSE](LICENSE).
-

config.sample.toml

@@ -56,3 +56,13 @@
 [files]
     # Storage location of the bucket on disk
     storageLocation = "/var/data/buckets/public"
+
+[thumbnails]
+    # Enable thumbnails? (add ?thumbnail to the end of a file object URL)
+    enable = true
+
+    # Enable thumbnail cache?
+    cacheEnable = true
+
+    # Thumbnail cache location (if enabled).
+    cacheLocation = "/tmp/thumbs"

go.mod

 module owo.codes/whats-this/cdn-origin
 
 require (
+	github.com/discordapp/lilliput v0.0.0-20190215194814-3d66dff0ee01
 	github.com/lib/pq v1.0.0
 	github.com/oschwald/maxminddb-golang v1.3.0
+	github.com/pkg/errors v0.8.0
 	github.com/rs/zerolog v1.11.0
 	github.com/spf13/pflag v1.0.3
 	github.com/spf13/viper v1.3.1

go.sum

@@ -3,6 +3,8 @@ github.com/coreos/etcd v3.3.10+incompatible/go.mod h1:uF7uidLiAD3TWHmW31ZFd/JWoc
 github.com/coreos/go-etcd v2.0.0+incompatible/go.mod h1:Jez6KQU2B/sWsbdaef3ED8NzMklzPG4d5KIOhIy30Tk=
 github.com/coreos/go-semver v0.2.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/discordapp/lilliput v0.0.0-20190215194814-3d66dff0ee01 h1:8iBDqApWovRWN5H5XG9RbaU2Jk0SJ9oPWS8yRwJqkrU=
+github.com/discordapp/lilliput v0.0.0-20190215194814-3d66dff0ee01/go.mod h1:Y/qDoFO8ipcjzUjR+FB3NrX2oJcp6EXQBLuW0nws2IY=
 github.com/fortytw2/leaktest v1.2.0/go.mod h1:jDsjWgpAGjm2CA7WthBh/CdZYEPF31XHquHwclZch5g=
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=

lib/db/models.go

@@ -11,4 +11,5 @@ type Object struct {
 	ObjectType   int        `json:"object_type"`
 	DeletedAt    *time.Time `json:"deleted_at"`
 	DeleteReason *string    `json:"delete_reason"`
+	MD5Hash      *string    `json:"md5_hash"`
 }

lib/db/queries.go

@@ -16,8 +16,9 @@ func SelectObjectByBucketKey(bucket, key string) (Object, error) {
 	var objectType int
 	var deletedAt pq.NullTime
 	var deleteReason sql.NullString
+	var md5Hash sql.NullString
 	err := DB.QueryRow(selectObjectByBucketKey, fmt.Sprintf("%s/%s", bucket, key)).
-		Scan(&contentType, &destURL, &objectType, &deletedAt, &deleteReason)
+		Scan(&contentType, &destURL, &objectType, &deletedAt, &deleteReason, &md5Hash)
 	if err != nil {
 		return object, err
 	}
@@ -27,7 +28,7 @@ func SelectObjectByBucketKey(bucket, key string) (Object, error) {
 		object.ContentType = &contentType.String
 	}
 	if destURL.Valid {
-		object.DestURL= &destURL.String
+		object.DestURL = &destURL.String
 	}
 	if deletedAt.Valid {
 		object.DeletedAt = &deletedAt.Time
@@ -35,6 +36,9 @@ func SelectObjectByBucketKey(bucket, key string) (Object, error) {
 			object.DeleteReason = &deleteReason.String
 		}
 	}
+	if md5Hash.Valid {
+		object.MD5Hash = &md5Hash.String
+	}
 	object.ObjectType = objectType
 	return object, nil
 }

lib/db/sql.go

@@ -6,7 +6,8 @@ SELECT
 	dest_url,
 	"type",
 	deleted_at,
-	delete_reason
+	delete_reason,
+	md5_hash
 FROM
 	objects
 WHERE

lib/thumbnailer/cache.go

+package thumbnailer
+
+import (
+	"io"
+	"os"
+	"path/filepath"
+)
+
+// ThumbnailCache allows access to thumbnails stored in a directory. Each
+// thumbnail has a key, which uniquely identifies it. The key should be a unique
+// ID from a database or the original file's hash.
+type ThumbnailCache struct {
+	Directory string
+}
+
+// NewThumbnailCache creates a new *ThumbnailCache.
+func NewThumbnailCache(directory string) *ThumbnailCache {
+	return &ThumbnailCache{
+		Directory: directory,
+	}
+}
+
+// GetThumbnail returns a thumbnail that is cached. If no cached copy exists, a
+// exists, a NoCachedCopy error is returned.
+func (c *ThumbnailCache) GetThumbnail(key string) (io.ReadCloser, error) {
+	path := filepath.Join(c.Directory, key)
+	data, err := os.Open(path)
+	if os.IsNotExist(err) {
+		return nil, NoCachedCopy
+	}
+	return data, err
+}
+
+// SetThumbnail stores a thumbnail with the specified key.
+func (c *ThumbnailCache) SetThumbnail(key string, data io.Reader) error {
+	path := filepath.Join(c.Directory, key)
+	file, err := os.Create(path)
+	defer file.Close()
+	if err != nil {
+		return err
+	}
+	_, err = io.Copy(file, data)
+	return err
+}
+
+// Transform generates a thumbnail and caches it.
+func (c *ThumbnailCache) Transform(key string, data io.Reader) error {
+	outputImage, err := Transform(data)
+	if err != nil {
+		return err
+	}
+	return c.SetThumbnail(key, outputImage)
+}
+
+// DeleteThumbnail deletes a thumbnail from the cache.
+func (c *ThumbnailCache) DeleteThumbnail(key string) error {
+	path := filepath.Join(c.Directory, key)
+	return os.Remove(path)
+}

lib/thumbnailer/errors.go

+package thumbnailer
+
+type thumbnailerError struct {
+	Err string
+}
+
+// Error implements error.
+func (e *thumbnailerError) Error() string {
+	return e.Err
+}
+
+// NoCachedCopy means there is no cached copy for the specified key available.
+var NoCachedCopy error = &thumbnailerError{"no cached copy of the thumbnail requested is available"}
+
+// InputTooLarge means that the pixel size of the input image is too big to be thumbnailed.
+var InputTooLarge error = &thumbnailerError{"the input size in pixels is too large"}

lib/thumbnailer/transform.go

+package thumbnailer
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"sync"
+
+	"github.com/discordapp/lilliput"
+	"github.com/pkg/errors"
+)
+
+// Transform settings. These should remain constant. If these are ever changed,
+// the thumbnail cache should be cleared.
+var (
+	maxInputResolution  = 10000           // maximum length in pixels of any input length
+	maxOutputResolution = 200             // maximum length in pixels of any output length
+	outputBufferSize    = 5 * 1024 * 1024 // bytes
+	outputType          = ".jpeg"         // with leading dot
+	encodeOptions       = map[int]int{
+		lilliput.JpegQuality: 85,
+	}
+)
+
+// Accepted MIME types for thumbnails in map for easy checking
+var thumbnailMIMETypes = map[string]struct{}{
+	"image/gif":  struct{}{},
+	"image/jpeg": struct{}{},
+	"image/png":  struct{}{},
+	"image/webp": struct{}{},
+}
+
+// AcceptedMIMEType checks if a MIME type is suitable for thumbnailing.
+func AcceptedMIMEType(mime string) bool {
+	_, ok := thumbnailMIMETypes[mime]
+	return ok
+}
+
+// Pool operations for reusing *lillput.ImageOps objects.
+var imageOpsPool = &sync.Pool{
+	New: func() interface{} {
+		return lilliput.NewImageOps(maxInputResolution)
+	},
+}
+
+func getImageOps() *lilliput.ImageOps {
+	return imageOpsPool.Get().(*lilliput.ImageOps)
+}
+func returnImageOps(imageOps *lilliput.ImageOps) {
+	imageOpsPool.Put(imageOps)
+}
+
+// Transform takes image data and resizes it to fit the maxOutputResolution above.
+func Transform(data io.Reader) (*bytes.Buffer, error) {
+	b, err := ioutil.ReadAll(data)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to read in image data")
+	}
+
+	// Create decoder and parse header
+	decoder, err := lilliput.NewDecoder(b)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to decode image data")
+	}
+	header, err := decoder.Header()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to parse image header")
+	}
+
+	// Check if the image is within the input limit
+	if header.Width() > maxInputResolution || header.Height() > maxInputResolution {
+		return nil, InputTooLarge
+	}
+
+	// Determine output resolution
+	width := -1
+	height := -1
+	if header.Width() == header.Height() {
+		width = maxOutputResolution
+		height = maxOutputResolution
+	} else if header.Width() > header.Height() {
+		width = 200
+		scale := header.Width() / maxOutputResolution
+		height = header.Height() / scale
+	} else {
+		height = 200
+		scale := header.Height() / maxOutputResolution
+		width = header.Width() / scale
+	}
+
+	// Prepare to resize image
+	ops := getImageOps()
+	outputImage := make([]byte, outputBufferSize)
+
+	// Resizing options
+	opts := &lilliput.ImageOptions{
+		FileType:             outputType,
+		Width:                width,
+		Height:               height,
+		ResizeMethod:         lilliput.ImageOpsFit,
+		NormalizeOrientation: true,
+		EncodeOptions:        encodeOptions,
+	}
+	if header.Width() <= maxOutputResolution && header.Height() <= maxOutputResolution {
+		opts.Width = header.Width()
+		opts.Height = header.Height()
+		opts.ResizeMethod = lilliput.ImageOpsNoResize
+	}
+
+	// Transform the image
+	outputImage, err = ops.Transform(decoder, opts, outputImage)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to transcode image")
+	}
+	returnImageOps(ops)
+	return bytes.NewBuffer(outputImage), err
+}

main.go

 package main
 
 import (
+	"bytes"
 	"database/sql"
 	"fmt"
 	"html/template"
+	"io"
 	"net"
 	"os"
 	"path/filepath"
@@ -13,23 +15,33 @@ import (
 
 	"owo.codes/whats-this/cdn-origin/lib/db"
 	"owo.codes/whats-this/cdn-origin/lib/metrics"
+	"owo.codes/whats-this/cdn-origin/lib/thumbnailer"
 
+	_ "github.com/lib/pq"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
-	_ "github.com/lib/pq"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 	"github.com/valyala/fasthttp"
 )
 
-// version is the current version of cdn-origin.
 // Build config
 const (
 	configLocationUnix = "/etc/whats-this/cdn-origin/config.toml"
-	shutdownTimeout = 10 *time.Second
-	version = "0.5.0"
+	version            = "0.7.0"
 )
 
+// readCloserBuffer is a *bytes.Buffer that implements io.ReadCloser.
+type readCloserBuffer struct {
+	*bytes.Buffer
+}
+
+func (b *readCloserBuffer) Close() error {
+	return nil
+}
+
+var _ io.ReadCloser = &readCloserBuffer{}
+
 // redirectHTML is the html/template template for generating redirect HTML.
 const redirectHTML = `<html><head><meta charset="UTF-8" /><meta http-equiv=refresh content="0; url={{.}}" /><script type="text/javascript">window.location.href="{{.}}"</script><title>Redirect</title></head><body><p>If you are not redirected automatically, click <a href="{{.}}">here</a> to go to the destination.</p></body></html>`
 
@@ -131,6 +143,9 @@ func init() {
 	if viper.GetString("files.storageLocation") == "" {
 		log.Fatal().Msg("Configuration: files.storageLocation is required")
 	}
+	if viper.GetBool("thumbnails.enable") && viper.GetBool("thumbnails.cacheEnable") && viper.GetString("thumbnails.cacheLocation") == "" {
+		log.Fatal().Msg("thumbnails.cacheLocation is required when thumbnails and thumbnails cache is enabled")
+	}
 
 	// Parse redirect templates
 	redirectHTMLTemplate, err = template.New("redirectHTML").Parse(redirectHTML)
@@ -144,6 +159,7 @@ func init() {
 }
 
 var collector *metrics.Collector
+var thumbnailCache *thumbnailer.ThumbnailCache
 
 func main() {
 	// Connect to PostgreSQL database
@@ -177,6 +193,11 @@ func main() {
 		}
 	}
 
+	// Setup thumbnail cache
+	if viper.GetBool("thumbnails.enable") && viper.GetBool("thumbnails.cacheEnable") {
+		thumbnailCache = thumbnailer.NewThumbnailCache(viper.GetString("thumbnails.cacheLocation"))
+	}
+
 	// Launch server
 	h := requestHandler
 	if viper.GetBool("http.compressResponse") {
@@ -252,6 +273,8 @@ func recordMetrics(ctx *fasthttp.RequestCtx) {
 }
 
 func requestHandler(ctx *fasthttp.RequestCtx) {
+	defer recordMetrics(ctx)
+
 	// Fetch object from database
 	key := string(ctx.Path()[1:])
 	object, err := db.SelectObjectByBucketKey(viper.GetString("database.objectBucket"), key)
@@ -260,14 +283,10 @@ func requestHandler(ctx *fasthttp.RequestCtx) {
 		ctx.SetStatusCode(fasthttp.StatusNotFound)
 		ctx.SetContentType("text/plain; charset=utf8")
 		fmt.Fprintf(ctx, "404 Not Found: %s", ctx.Path())
-		recordMetrics(ctx)
 		return
 	case err != nil:
 		log.Error().Err(err).Msg("failed to run SELECT query on database")
-		ctx.SetStatusCode(fasthttp.StatusInternalServerError)
-		ctx.SetContentType("text/plain; charset=utf8")
-		fmt.Fprint(ctx, "500 Internal Server Error")
-		recordMetrics(ctx)
+		internalServerError(ctx)
 		return
 	}
 
@@ -275,6 +294,98 @@ func requestHandler(ctx *fasthttp.RequestCtx) {
 	case 0: // file
 		ctx.SetUserValue("object_type", "file")
 
+		// Thumbnails
+		if viper.GetBool("thumbnails.enable") && ctx.QueryArgs().Has("thumbnail") {
+			thumbnailKey := *object.MD5Hash
+			if !thumbnailer.AcceptedMIMEType(*object.ContentType) || thumbnailKey == "" {
+				ctx.SetStatusCode(fasthttp.StatusNotFound)
+				ctx.SetContentType("text/plain; charset=utf8")
+				fmt.Fprintf(ctx, "404 Not Found: %s?thumbnail (cannot generate thumbnail)", ctx.Path())
+				return
+			}
+
+			// Get thumbnail
+			var thumb io.ReadCloser
+			if viper.GetBool("thumbnails.cacheEnable") {
+				thumb, err = thumbnailCache.GetThumbnail(thumbnailKey)
+				if thumb != nil {
+					defer thumb.Close()
+				}
+				if err == thumbnailer.NoCachedCopy {
+					fPath := filepath.Join(viper.GetString("files.storageLocation"), key)
+					file, err := os.Open(fPath)
+					if file != nil {
+						defer file.Close()
+					}
+					if err != nil {
+						log.Warn().Err(err).Msg("failed to open original file to generate thumbnail")
+						internalServerError(ctx)
+						return
+					}
+					err = thumbnailCache.Transform(thumbnailKey, file)
+					if err == thumbnailer.InputTooLarge {
+						ctx.SetStatusCode(fasthttp.StatusNotFound)
+						ctx.SetContentType("text/plain; charset=utf8")
+						fmt.Fprintf(ctx, "404 Not Found: %s?thumbnail (cannot generate thumbnail)", ctx.Path())
+						return
+					} else if err != nil {
+						log.Warn().Err(err).Msg("failed to generate new thumbnail")
+						internalServerError(ctx)
+						return
+					}
+					thumb, err = thumbnailCache.GetThumbnail(thumbnailKey)
+					if thumb != nil {
+						defer thumb.Close()
+					}
+					if err != nil {
+						log.Warn().Err(err).Msg("failed to get thumbnail from cache")
+						internalServerError(ctx)
+						return
+					}
+				} else if err != nil {
+					log.Warn().Err(err).Msg("failed to get thumbnail from cache")
+					internalServerError(ctx)
+					return
+				}
+			} else {
+				fPath := filepath.Join(viper.GetString("files.storageLocation"), key)
+				file, err := os.Open(fPath)
+				if file != nil {
+					defer file.Close()
+				}
+				if err != nil {
+					log.Warn().Err(err).Msg("failed to open original file to generate thumbnail")
+					internalServerError(ctx)
+					return
+				}
+				thumbR, err := thumbnailer.Transform(file)
+				if err == thumbnailer.InputTooLarge {
+					ctx.SetStatusCode(fasthttp.StatusNotFound)
+					ctx.SetContentType("text/plain; charset=utf8")
+					fmt.Fprintf(ctx, "404 Not Found: %s?thumbnail (cannot generate thumbnail)", ctx.Path())
+					return
+				} else if err != nil {
+					log.Warn().Err(err).Msg("failed to generate new thumbnail")
+					internalServerError(ctx)
+					return
+				}
+				// Turn the *bytes.Buffer from thumbnailer.Transform into a fake io.ReadCloser.
+				thumb = &readCloserBuffer{thumbR}
+			}
+
+			// Send response
+			ctx.SetStatusCode(fasthttp.StatusOK)
+			ctx.SetContentType("image/jpeg")
+			ctx.Response.Header.Set("Content-Disposition", fmt.Sprintf(`filename="%s.thumbnail.jpeg"`, key))
+			_, err = io.Copy(ctx, thumb)
+			if err != nil {
+				log.Warn().Err(err).Msg("failed to send thumbnail response")
+				ctx.Response.Header.Del("Content-Disposition")
+				internalServerError(ctx)
+			}
+			return
+		}
+
 		// Serve file to client
 		fPath := filepath.Join(viper.GetString("files.storageLocation"), key)
 		ctx.SetStatusCode(fasthttp.StatusOK)
@@ -284,17 +395,13 @@ func requestHandler(ctx *fasthttp.RequestCtx) {
 			ctx.SetContentType("application/octet-stream")
 		}
 		fasthttp.ServeFileUncompressed(ctx, fPath)
-		recordMetrics(ctx)
 
 	case 1: // redirect
 		ctx.SetUserValue("object_type", "redirect")
 
 		if object.DestURL == nil {
 			log.Warn().Str("key", key).Msg("encountered redirect object with NULL dest_url")
-			ctx.SetStatusCode(fasthttp.StatusInternalServerError)
-			ctx.SetContentType("text/plain; charset=utf8")
-			fmt.Fprint(ctx, "500 Internal Server Error")
-			recordMetrics(ctx)
+			internalServerError(ctx)
 			return
 		}
 
@@ -309,11 +416,10 @@ func requestHandler(ctx *fasthttp.RequestCtx) {
 		if err != nil {
 			log.Warn().Err(err).
 				Str("dest_url", *object.DestURL).
-				Bool("preview", ctx.QueryArgs().
-				Has("preview")).Msg("failed to generate HTML redirect page to send to client")
+				Bool("preview", ctx.QueryArgs().Has("preview")).
+				Msg("failed to generate HTML redirect page to send to client")
 			ctx.SetContentType("text/plain; charset=utf8")
 			fmt.Fprintf(ctx, "Failed to generate HTML redirect page, destination URL: %s", *object.DestURL)
-			recordMetrics(ctx)
 			return
 		}
 
@@ -324,7 +430,6 @@ func requestHandler(ctx *fasthttp.RequestCtx) {
 		} else {
 			ctx.SetStatusCode(fasthttp.StatusOK)
 		}
-		recordMetrics(ctx)
 
 	case 2: // tombstone
 		ctx.SetUserValue("object_type", "tombstone")
@@ -337,6 +442,12 @@ func requestHandler(ctx *fasthttp.RequestCtx) {
 			reason = *object.DeleteReason
 		}
 		fmt.Fprintf(ctx, "410 Gone: %s\n\nReason: %s", ctx.Path(), reason)
-		recordMetrics(ctx)
 	}
 }
+
+// internalServerError returns a 500 Internal Server Response.
+func internalServerError(ctx *fasthttp.RequestCtx) {
+	ctx.SetStatusCode(fasthttp.StatusInternalServerError)
+	ctx.SetContentType("text/plain; charset=utf8")
+	fmt.Fprint(ctx, "500 Internal Server Error")
+}