goredirect/main.go

package main

import (
	"context"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
)

type inventoryItemType struct {
	ID        primitive.ObjectID `bson:"_id" json:"_id,omitempty"`
	Shortener struct {
		Slug string `bson:"slug" json:"slug"`
		URL  string `bson:"url" json:"url"`
	} `bson:"shortener" json:"shortener"`
}

var (
	mongoURI        = "mongodb://127.0.0.1:27017/default?replicaSet=rs0"
	mongoCollection = "inventory"
)

var (
	regexValid       = regexp.MustCompile("^[a-zA-Z0-9]{4,6}$")
	redirectFound    = os.Getenv("GOREDIRECT_FOUND")
	redirectNotFound = os.Getenv("GOREDIRECT_NOT_FOUND")
	redirectNoPath   = os.Getenv("GOREDIRECT_NOPATH")
)

func lookupHandler(ctx context.Context, coll *mongo.Collection) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		counterQueries.Inc()
		slug := r.URL.Path[1:]

		// Prevalidation

		if slug == "" && redirectNoPath != "" {
			counterNoPath.Inc()

			log.Printf("Redirecting empty slug to %s\n", redirectNoPath)
			http.Redirect(w, r, redirectNoPath, 302)

			return
		}

		if match := regexValid.MatchString(slug); !match {
			counterInvalidSlug.Inc()

			http.Error(w, "Invalid slug", 400)
			return
		}

		// Query

		var doc inventoryItemType
		if err := coll.FindOne(ctx, bson.M{"shortener.slug": slug}).Decode(&doc); err != nil {
			counterNotFound.Inc()

			if redirectNotFound == "" {
				http.NotFound(w, r)
				return
			}

			http.Redirect(w, r, strings.Replace(redirectNotFound, "%s", slug, 1), 302)
			return
		}

		counterFound.Inc()

		redirURL := doc.Shortener.URL
		if redirURL == "" {
			redirURL = strings.Replace(redirectFound, "%s", doc.ID.Hex(), 1)
		}

		http.Redirect(w, r, redirURL, 302)
		log.Printf("Redirecting %s to %s\n", slug, redirURL)
	}
}

var (
	counterQueries = promauto.NewCounter(prometheus.CounterOpts{
		Name: "goredirect_queries",
		Help: "The total number of queries.",
	})
	counterNoPath = promauto.NewCounter(prometheus.CounterOpts{
		Name: "goredirect_not_found",
		Help: "The total number of queries with unknown slug.",
	})
	counterInvalidSlug = promauto.NewCounter(prometheus.CounterOpts{
		Name: "goredirect_invalid_slug",
		Help: "The total number of queries that did not match regex.",
	})
	counterNotFound = promauto.NewCounter(prometheus.CounterOpts{
		Name: "goredirect_not_found",
		Help: "The total number of queries with unknown slug.",
	})
	counterFound = promauto.NewCounter(prometheus.CounterOpts{
		Name: "goredirect_found",
		Help: "The total number of queries that resulted in successful redirect.",
	})
)

func main() {
	ctx := context.Background()

	if val := os.Getenv("GOREDIRECT_REGEX"); val != "" {
		regexValid = regexp.MustCompile(val)
	}

	if val := os.Getenv("MONGO_URI"); val != "" {
		mongoURI = val
	}

	if val := os.Getenv("GOREDIRECT_COLLECTION"); val != "" {
		mongoCollection = val
	}

	// Mongo database //

	cs, err := connstring.ParseAndValidate(mongoURI)
	if err != nil {
		log.Fatal(err)

		os.Exit(1)
	}

	client, err := mongo.NewClient(options.Client().ApplyURI(mongoURI))
	if err != nil {
		log.Fatal(err)

		os.Exit(1)
	}

	connectCtx, cancel := context.WithTimeout(ctx, 10*time.Second)

	if err := client.Connect(connectCtx); err != nil {
		log.Fatal(err)

		os.Exit(1)
	}
	cancel()

	coll := client.Database(cs.Database).Collection(mongoCollection)
	defer client.Disconnect(ctx)

	// HTTP Server //

	http.Handle("/metrics", promhttp.Handler())
	http.HandleFunc("/", lookupHandler(ctx, coll))

	log.Printf("Starting HTTP server on :8080\n")

	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatal("ListenAndServe: ", err)

		os.Exit(1)
	}
}