kos/entry/http/server.go

package http

import (
	"context"
	"embed"
	"git.nobla.cn/golang/kos/entry/http/router"
	"net"
	"net/http"
	"path"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var (
	ctxPool sync.Pool
)

type Server struct {
	ctx                context.Context
	serve              *http.Server
	router             *router.Router
	middleware         []Middleware
	uptime             time.Time
	enableDocumentRoot bool
	fileSystem         http.FileSystem
	beforeRequests     []HandleFunc
	anyRequests        map[string]http.Handler
	exitFlag           int32
}

func (svr *Server) applyContext() *Context {
	if v := ctxPool.Get(); v != nil {
		if ctx, ok := v.(*Context); ok {
			return ctx
		}
	}
	return &Context{}
}

func (svr *Server) releaseContext(ctx *Context) {
	ctxPool.Put(ctx)
}

func (svr *Server) wrapHandle(cb HandleFunc, middleware ...Middleware) router.Handle {
	return func(writer http.ResponseWriter, request *http.Request, params router.Params) {
		ctx := svr.applyContext()
		ps := make(map[string]string, 4)
		defer func() {
			svr.releaseContext(ctx)
			ps = make(map[string]string, 0)
		}()
		for _, v := range params {
			ps[v.Key] = v.Value
		}
		ctx.reset(request, writer, ps)
		if len(svr.beforeRequests) > 0 {
			for i := len(svr.beforeRequests) - 1; i >= 0; i-- {
				if err := svr.beforeRequests[i](ctx); err != nil {
					ctx.Status(http.StatusServiceUnavailable)
					return
				}
			}
		}
		for i := len(svr.middleware) - 1; i >= 0; i-- {
			cb = svr.middleware[i](cb)
		}
		for i := len(middleware) - 1; i >= 0; i-- {
			cb = middleware[i](cb)
		}
		if err := cb(ctx); err != nil {
			ctx.Status(http.StatusServiceUnavailable)
		}
	}
}

func (svr *Server) Before(cb ...HandleFunc) {
	svr.beforeRequests = append(svr.beforeRequests, cb...)
}

func (svr *Server) Use(middleware ...Middleware) {
	svr.middleware = append(svr.middleware, middleware...)
}

func (svr *Server) Any(prefix string, handle http.Handler) {
	if !strings.HasPrefix(prefix, "/") {
		prefix = "/" + prefix
	}
	svr.anyRequests[prefix] = handle
}

func (svr *Server) Handle(method string, path string, cb HandleFunc, middleware ...Middleware) {
	if method == "" {
		method = http.MethodPost
	}
	if path == "" {
		path = "/"
	}
	if !strings.HasPrefix(path, "/") {
		path = "/" + path
	}
	svr.router.Replace(method, path, svr.wrapHandle(cb, middleware...))
}

func (svr *Server) Group(prefix string, routes []Route, middleware ...Middleware) {
	for _, route := range routes {
		svr.Handle(route.Method, path.Join(prefix, route.Path), route.Handle, middleware...)
	}
}

func (svr *Server) Root(prefix string, fs http.FileSystem) {
	svr.enableDocumentRoot = true
	s := newFS(svr.uptime, fs)
	s.SetPrefix(prefix)
	s.DenyAccessDirectory()
	s.SetIndexFile("/index.html")
	svr.fileSystem = s
}

func (svr *Server) Embed(prefix string, root string, embedFs embed.FS) {
	routePath := prefix
	if !strings.HasSuffix(routePath, "/*filepath") {
		if strings.HasSuffix(routePath, "/") {
			routePath += "/"
		} else {
			routePath += "/*filepath"
		}
	}
	httpFs := http.FS(embedFs)
	svr.Handle(MethodGet, routePath, func(ctx *Context) (err error) {
		filename := strings.TrimPrefix(ctx.Request().URL.Path, prefix)
		if filename == "" || filename == "/" {
			filename = root + "/"
			if !strings.HasSuffix(filename, "/") {
				filename = filename + "/"
			}
		} else {
			if !strings.HasPrefix(filename, root) {
				filename = path.Clean(path.Join(root, filename))
			}
		}
		if !strings.HasPrefix(filename, "/") {
			filename = "/" + filename
		}
		ctx.Request().URL.Path = filename
		http.FileServer(newFS(svr.uptime, httpFs)).ServeHTTP(ctx.Response(), ctx.Request())
		return
	})
}

func (svr *Server) Static(path string, root http.FileSystem) {
	if !strings.HasSuffix(path, "/*filepath") {
		if strings.HasSuffix(path, "/") {
			path += "/"
		} else {
			path += "/*filepath"
		}
	}
	svr.router.ServeFiles(path, root)
}

func (svr *Server) handleOption(res http.ResponseWriter, req *http.Request) {
	res.Header().Add("Vary", "Origin")
	res.Header().Add("Vary", "Access-Control-Request-Method")
	res.Header().Add("Vary", "Access-Control-Request-Headers")
	res.Header().Set("Access-Control-Allow-Origin", "*")
	res.Header().Set("Access-Control-Allow-Credentials", "true")
	res.Header().Set("Access-Control-Allow-Methods", "GET,HEAD,PUT,PATCH,POST,DELETE")
	h := req.Header.Get("Access-Control-Request-Headers")
	if h != "" {
		res.Header().Set("Access-Control-Allow-Headers", h)
	}
	res.WriteHeader(http.StatusNoContent)
}

func (svr *Server) handleRequest(res http.ResponseWriter, req *http.Request) {
	res.Header().Add("Vary", "Origin")
	res.Header().Set("Access-Control-Allow-Origin", "*")
	res.Header().Set("Access-Control-Allow-Credentials", "true")
	h := req.Header.Get("Access-Control-Request-Headers")
	if h != "" {
		res.Header().Set("Access-Control-Allow-Headers", h)
	}
	svr.router.ServeHTTP(res, req)
}

func (svr *Server) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	var (
		err  error
		file http.File
	)
	for prefix, handle := range svr.anyRequests {
		if strings.HasPrefix(request.URL.Path, prefix) {
			handle.ServeHTTP(writer, request)
			return
		}
	}
	if svr.enableDocumentRoot && request.Method == http.MethodGet {
		uri := path.Clean(request.URL.Path)
		if file, err = svr.fileSystem.Open(uri); err == nil {
			http.ServeContent(writer, request, path.Base(uri), svr.uptime, file)
			err = file.Close()
			return
		}
	}
	switch request.Method {
	case http.MethodOptions:
		svr.handleOption(writer, request)
	default:
		svr.handleRequest(writer, request)
	}
}

func (svr *Server) Serve(l net.Listener) (err error) {
	svr.serve = &http.Server{
		Handler: svr,
	}
	svr.router.NotFound = NotFound{}
	svr.router.MethodNotAllowed = NotAllowed{}
	atomic.StoreInt32(&svr.exitFlag, 0)
	return svr.serve.Serve(l)
}

func (svr *Server) Shutdown() (err error) {
	if !atomic.CompareAndSwapInt32(&svr.exitFlag, 0, 1) {
		return
	}
	if svr.serve != nil {
		err = svr.serve.Shutdown(svr.ctx)
	}
	return
}

func New(ctx context.Context) *Server {
	svr := &Server{
		ctx:            ctx,
		uptime:         time.Now(),
		router:         router.New(),
		beforeRequests: make([]HandleFunc, 0, 10),
		anyRequests:    make(map[string]http.Handler),
		middleware:     make([]Middleware, 0, 10),
	}
	return svr
}