hswaw/teleimg/main.go

package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"regexp"

	"code.hackerspace.pl/hscloud/go/mirko"
	"github.com/dgraph-io/ristretto"
	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	"github.com/golang/glog"
	"github.com/ulule/limiter/v3"
	"github.com/ulule/limiter/v3/drivers/store/memory"
)

func init() {
	flag.Set("logtostderr", "true")
}

var (
	flagPublicListen  string
	flagTelegramToken string
	reTelegram        = regexp.MustCompile(`/fileid/([a-zA-Z0-9_-]+).([a-z0-9]+)`)
)

type server struct {
	cache   *ristretto.Cache
	limiter *limiter.Limiter
	tel     *tgbotapi.BotAPI
}

func main() {
	flag.StringVar(&flagPublicListen, "public_listen", "127.0.0.1:5000", "Listen address for public HTTP handler")
	flag.StringVar(&flagTelegramToken, "telegram_token", "", "Telegram Bot API Token")
	flag.Parse()

	if flagTelegramToken == "" {
		glog.Exitf("telegram_token must be set")
	}

	cache, err := ristretto.NewCache(&ristretto.Config{
		NumCounters: 1e7,     // number of keys to track frequency of (10M).
		MaxCost:     1 << 30, // maximum cost of cache (1GB).
		BufferItems: 64,      // number of keys per Get buffer.
	})
	if err != nil {
		glog.Exit(err)
	}

	tel, err := tgbotapi.NewBotAPI(flagTelegramToken)
	if err != nil {
		glog.Exitf("Error when creating telegram bot: %v", err)
	}

	rate, err := limiter.NewRateFromFormatted("10-M")
	if err != nil {
		glog.Exit(err)
	}

	store := memory.NewStore()
	instance := limiter.New(store, rate, limiter.WithTrustForwardHeader(true))

	s := &server{
		cache:   cache,
		limiter: instance,
		tel:     tel,
	}

	m := mirko.New()
	if err := m.Listen(); err != nil {
		glog.Exitf("Listen(): %v", err)
	}

	if err := m.Serve(); err != nil {
		glog.Exitf("Serve(): %v", err)
	}

	publicMux := http.NewServeMux()
	publicMux.HandleFunc("/", s.publicHandler)
	publicSrv := http.Server{
		Addr:    flagPublicListen,
		Handler: publicMux,
	}
	go func() {
		if err := publicSrv.ListenAndServe(); err != nil {
			glog.Exitf("public ListenAndServe: %v", err)
		}
	}()

	<-m.Done()
}

func setMime(w http.ResponseWriter, ext string) {
	switch ext {
	case "jpg":
		w.Header().Set("Content-Type", "image/jpeg")
	case "mp4":
		w.Header().Set("Content-Type", "video/mp4")
	}
}

func (s *server) publicHandler(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	if !reTelegram.MatchString(r.URL.Path) {
		http.NotFound(w, r)
		return
	}
	parts := reTelegram.FindStringSubmatch(r.URL.Path)
	fileid := parts[1]
	fileext := parts[2]
	glog.Infof("FileID: %s", fileid)

	c, ok := s.cache.Get(fileid)
	if ok {
		glog.Infof("Get %q - cache hit", fileid)
		// cache hit
		setMime(w, fileext)
		w.Write(c.([]byte))
		return
	}

	glog.Infof("Get %q - cache miss", fileid)

	limit, err := s.limiter.Get(ctx, s.limiter.GetIPKey(r))
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, ":(")
		glog.Errorf("limiter.Get(%q): %v", s.limiter.GetIPKey(r), err)
		return
	}

	if limit.Reached {
		w.WriteHeader(420)
		fmt.Fprintf(w, "enhance your calm")
		glog.Warningf("Limit reached by %q", s.limiter.GetIPKey(r))
		return
	}

	f, err := s.tel.GetFile(tgbotapi.FileConfig{fileid})
	if err != nil {
		w.WriteHeader(502)
		fmt.Fprintf(w, "telegram mumbles.")
		glog.Errorf("tel.GetFile(%q): %v", fileid, err)
		return
	}

	target := f.Link(flagTelegramToken)

	req, err := http.NewRequest("GET", target, nil)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, ":(")
		glog.Errorf("NewRequest(GET, %q, nil): %v", target, err)
		return
	}

	req = req.WithContext(ctx)
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, ":(")
		glog.Errorf("GET(%q): %v", target, err)
		return
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		// do not cache errors
		w.WriteHeader(res.StatusCode)
		io.Copy(w, res.Body)
		return
	}

	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		w.WriteHeader(500)
		fmt.Fprintf(w, ":(")
		glog.Errorf("Read(%q): %v", target, err)
		return
	}

	s.cache.Set(fileid, b, int64(len(b)))

	setMime(w, fileext)
	w.Write(b)
}