package main

import (
	"context"
	"fmt"
	"io"
  "strings"

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

// BackendName is the name of our service backend.
const BackendName = "origin_0"

func main() {
	fsthttp.ServeFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
		// This requires your service to be configured with a backend
		// named "origin_0" and pointing to "https://http-me.glitch.me".

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

    if (resp.StatusCode == 200 && 
    strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") && 
    len(r.Header.Get("referer")) > 0) {
      // Check if the host header does not contain the referer.
      if (!strings.Contains(r.Header.Get("referer"), r.Header.Get("host"))){
        //Generate a synthetic 403 response
          resp := &fsthttp.Response{
          StatusCode: fsthttp.StatusForbidden,
          Body: io.NopCloser(strings.NewReader("Hotlinking this image is not allowed. Naughty!")),
        }
        flush(resp,w)
        return
      }
    }

		flush(resp,w)
	})
}

func flush(resp *fsthttp.Response, w fsthttp.ResponseWriter) {
	w.Header().Reset(resp.Header)
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}