go-base/webx/web_router.go

package webx

import (
	"net/http"
)

// HandlerFunc 通用的处理器函数类型
type HandlerFunc interface{}

// Router 路由器
type Router struct {
	mux          *http.ServeMux
	globalBefore []func(http.Handler) http.Handler
	globalAfter  []func(http.ResponseWriter, *http.Request)
}

func NewRouter(mux *http.ServeMux) *Router {
	return &Router{
		mux: mux,
	}
}

// UseBefore 在处理器前执行的中间件
func (r *Router) UseBefore(middleware ...func(http.Handler) http.Handler) *Router {
	r.globalBefore = append(r.globalBefore, middleware...)
	return r
}

// UseAfter 在处理器后执行的中间件
func (r *Router) UseAfter(middleware ...func(http.ResponseWriter, *http.Request)) *Router {
	r.globalAfter = append(r.globalAfter, middleware...)
	return r
}

// GET 注册 GET 请求
func (r *Router) GET(path string, handler HandlerFunc) *routeBuilder {
	return r.handle("GET", path, handler)
}

// POST 注册 POST 请求
func (r *Router) POST(path string, handler HandlerFunc) *routeBuilder {
	return r.handle("POST", path, handler)
}

// handle 内部处理方法
func (r *Router) handle(method, path string, handler HandlerFunc) *routeBuilder {
	return &routeBuilder{
		router:  r,
		method:  method,
		path:    path,
		handler: handler,
	}
}

// 路由构建器
type routeBuilder struct {
	router      *Router
	method      string
	path        string
	handler     HandlerFunc
	before      []func(http.Handler) http.Handler
	after       []func(http.ResponseWriter, *http.Request)
	description string
}

// Before 添加路由前中间件
func (rb *routeBuilder) Before(middleware ...func(http.Handler) http.Handler) *routeBuilder {
	rb.before = append(rb.before, middleware...)
	return rb
}

// After 添加路由后中间件
func (rb *routeBuilder) After(middleware ...func(http.ResponseWriter, *http.Request)) *routeBuilder {
	rb.after = append(rb.after, middleware...)
	return rb
}

// Desc 添加描述
func (rb *routeBuilder) Desc(desc string) *routeBuilder {
	rb.description = desc
	return rb
}

// Register 注册路由
func (rb *routeBuilder) Register() {
	// 创建最终的 HTTP 处理器
	finalHandler := rb.createHandler()

	// 应用路由级 before 中间件
	for i := len(rb.before) - 1; i >= 0; i-- {
		finalHandler = rb.before[i](finalHandler)
	}

	// 应用全局 before 中间件
	for i := len(rb.router.globalBefore) - 1; i >= 0; i-- {
		finalHandler = rb.router.globalBefore[i](finalHandler)
	}

	// 注册到 ServeMux
	rb.router.mux.Handle(rb.path, finalHandler)
}

// 创建处理器
func (rb *routeBuilder) createHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		// 方法检查
		if rb.method != "" && req.Method != rb.method {
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}

		// 创建上下文
		ctx := req.Context()

		// 执行处理器
		rb.invokeHandler(ctx, w, req)

		// 执行路由级 after 中间件
		for _, after := range rb.after {
			after(w, req)
		}

		// 执行全局 after 中间件
		for _, after := range rb.router.globalAfter {
			after(w, req)
		}
	})
}