package main

import (
	"context"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"regexp"
	"strconv"
	"time"

	"github.com/fastly/compute-sdk-go/fsthttp"
)

// Secret hmac key value.
const secret = "iqFPeN2uZ0Lm5IrsKaOFKR"

func main() {
	fsthttp.ServeFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
		// Extract the token from the URL or a cookie
		token := r.URL.Query().Get("token")
		if token == "" {
			if cookie, err := r.Cookie("token"); err == nil {
				token = cookie.Value
			}
			if token == "" {
				fmt.Fprintln(os.Stdout, `Missing "token" query param or cookie`)
				w.WriteHeader(fsthttp.StatusForbidden)
				return
			}
		}

		// Extract token expiration and signature.
		matches := regexp.MustCompile(`^(\d+)_([\w-]+)$`).FindStringSubmatch(token)
		if matches == nil {
			fmt.Fprintln(os.Stdout, "Token format is invalid")
			w.WriteHeader(fsthttp.StatusForbidden)
			return
		}

		expiryTime := matches[1]
		signature := matches[2]
		message := fmt.Sprintf("%s%s%s", r.URL.Path, expiryTime, r.Header.Get("User-Agent"))
    
		mac := hmac.New(sha1.New, []byte(secret))
		mac.Write([]byte(message))
		expectedSignature := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))

		// Validate signature.
		if !hmac.Equal([]byte(signature), []byte(expectedSignature)) {
			fmt.Fprintf(os.Stdout, "Token is incorrect, expected %s\n", expectedSignature)
			w.WriteHeader(fsthttp.StatusGone)
			return
		}

		// Check that expiration time has not elapsed.
		i, _ := strconv.ParseInt(expiryTime, 10, 64)
		if time.Now().After(time.Unix(i, 0)) {
			fmt.Fprintln(os.Stdout, "Token has expired")
			w.WriteHeader(fsthttp.StatusGone)
			return
		}

		// Remove the token from the URL before lookup, preventing the cache from fragmenting.
		qs := r.URL.Query()
		qs.Del("token")
		r.URL.RawQuery = qs.Encode()

		resp, err := r.Send(ctx, "origin_0")
		if err != nil {
			fmt.Fprintln(os.Stdout, err.Error())
			w.WriteHeader(fsthttp.StatusBadGateway)
			return
		}

		w.Header().Reset(resp.Header)
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	})
}