...

Source file src/edge-infra.dev/pkg/x/tonic/server.go

Documentation: edge-infra.dev/pkg/x/tonic

     1  package tonic
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  	"time"
     8  
     9  	"github.com/gin-contrib/cors"
    10  	"github.com/gin-gonic/gin"
    11  	"github.com/go-logr/logr"
    12  )
    13  
    14  const (
    15  	// defaultPort is the default port for the http server.
    16  	defaultPort = "9003"
    17  	// MethodAny http method for catch all route.
    18  	MethodAny = "ANY"
    19  )
    20  
    21  var (
    22  	// NoRoute http route for not found (404).
    23  	NoRoute = Route{
    24  		Path:   "",
    25  		Action: "",
    26  	}
    27  	// ReadyRoute http ready route.
    28  	ReadyRoute = Route{
    29  		Path:   "/ready",
    30  		Action: MethodAny,
    31  		Handlers: []gin.HandlerFunc{func(c *gin.Context) {
    32  			c.String(http.StatusOK, "ok")
    33  		}},
    34  	}
    35  	// HealthRoute http health route.
    36  	HealthRoute = Route{
    37  		Path:   "/health",
    38  		Action: MethodAny,
    39  		Handlers: []gin.HandlerFunc{func(c *gin.Context) {
    40  			c.String(http.StatusOK, "ok")
    41  		}},
    42  	}
    43  )
    44  
    45  // Server represents a http server that is run using the runnable lib.
    46  type Server struct {
    47  	port        string
    48  	mode        string
    49  	routes      []Route
    50  	middlewares []gin.HandlerFunc
    51  	router      *gin.Engine
    52  	server      *http.Server
    53  	logger      logr.Logger
    54  }
    55  
    56  // Route represents a http route.
    57  type Route struct {
    58  	// Path is the http route path e.g. /login
    59  	Path string
    60  	// Action is the http route action e.g. POST
    61  	Action string
    62  	// Handlers is the list of gin handler funcs called when the route is accessed.
    63  	Handlers []gin.HandlerFunc
    64  }
    65  
    66  // New returns a new instance of the Server with defaults.
    67  func New() *Server {
    68  	return NewWithOptions(defaultPort, gin.ReleaseMode)
    69  }
    70  
    71  // NewWithOptions returns a new instance of Server with the specified options.
    72  func NewWithOptions(port string, mode string) *Server {
    73  	return &Server{
    74  		mode:   mode,
    75  		port:   port,
    76  		routes: make([]Route, 0),
    77  	}
    78  }
    79  
    80  // SetRoutes sets the routes that can be accessed in the http server.
    81  func (b *Server) SetRoutes(routes ...Route) *Server {
    82  	b.routes = append(b.routes, routes...)
    83  	return b
    84  }
    85  
    86  // SetMiddlewares sets the middleware for the http server.
    87  func (b *Server) SetMiddlewares(middlewares ...gin.HandlerFunc) *Server {
    88  	b.middlewares = middlewares
    89  	return b
    90  }
    91  
    92  // Start runs the gin server.
    93  func (b *Server) Start(_ context.Context) error {
    94  	gin.SetMode(b.mode)
    95  	gin.DisableConsoleColor()
    96  	b.router = gin.New()
    97  	b.registerMiddlewares()
    98  	b.registerRoutes()
    99  	b.registerServer()
   100  	b.logger.Info("running http server", "port", b.port, "mode", b.mode)
   101  	return b.server.ListenAndServe()
   102  }
   103  
   104  // SetLogger sets the logger.
   105  func (b *Server) SetLogger(logger logr.Logger) *Server {
   106  	b.logger = logger
   107  	return b
   108  }
   109  
   110  // registerMiddlewares registers the middlewares to the gin server.
   111  func (b *Server) registerMiddlewares() {
   112  	for _, middleware := range b.middlewares {
   113  		b.router.Use(middleware)
   114  	}
   115  }
   116  
   117  // registerRoutes registers the routes to the gin server.
   118  func (b *Server) registerRoutes() {
   119  	for _, route := range b.routes {
   120  		switch route.Action {
   121  		case http.MethodGet:
   122  			b.router.GET(route.Path, route.Handlers...)
   123  		case http.MethodPost:
   124  			b.router.POST(route.Path, route.Handlers...)
   125  		case http.MethodPatch:
   126  			b.router.PATCH(route.Path, route.Handlers...)
   127  		case http.MethodDelete:
   128  			b.router.DELETE(route.Path, route.Handlers...)
   129  		case http.MethodPut:
   130  			b.router.PUT(route.Path, route.Handlers...)
   131  		case http.MethodHead:
   132  			b.router.HEAD(route.Path, route.Handlers...)
   133  		case http.MethodOptions:
   134  			b.router.OPTIONS(route.Path, route.Handlers...)
   135  		case MethodAny:
   136  			b.router.Any(route.Path, route.Handlers...)
   137  		default:
   138  			if route.Path == "" {
   139  				b.router.NoRoute(func(c *gin.Context) {
   140  					c.AbortWithStatus(404)
   141  				})
   142  			}
   143  		}
   144  	}
   145  }
   146  
   147  // registerServer registers the http server.
   148  func (b *Server) registerServer() {
   149  	b.server = &http.Server{
   150  		Addr:              fmt.Sprintf(":%s", b.port),
   151  		Handler:           b.router,
   152  		ReadHeaderTimeout: 60 * time.Second,
   153  	}
   154  }
   155  
   156  // WithHealthRoute adds the default health route to the server.
   157  func (b *Server) WithHealthRoute() *Server {
   158  	b.routes = append(b.routes, HealthRoute)
   159  	return b
   160  }
   161  
   162  // WithReadyRoute adds the default ready route to the server.
   163  func (b *Server) WithReadyRoute() *Server {
   164  	b.routes = append(b.routes, ReadyRoute)
   165  	return b
   166  }
   167  
   168  // With404Route adds a not found (404) route to the server.
   169  func (b *Server) With404Route() *Server {
   170  	b.routes = append(b.routes, NoRoute)
   171  	return b
   172  }
   173  
   174  // CorsMiddleware returns an initialized cors middleware.
   175  func CorsMiddleware() gin.HandlerFunc {
   176  	corsConfig := cors.DefaultConfig()
   177  	corsConfig.AllowAllOrigins = true
   178  	corsConfig.AllowHeaders = []string{"*"}
   179  	corsConfig.AllowCredentials = true
   180  	return cors.New(corsConfig)
   181  }
   182  

View as plain text