Skip to content

Instantly share code, notes, and snippets.

@CarsonSlovoka
Created August 2, 2024 03:18
Show Gist options
  • Save CarsonSlovoka/26434d56b93a1e3b7f6c28446cb24c48 to your computer and use it in GitHub Desktop.
Save CarsonSlovoka/26434d56b93a1e3b7f6c28446cb24c48 to your computer and use it in GitHub Desktop.
go http chain-middleware
package mw_test
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
)
type Middleware func(http.Handler) http.Handler
func NewHandler(mw ...Middleware) func(http.Handler) http.Handler {
return func(h http.Handler) http.Handler {
next := h
for k := len(mw) - 1; k >= 0; k-- {
next = mw[k](next)
// NewHandler(m1, m2, m3)(myHandler)
// => m1(m2(m3(myHandler)))
}
return next
}
}
func NewHandleFunc(mw ...Middleware) func(http.HandlerFunc) http.Handler {
return func(h http.HandlerFunc) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
handler := h
for k := len(mw) - 1; k >= 0; k-- {
curH := mw[k]
nextH := handler
// update the chain
handler = func(w http.ResponseWriter, r *http.Request) {
curH(nextH).ServeHTTP(w, r)
}
}
// Execute the assembled processor chain
handler.ServeHTTP(w, r)
})
}
}
func WithTime(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t := time.Now()
orgContext := r.Context()
newContext := context.WithValue(orgContext, "time", &t)
newRequest := r.WithContext(newContext)
next.ServeHTTP(w, newRequest)
})
}
// WithMsg middleware with decorators
func WithMsg(msg string) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fmt.Println(msg)
next.ServeHTTP(w, r)
})
}
}
func WithLogging(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fmt.Println(r.URL.Path)
next.ServeHTTP(w, r)
})
}
func WithAuth(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
apiKey := r.Header.Get("Authorization")
if apiKey != "api-key-test" {
http.Error(w, "invalid api-key", http.StatusForbidden)
return
}
next.ServeHTTP(w, r)
})
}
func Test_newHandler(t *testing.T) {
request := httptest.NewRequest(http.MethodGet, "/", nil)
request.Header.Add("Authorization", "api-key-test")
finalHandle := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
myTime, ok := ctx.Value("time").(*time.Time)
if !ok {
t.Fatal()
}
fmt.Println(myTime)
})
handler := NewHandler(WithAuth, WithTime, WithLogging)(finalHandle)
handler.ServeHTTP(httptest.NewRecorder(), request)
handler2 := NewHandler(WithAuth, WithTime, WithLogging)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Fatal("Because the Auth verification did not pass, it will not run to this point")
}))
handler2.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/", nil))
}
func Example_newHandlerFunc() {
NewHandler(WithTime, WithLogging)(
http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fmt.Println("hello") }),
).ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/H", nil))
NewHandleFunc(WithTime, WithLogging)( // Using NewHandleFunc can eliminate the need for http.HandlerFunc
func(w http.ResponseWriter, r *http.Request) { fmt.Println("world") },
).ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/W", nil))
// Output:
// /H
// hello
// /W
// world
}
// middleware with decorators
func Example_withMsg() {
NewHandleFunc(WithLogging, WithMsg("123"))(
func(w http.ResponseWriter, r *http.Request) { fmt.Println("hello") },
).ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/W", nil))
// Output:
// /W
// 123
// hello
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment