package main

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"strings"

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

const BackendName = "origin_0"
const UnauthorizedResp = `
<!DOCTYPE HTML PUBLIC "-//W3C//DTD HTML 4.01 Transitional//EN" "http://www.w3.org/TR/html4/loose.dtd">
<html>
    <head>
        <title>Error</title>
        <meta HTTP-EQUIV='Content-Type' CONTENT='text/html;'>
    </head>
    <body><h1>401 Unauthorized (Fastly)</h1></body>
</html>`

func main() {
	fsthttp.ServeFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
		// Get Authorization header
		authHdr := r.Header.Get("Authorization")
		if authHdr == "" {
			writeUnauthorizedResponse(w)
			return
		}

		// Get just the credential part of the header
		const prefix = "Basic "
		encodedCredential := authHdr[len(prefix):]

		// Open our credential dictionary
		store, err := edgedict.Open("username_password")
		if err != nil {
			w.WriteHeader(fsthttp.StatusInternalServerError)
			return
		}

		// Check edgedict for a key matching our base64 encoded credential
		_, err = store.Get(encodedCredential)
		if err != nil {
			writeUnauthorizedResponse(w)
			return
		}

		// Authorization was successful! Now we set the authorized-user
		// header on our request before sending it along to the backend

		username := getUsername(encodedCredential)

		r.Header.Add("authorized-user", username)
		r.Header.Del("Authorization")
		resp, err := r.Send(ctx, BackendName)

    // Check the backend response and forward it along to the client
		if err != nil {
			w.WriteHeader(fsthttp.StatusBadGateway)
			fmt.Fprintln(w, err.Error())
			return
		}
		w.Header().Reset(resp.Header)
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	})
}

func unauthorizedResponse() *fsthttp.Response {
	h := fsthttp.NewHeader()
	h.Add("Content-Type", "text/html; charset=utf-8")
	return &fsthttp.Response{
		Header:     h,
		StatusCode: fsthttp.StatusUnauthorized,
		Body:       io.NopCloser(strings.NewReader(UnauthorizedResp)),
	}
}

func writeUnauthorizedResponse(w fsthttp.ResponseWriter) {
	resp := unauthorizedResponse()
	w.Header().Reset(resp.Header)
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func getUsername(encodedCredential string) string {
	data, _ := base64.StdEncoding.DecodeString(encodedCredential)
	dataStr := string(data)
	return dataStr[:strings.Index(dataStr, ":")]
}