go http chain-middleware

README.md

go-playground

mw_test.go

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
}