quartzgun/middleware/middleware.go

package middleware

import (
	"context"
	"fmt"
	"hacklab.nilfm.cc/quartzgun/auth"
	"hacklab.nilfm.cc/quartzgun/cookie"
	"hacklab.nilfm.cc/quartzgun/renderer"
	"hacklab.nilfm.cc/quartzgun/util"
	"net/http"
	"strings"
)

type TokenPayload struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
}

func Protected(next http.Handler, method string, userStore auth.UserStore, login string) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		user, err := cookie.GetToken("user", req)
		if err == nil {
			session, err := cookie.GetToken("session", req)
			if err == nil {
				login, err := userStore.ValidateUser(user, session)
				if err == nil && login {
					fmt.Printf("authorized user: %s\n", user)
					req.Method = method
					next.ServeHTTP(w, req)
					return
				} else if err != nil && err.Error() == "Cookie or token expired" {
					auth.Logout(user, userStore, w)
				}
			}
		}
		fmt.Printf("unauthorized...\n")
		req.Method = http.MethodGet
		http.Redirect(w, req, login, http.StatusSeeOther)
	}

	return http.HandlerFunc(handlerFunc)
}

func Bunt(next string, userStore auth.UserStore, denied string) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		user, err := cookie.GetToken("user", req)
		if err == nil {
			err := auth.Logout(
				user,
				userStore,
				w)
			if err == nil {
				req.Method = http.MethodGet
				http.Redirect(w, req, next, http.StatusSeeOther)
			}
		}
		req.Method = http.MethodGet
		http.Redirect(w, req, denied, http.StatusUnauthorized)
	}

	return http.HandlerFunc(handlerFunc)
}

func Authorize(next string, userStore auth.UserStore, denied string, ttl int) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		err := auth.Login(
			req.FormValue("user"),
			req.FormValue("password"),
			userStore,
			w,
			ttl)
		if err == nil {
			req.Method = http.MethodGet
			fmt.Printf("logged in as %s\n", req.FormValue("user"))
			http.Redirect(w, req, next, http.StatusSeeOther)
		} else {
			fmt.Printf("login failed!\n")
			req.Method = http.MethodGet
			http.Redirect(w, req, denied, http.StatusSeeOther)
		}
	}

	return http.HandlerFunc(handlerFunc)
}

func Provision(userStore auth.UserStore, ttl int) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		user, password, ok := req.BasicAuth()
		if ok {
			token, err := userStore.GrantToken(user, password, ttl)
			if err == nil {
				token := TokenPayload{
					AccessToken: token,
					TokenType:   "bearer",
					ExpiresIn:   ttl,
				}
				util.AddContextValue(req, "token", token)
				renderer.JSON("token").ServeHTTP(w, req)
				return
			}
		}
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	return http.HandlerFunc(handlerFunc)
}

func Validate(next http.Handler, userStore auth.UserStore, scopes map[string]string) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		errString := ""
		authHeader := req.Header.Get("Authorization")
		if strings.HasPrefix(authHeader, "Bearer ") {
			authToken := strings.Split(authHeader, "Bearer ")[1]
			validated, err := userStore.ValidateTokenWithScopes(authToken, scopes)
			if validated && err == nil {
				next.ServeHTTP(w, req)
				return
			} else {
				errString = err.Error()
			}
		} else {
			errString = "No authentication data"
		}
		w.Header().Add("Quartzgun-Error", errString)
		w.WriteHeader(http.StatusUnauthorized)
	}

	return http.HandlerFunc(handlerFunc)
}

func Fortify(next http.Handler) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		token, err := cookie.GetToken("csrfToken", req)
		if err == nil {
			*req = *req.WithContext(
				context.WithValue(
					req.Context(),
					"csrfToken",
					token))
		}
		next.ServeHTTP(w, req)
	}

	return http.HandlerFunc(handlerFunc)
}

func Defend(next http.Handler, userStore auth.UserStore, denied string) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		user, err := cookie.GetToken("user", req)
		if err == nil {
			masterToken, err := userStore.GetData(user, "csrfToken")
			if err == nil {
				cookieToken, err := cookie.GetToken("csrfToken", req)
				if err == nil {
					formToken := req.FormValue("csrfToken")
					if formToken == cookieToken && formToken == masterToken.(string) {
						next.ServeHTTP(w, req)
						return
					}
				}
			}
		}
		http.Redirect(w, req, denied, http.StatusUnauthorized)
	}

	return http.HandlerFunc(handlerFunc)
}

func Throttle(next http.Handler, bouncer func(*http.Request) bool) http.Handler {
	handlerFunc := func(w http.ResponseWriter, req *http.Request) {
		allowed := bouncer(req)
		if allowed {
			next.ServeHTTP(w, req)
			return
		}
		w.WriteHeader(http.StatusTooManyRequests)
	}
	return http.HandlerFunc(handlerFunc)
}