micro/gateway/http/server.go

package http

import (
	"context"
	"git.nspix.com/golang/micro/gateway/http/router"
	"golang.org/x/net/websocket"
	"net"
	"net/http"
	"sync"
)

var (
	contextPool sync.Pool
)

type contextRouteKey struct{}

type HandleFunc func(ctx *Context) (err error)

type Middleware func(next HandleFunc) HandleFunc

type Node struct {
	Method string
	Path   string
	Handle HandleFunc
}

type Server struct {
	ctx            context.Context
	svr            *http.Server
	middleware     []Middleware
	keepAlive      bool
	isSetKeepAlive bool
	hookHandle     http.Handler
	router         *router.Router
}

func getContext() *Context {
	if v := contextPool.Get(); v != nil {
		return v.(*Context)
	}
	return &Context{}
}

func releaseContext(c *Context) {
	contextPool.Put(c)
}

func (r *Server) Group(prefix string, nodes []Node, middleware ...Middleware) {
	for _, node := range nodes {
		r.Handle(node.Method, prefix+node.Path, node.Handle, middleware...)
	}
}

func (r *Server) Handler(method, path string, h http.Handler) {
	r.router.Handler(method, path, h)
}

func (r *Server) createWebsocket(h websocket.Handler) *websocket.Server {
	return &websocket.Server{
		Handshake: func(cfg *websocket.Config, req *http.Request) (err error) {
			cfg.Origin, err = websocket.Origin(cfg, req)
			if err == nil && cfg != nil {
				if cfg.Header == nil {
					cfg.Header = http.Header(make(map[string][]string))
				}
				if cfg.Origin != nil {
					cfg.Header.Set("Access-Control-Allow-Origin", cfg.Origin.Host)
				}
				cfg.Header.Set("Access-Control-Request-Headers", "GET,POST,PUT,OPTIONS")
				cfg.Header.Set("Access-Control-Allow-Credentials", "true")
			}
			return
		},
		Handler: h,
	}
}

func (r *Server) Websocket(path string, h websocket.Handler) {
	r.Handler("GET", path, r.createWebsocket(h))
}

func (r *Server) Handle(method string, path string, h HandleFunc, middleware ...Middleware) {
	r.router.Handle(method, path, func(writer http.ResponseWriter, request *http.Request, params router.Params) {
		ctx := getContext()
		ps := make(map[string]string)
		for _, v := range params {
			ps[v.Key] = v.Value
		}
		ctx.Reset(request, writer, ps)
		for i := len(r.middleware) - 1; i >= 0; i-- {
			h = r.middleware[i](h)
		}
		for i := len(middleware) - 1; i >= 0; i-- {
			h = middleware[i](h)
		}
		if err := h(ctx); err != nil {
			writer.WriteHeader(http.StatusInternalServerError)
		}
		releaseContext(ctx)
	})
}

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

func (r *Server) WithHandle(h http.Handler) {
	r.hookHandle = h
}

func (r *Server) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	//hook all request
	if r.hookHandle != nil {
		r.hookHandle.ServeHTTP(res, req)
		return
	}
	if req.Method == http.MethodOptions {
		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)
		return
	} else {
		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)
		}
		r.router.ServeHTTP(res, req)
	}
}

func (r *Server) Server() *http.Server {
	return r.svr
}

func (r *Server) SetKeepAlivesEnabled(enable bool) {
	r.isSetKeepAlive = true
	r.keepAlive = enable
}

func (r *Server) ListenAndServe(addr string) (err error) {
	r.svr = &http.Server{
		Addr:    addr,
		Handler: r,
	}
	if r.isSetKeepAlive {
		r.svr.SetKeepAlivesEnabled(r.keepAlive)
	}
	return r.svr.ListenAndServe()
}

func (r *Server) Serve(l net.Listener) (err error) {
	r.svr = &http.Server{
		Handler: r,
	}
	if r.isSetKeepAlive {
		r.svr.SetKeepAlivesEnabled(r.keepAlive)
	}
	return r.svr.Serve(l)
}

func (r *Server) Shutdown(ctx context.Context) (err error) {
	if r.svr != nil {
		return r.svr.Shutdown(ctx)
	}
	return
}

func New(ctx context.Context) *Server {
	return &Server{
		ctx:    ctx,
		router: router.New(),
	}
}

func WithContext(ctx context.Context, route *Server) context.Context {
	return context.WithValue(ctx, contextRouteKey{}, route)
}

func FromContext(ctx context.Context) *Server {
	if v := ctx.Value(contextRouteKey{}); v != nil {
		return v.(*Server)
	}
	return nil
}