package http import ( "context" "errors" "net" "net/http" "net/url" "strings" "time" ic "github.com/go-kratos/kratos/v2/internal/context" "github.com/go-kratos/kratos/v2/internal/host" "github.com/go-kratos/kratos/v2/log" "github.com/go-kratos/kratos/v2/transport" "github.com/gorilla/mux" ) var _ transport.Server = (*Server)(nil) var _ transport.Endpointer = (*Server)(nil) // ServerOption is an HTTP server option. type ServerOption func(*Server) // Network with server network. func Network(network string) ServerOption { return func(s *Server) { s.network = network } } // Address with server address. func Address(addr string) ServerOption { return func(s *Server) { s.address = addr } } // Timeout with server timeout. func Timeout(timeout time.Duration) ServerOption { return func(s *Server) { s.timeout = timeout } } // Logger with server logger. func Logger(logger log.Logger) ServerOption { return func(s *Server) { s.log = log.NewHelper(logger) } } // Server is an HTTP server wrapper. type Server struct { *http.Server ctx context.Context lis net.Listener network string address string endpoint *url.URL timeout time.Duration router *mux.Router log *log.Helper } // NewServer creates an HTTP server by options. func NewServer(opts ...ServerOption) *Server { srv := &Server{ network: "tcp", address: ":0", timeout: 1 * time.Second, log: log.NewHelper(log.DefaultLogger), } for _, o := range opts { o(srv) } srv.router = mux.NewRouter() srv.Server = &http.Server{Handler: srv} return srv } // Handle registers a new route with a matcher for the URL path. func (s *Server) Handle(path string, h http.Handler) { s.router.Handle(path, h) } // HandlePrefix registers a new route with a matcher for the URL path prefix. func (s *Server) HandlePrefix(prefix string, h http.Handler) { s.router.PathPrefix(prefix).Handler(h) } // HandleFunc registers a new route with a matcher for the URL path. func (s *Server) HandleFunc(path string, h http.HandlerFunc) { s.router.HandleFunc(path, h) } // ServeHTTP should write reply headers and data to the ResponseWriter and then return. func (s *Server) ServeHTTP(res http.ResponseWriter, req *http.Request) { ctx, cancel := ic.Merge(req.Context(), s.ctx) defer cancel() ctx = transport.NewContext(ctx, transport.Transport{Kind: transport.KindHTTP}) ctx = NewServerContext(ctx, ServerInfo{Request: req, Response: res, Endpoint: s.endpoint}) if s.timeout > 0 { ctx, cancel = context.WithTimeout(ctx, s.timeout) defer cancel() } s.router.ServeHTTP(res, req.WithContext(ctx)) } // Endpoint return a real address to registry endpoint. // examples: // http://127.0.0.1:8000?isSecure=false func (s *Server) Endpoint() (*url.URL, error) { if s.lis == nil && strings.HasSuffix(s.address, ":0") { lis, err := net.Listen(s.network, s.address) if err != nil { return nil, err } s.lis = lis } addr, err := host.Extract(s.address, s.lis) if err != nil { return nil, err } u := &url.URL{ Scheme: "http", Host: addr, } s.endpoint = u return u, nil } // Start start the HTTP server. func (s *Server) Start(ctx context.Context) error { s.ctx = ctx if s.lis == nil { lis, err := net.Listen(s.network, s.address) if err != nil { return err } s.lis = lis } s.log.Infof("[HTTP] server listening on: %s", s.lis.Addr().String()) if err := s.Serve(s.lis); !errors.Is(err, http.ErrServerClosed) { return err } return nil } // Stop stop the HTTP server. func (s *Server) Stop(ctx context.Context) error { s.log.Info("[HTTP] server stopping") return s.Shutdown(context.Background()) }