string -> color

lifecolors.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/mitchellh/hashstructure"
	"github.com/patrickmn/go-cache"

	"encoding/base64"
	"encoding/json"

	"github.com/EdlinOrg/prominentcolor"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/gregjones/httpcache"
	"github.com/gregjones/httpcache/diskcache"
)

type imageUploadType string

const (
	imgURL    imageUploadType = "url"
	imgBase64                 = "base64"
	imgUpload                 = "file-upload"
)

// make these configurable
const (
	defaultMaxRequestBodySize int64  = 1 // in mb
	defaultMaxProminentColors int    = 5
	defaultPort               string = ":8080"
)

type errorType string

const (
	errNone              errorType = "no_error"
	errSerialization               = "serialization_error"
	errSizeTooLarge                = "size_too_large_error"
	errUnknownDataFormat           = "unknown_data_format_error"
	errOthers                      = "other_error"
)

// App holds things that are required for the app
type App struct {
	MaxBodySizeInBytes int64
	MaxProminentColors int
}

var c = cache.New(5*time.Minute, 10*time.Minute)

// HandleImgRequestBody is the request body when type is imgURL value is url when type is imgUpload value is base64 encoded image
type HandleImgRequestBody struct {
	Type            imageUploadType `json:"type"`
	Value           string          `json:"value"`
	ProminentColors int             `json:"num_prominent_colors"`
}

// HandleImgResponseBody is the reponse body of the hanldeImg handler
type HandleImgResponseBody struct {
	ProminentColors []string  `json:"prominent_colors"`
	Error           *AppError `json:"error"`
}

// AppError is error that has fields to contain about application errors
type AppError struct {
	Error     string    `json:"msg"`
	ErrorType errorType `json:"type"`
}

func createAppError(err error, typ errorType) *AppError {
	return &AppError{
		Error:     err.Error(),
		ErrorType: typ,
	}
}

// Hash gives a hash to h
func (h *HandleImgRequestBody) Hash(opts *hashstructure.HashOptions) (uint64, error) {
	return hashstructure.Hash(h, opts)
}

func (a *App) handleImg(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()
	var response HandleImgResponseBody
	// handle response forming here
	defer func() {
		responseBytes, err := json.Marshal(response)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "error: %s", err)
			return
		}

		if response.Error != nil {
			w.WriteHeader(http.StatusInternalServerError)
		} else {
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusOK)
		}

		w.Write(responseBytes)
	}()

	var buf bytes.Buffer
	handleImgRequestBody := new(HandleImgRequestBody)

	lr := io.LimitReader(io.TeeReader(r.Body, &buf), a.MaxBodySizeInBytes)
	dec := json.NewDecoder(lr)
	err := dec.Decode(handleImgRequestBody)
	if err != nil {
		_, _ = io.Copy(ioutil.Discard, r.Body)
		go func() {
			log.Printf("could not decode request body, failed with '%s'\n", err.Error())
			log.Println(buf.String()) // can be big
		}()
		// decode failed
		response.Error = createAppError(fmt.Errorf("could not decode request body, failed with '%s'", err.Error()), errSerialization)
		return
	}

	// look in cache
	if handleImgRequestBody.ProminentColors == 0 || handleImgRequestBody.ProminentColors > a.MaxProminentColors {
		handleImgRequestBody.ProminentColors = a.MaxProminentColors
	}

	var inCache bool

	defer func() {
		if inCache {
			return
		}
		uid, err := handleImgRequestBody.Hash(nil)
		if err == nil && response.Error == nil && len(response.ProminentColors) > 0 {
			c.Set(strconv.Itoa(int(uid)), response.ProminentColors, cache.DefaultExpiration)
		} else {
			log.Println(err)
		}
	}()

	h, err := handleImgRequestBody.Hash(nil)
	if err == nil {
		var cachedProminentColors []prominentcolor.ColorItem
		cachedData, inCache := c.Get(strconv.Itoa(int(h)))
		cachedProminentColors, ok := cachedData.([]prominentcolor.ColorItem)
		if inCache && ok && len(cachedProminentColors) > 0 {
			response.ProminentColors = topColors(cachedProminentColors)
			return
		}
	}

	colors, errTyp, err := findProminentColors(handleImgRequestBody, a.MaxBodySizeInBytes)
	if err != nil {
		response.Error = createAppError(err, errTyp)
	} else {
		response.ProminentColors = colors
	}
	return
}

func main() {
	// receive flags
	portPtr := flag.String("port", defaultPort, "port number for the server to run on")
	maxReqBodySizePtr := flag.Int64("max-req-size", defaultMaxRequestBodySize, "maximum request body size in mb")
	maxProminentColorsPtr := flag.Int("max-prominent-colors", defaultMaxProminentColors, "maximum promiment colors that can be used to limit the user's choice")

	flag.Parse()

	app := &App{
		MaxBodySizeInBytes: *maxReqBodySizePtr << 20, // in bytes
		MaxProminentColors: *maxProminentColorsPtr,
	}

	r := mux.NewRouter()
	r.HandleFunc("/", app.handleImg).Methods(http.MethodPost).Headers("Content-Type", "application/json") // add rate limiting

	srv := &http.Server{
		Handler: handlers.CORS()(r),
		Addr:    *portPtr,
		// Good practice: enforce timeouts for servers you create!
		WriteTimeout: 2 * time.Second,
		ReadTimeout:  2 * time.Second,
	}

	log.Fatal(srv.ListenAndServe())
}

func findProminentColors(handleImgRequestBody *HandleImgRequestBody, maxRequestBodySize int64) ([]string, errorType, error) {
	var imgReader io.ReadCloser
	var err error
	switch handleImgRequestBody.Type {
	case imgURL:
		imgReader, err = getReaderFromURL(handleImgRequestBody.Value)
		if err != nil {
			return nil, errOthers, err
		}
	case imgBase64:
		imgReader, err = getReaderFromBase64Data(handleImgRequestBody.Value)
		if err != nil {
			return nil, errOthers, err
		}
	default:
		return nil, errOthers, fmt.Errorf("requested type %s is not implemented", handleImgRequestBody.Type)
	}
	defer imgReader.Close()
	bufr := new(bytes.Buffer)
	const headerBufLen int64 = 512
	_, err = io.CopyN(bufr, imgReader, headerBufLen)
	if err != nil {
		return nil, errOthers, err
	}
	detectContentType := http.DetectContentType(bufr.Bytes())
	if !strings.HasPrefix(detectContentType, "image/") {
		return nil, errUnknownDataFormat, fmt.Errorf("%s may not be image", detectContentType)
	}
	_, err = io.CopyN(bufr, imgReader, maxRequestBodySize-headerBufLen)
	if err != nil && err != io.EOF {
		return nil, errSizeTooLarge, fmt.Errorf("%s %dmb is the limit of the acceptable image size", err, maxRequestBodySize>>20)
	}

	_, err = io.CopyN(ioutil.Discard, imgReader, 1)
	if err == nil {
		return nil, errSizeTooLarge, fmt.Errorf("%dmb is the limit of the acceptable image size", maxRequestBodySize>>20)
	}

	prominentColors, err := getProminentColorsFromReader(bufr, handleImgRequestBody.ProminentColors)
	if err != nil {
		return nil, errOthers, err
	}

	return topColors(prominentColors), errNone, nil
}

func topColors(prominentColors []prominentcolor.ColorItem) (colorsInHex []string) {
	for _, prominentColor := range prominentColors {
		colorsInHex = append(colorsInHex, prominentColor.AsString())
	}
	return colorsInHex
}

func getProminentColorsFromReader(r io.Reader, n int) (rgbColors []prominentcolor.ColorItem, err error) {
	img, _, err := image.Decode(r)

	if err != nil {
		return nil, err
	}

	centroids, err := prominentcolor.KmeansWithAll(n, img, prominentcolor.ArgumentNoCropping, prominentcolor.DefaultSize, prominentcolor.GetDefaultMasks())
	if err != nil {
		return nil, err
	}

	return centroids, nil
}

func getReaderFromFile(path string) (io.ReadCloser, error) {
	f, err := os.Open(path)

	return f, err
}

func getReaderFromBase64Data(data string) (r io.ReadCloser, err error) {
	r = ioutil.NopCloser(base64.NewDecoder(base64.StdEncoding, strings.NewReader(data)))
	return
}

func getReaderFromURL(url string) (r io.ReadCloser, err error) {
	cli := http.Client{
		Transport: httpcache.NewTransport(diskcache.New("./cache")),
	}

	resp, err := cli.Get(url)
	if err != nil {
		return nil, err
	}

	contentType := resp.Header.Get("Content-Type")
	if contentType != "" && !strings.HasPrefix(contentType, "image/") {
		err = fmt.Errorf("%s may not be an image", contentType)
	}

	return resp.Body, err
}