Permalink
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
119 lines (107 sloc) 3.34 KB
package main
import (
"fmt"
"log"
"net/http"
"os"
"strings"
)
// LogResp logs all the responses to clients
// It does this by replacing the default ResponseWriter with a custom one
func LogResp(fn http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
ip := getIPAddress(r)
log.Printf("%s - Received req for %s\n", ip, r.URL.Path)
logger := responseLogger{w: w, ip: ip}
fn(logger, r)
}
}
// MustParams checks for the existence of all specified params in the request
// If they don't exist, the request is declined as BadRequest (400)
func MustParams(fn http.HandlerFunc, params ...string) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
r.ParseForm()
for _, param := range params {
if len(r.Form.Get(param)) == 0 {
errortext := fmt.Sprintf("Expected parameter %s not found", param)
http.Error(w, errortext, http.StatusBadRequest)
return
}
}
fn(w, r)
}
}
// SimpleAuth checks if a password has been set in the ENV and
// compares it to the request's password
// If they don't match, the request is declined as Unauthorized (401)
func SimpleAuth(fn http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
linktopassword := os.Getenv(passwordEnv)
if len(linktopassword) > 0 {
password := r.Form.Get("password")
if linktopassword != password {
errortext := "Password param incorrect"
http.Error(w, errortext, http.StatusUnauthorized)
return
}
}
fn(w, r)
}
}
// RateLimit checks if the client making the request has been making
// more than numreq requests in the specified duration
// If it has, the request is declined as TooManyRequests (429)
func RateLimit(fn http.HandlerFunc, duration, numreq int, store ExpireStore) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
ip := getIPAddress(r)
requests, err := store.Get(ip)
if err != nil || requests < numreq {
store.Incr(ip, duration)
fn(w, r)
return
}
ttl, _ := store.TTL(ip)
errortext := fmt.Sprintf("You've sent too many requests in a short span of time. Try again after %d seconds", ttl)
http.Error(w, errortext, 429)
}
}
// CORS enables Cross-origin resource sharing for this request.
func CORS(fn http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Access-Control-Allow-Origin", "*")
w.Header().Set("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept")
fn(w, r)
}
}
func getIPAddress(r *http.Request) string {
ip := r.Header.Get("X-Forwarded-For") // assuming its behind nginx
if len(ip) > 0 {
return ip
}
return strings.Split(r.RemoteAddr, ":")[0]
}
type responseLogger struct {
w http.ResponseWriter
ip string
}
func (r responseLogger) Header() http.Header {
return r.w.Header()
}
// Write logs the first line of the response and writes the full response to the ResponseWriter
func (r responseLogger) Write(b []byte) (int, error) {
logtext := string(b)
endl := strings.IndexRune(logtext, '\n')
if endl > 0 {
if len(logtext) > endl+2 {
logtext = logtext[0:endl] + " ... (truncated)\n"
}
} else {
logtext = logtext + "\n"
}
log.Printf(fmt.Sprintf("%s - %s", r.ip, logtext))
return r.w.Write(b)
}
func (r responseLogger) WriteHeader(code int) {
log.Printf("%s - Returned code %d\n", r.ip, code)
r.w.WriteHeader(code)
}