imgcdn/main.go

package main

import (
	"image"
	"image/jpeg"
	"image/png"
	"net/http"
	"strconv"
	"strings"
	"log"
	"os"
	"bufio"
	"flag"
	"bytes"

	"github.com/nfnt/resize"
	"golang.org/x/image/webp"
	"github.com/go-redis/redis"
)

var (
	config		= flag.String("config", "", "a file to add your authorized hosts")
	useredis	= flag.Bool("redis", true, "you can disable redis with this variable")
	redishost	= flag.String("redishost", "localhost", "redis host")
	redisport	= flag.String("redisport", "6379", "redis host")
	redisClient *redis.Client
)

func main() {
	var (
		hosts []string
		err error
	)

	flag.Parse()

	if *config != "" {
		hosts, err = parseHosts(*config)
		if err != nil {
			log.Fatalf("Error reading config file: %s\n", err.Error())
			return
		}
	}

	redisClient = redis.NewClient(&redis.Options{
		Addr:		*redishost + ":" + *redisport,
		Password:	"",
		DB:			0,
	})

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}

		referrer := r.Referer()
		if referrer == "" {
			log.Println("Referrer not found")
			http.Error(w, "Referrer not found", http.StatusBadRequest)
			return
		}

		if len(hosts) > 0 && !hostMatch(referrer, hosts) {
			log.Println("Referrer unauthorized")
			http.Error(w, "Referrer unauthorized", http.StatusForbidden)
			return
		}

		params := r.URL.Query()
		width, err := strconv.Atoi(params.Get("w"))
		if err != nil {
			width = 0
		}

		height, err := strconv.Atoi(params.Get("h"))
		if err != nil {
			height = 0
		}

		percent, err := strconv.Atoi(params.Get("p"))
		if err != nil {
			percent = 0
		}

		path := r.URL.Path[1:]
		url := referrer + "/" + path
		ext := getExt(path)

		var (
			imgKey string
			imgData []byte
		)

		// Check if image is in Redis cache
		if *useredis == true {
			imgKey = "image:" + url
			imgData, err = redisClient.Get(imgKey).Bytes()
			if err == nil {
				// Image found in cache, decode and resize
				var img image.Image
				switch ext {
				case "png":
					img, err = png.Decode(bytes.NewReader(imgData))
				case "webp":
					img, err = webp.Decode(bytes.NewReader(imgData))
				default:
					img, _, err = image.Decode(bytes.NewReader(imgData))
				}

				if err != nil {
					log.Printf("Error decoding image: %s\n", err.Error())
					http.Error(w, "Error decoding image", http.StatusInternalServerError)
					return
				}

				var resizedImg image.Image
				if percent > 0 {
					resizedImg = resize.Resize(uint(img.Bounds().Dx()*percent/100), uint(img.Bounds().Dy()*percent/100), img, resize.Lanczos3)
				} else if width > 0 || height > 0 {
					resizedImg = resize.Resize(uint(width), uint(height), img, resize.Lanczos3)
				} else {
					resizedImg = resize.Resize(0, 500, img, resize.Lanczos3)
				}

				switch ext {
				case "png":
					w.Header().Set("Content-Type", "image/png")
					png.Encode(w, resizedImg)
				default:
					w.Header().Set("Content-Type", "image/jpeg")
					jpeg.Encode(w, resizedImg, nil)
				}

				return
			}
		}

		// Image not found in cache, fetch from URL
		resp, err := http.Get(url)
		if err != nil {
			log.Printf("Error getting image: %s\n", err.Error())
			http.Redirect(w, r, url, http.StatusSeeOther)
			return
		}
		defer resp.Body.Close()

		var img image.Image
		switch ext {
		case "png":
			img, err = png.Decode(resp.Body)
		case "webp":
			img, err = webp.Decode(resp.Body)
		default:
			img, _, err = image.Decode(resp.Body)
		}

		if err != nil {
			log.Printf("Error decoding image: %s\n", err.Error())
			http.Error(w, "Error decoding image", http.StatusInternalServerError)
			return
		}

		// Store image in Redis cache
		if *useredis == true {
			buf := new(bytes.Buffer)
			switch ext {
			case "png":
				png.Encode(buf, img)
			default:
				jpeg.Encode(buf, img, nil)
			}
			imgData = buf.Bytes()
			go redisClient.Set(imgKey, imgData, 0)
		}

		var resizedImg image.Image
		if percent > 0 {
			resizedImg = resize.Resize(uint(img.Bounds().Dx()*percent/100), uint(img.Bounds().Dy()*percent/100), img, resize.Lanczos3)
		} else if width > 0 || height > 0 {
			resizedImg = resize.Resize(uint(width), uint(height), img, resize.Lanczos3)
		} else {
			resizedImg = resize.Resize(0, 500, img, resize.Lanczos3)
		}

		switch ext {
		case "png":
			w.Header().Set("Content-Type", "image/png")
			png.Encode(w, resizedImg)
		default:
			w.Header().Set("Content-Type", "image/jpeg")
			jpeg.Encode(w, resizedImg, nil)
		}
		return
	})

	log.Println("Server started on port 8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func getExt(s string) string {
	parts := strings.Split(s, ".")
	return parts[len(parts)-1]
}

func parseHosts(filePath string) ([]string, error) {
	var hosts []string

	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		hosts = append(hosts, scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return hosts, nil
}

func hostMatch(s string, list []string) bool {
	for _, host := range list {
		if strings.Contains(s, host) {
			return true
		}
	}
	return false
}