123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197 |
- package server
- import (
- "context"
- "fmt"
- "net/http"
- "strings"
- "time"
- "github.com/google/uuid"
- "github.com/labstack/echo/v4"
- "github.com/labstack/echo/v4/middleware"
- "github.com/pkg/errors"
- apiv1 "github.com/usememos/memos/api/v1"
- apiv2 "github.com/usememos/memos/api/v2"
- "github.com/usememos/memos/plugin/telegram"
- "github.com/usememos/memos/server/frontend"
- "github.com/usememos/memos/server/integration"
- "github.com/usememos/memos/server/profile"
- versionchecker "github.com/usememos/memos/server/service/version_checker"
- "github.com/usememos/memos/store"
- )
- type Server struct {
- e *echo.Echo
- ID string
- Secret string
- Profile *profile.Profile
- Store *store.Store
- // Asynchronous runners.
- telegramBot *telegram.Bot
- }
- func NewServer(ctx context.Context, profile *profile.Profile, store *store.Store) (*Server, error) {
- e := echo.New()
- e.Debug = true
- e.HideBanner = true
- e.HidePort = true
- s := &Server{
- e: e,
- Store: store,
- Profile: profile,
- // Asynchronous runners.
- telegramBot: telegram.NewBotWithHandler(integration.NewTelegramHandler(store)),
- }
- e.Use(middleware.LoggerWithConfig(middleware.LoggerConfig{
- Format: `{"time":"${time_rfc3339}","latency":"${latency_human}",` +
- `"method":"${method}","uri":"${uri}",` +
- `"status":${status},"error":"${error}"}` + "\n",
- }))
- e.Use(CORSMiddleware())
- e.Use(middleware.TimeoutWithConfig(middleware.TimeoutConfig{
- Skipper: grpcRequestSkipper,
- Timeout: 30 * time.Second,
- }))
- serverID, err := s.getSystemServerID(ctx)
- if err != nil {
- return nil, errors.Wrap(err, "failed to retrieve system server ID")
- }
- s.ID = serverID
- // Only serve frontend when it's enabled.
- if profile.Frontend {
- frontendService := frontend.NewFrontendService(profile, store)
- frontendService.Serve(ctx, e)
- }
- secret := "usememos"
- if profile.Mode == "prod" {
- secret, err = s.getSystemSecretSessionName(ctx)
- if err != nil {
- return nil, errors.Wrap(err, "failed to retrieve system secret session name")
- }
- }
- s.Secret = secret
- // Register healthz endpoint.
- e.GET("/healthz", func(c echo.Context) error {
- return c.String(http.StatusOK, "Service ready.")
- })
- // Register API v1 endpoints.
- rootGroup := e.Group("")
- apiV1Service := apiv1.NewAPIV1Service(s.Secret, profile, store, s.telegramBot)
- apiV1Service.Register(rootGroup)
- apiV2Service := apiv2.NewAPIV2Service(s.Secret, profile, store, s.Profile.Port+1)
- // Register gRPC gateway as api v2.
- if err := apiV2Service.RegisterGateway(ctx, e); err != nil {
- return nil, errors.Wrap(err, "failed to register gRPC gateway")
- }
- return s, nil
- }
- func (s *Server) Start(ctx context.Context) error {
- go versionchecker.NewVersionChecker(s.Store, s.Profile).Start(ctx)
- go s.telegramBot.Start(ctx)
- return s.e.Start(fmt.Sprintf("%s:%d", s.Profile.Addr, s.Profile.Port))
- }
- func (s *Server) Shutdown(ctx context.Context) {
- ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
- defer cancel()
- // Shutdown echo server
- if err := s.e.Shutdown(ctx); err != nil {
- fmt.Printf("failed to shutdown server, error: %v\n", err)
- }
- // Close database connection
- if err := s.Store.Close(); err != nil {
- fmt.Printf("failed to close database, error: %v\n", err)
- }
- fmt.Printf("memos stopped properly\n")
- }
- func (s *Server) GetEcho() *echo.Echo {
- return s.e
- }
- func (s *Server) getSystemServerID(ctx context.Context) (string, error) {
- serverIDSetting, err := s.Store.GetWorkspaceSetting(ctx, &store.FindWorkspaceSetting{
- Name: apiv1.SystemSettingServerIDName.String(),
- })
- if err != nil {
- return "", err
- }
- if serverIDSetting == nil || serverIDSetting.Value == "" {
- serverIDSetting, err = s.Store.UpsertWorkspaceSetting(ctx, &store.WorkspaceSetting{
- Name: apiv1.SystemSettingServerIDName.String(),
- Value: uuid.NewString(),
- })
- if err != nil {
- return "", err
- }
- }
- return serverIDSetting.Value, nil
- }
- func (s *Server) getSystemSecretSessionName(ctx context.Context) (string, error) {
- secretSessionNameValue, err := s.Store.GetWorkspaceSetting(ctx, &store.FindWorkspaceSetting{
- Name: apiv1.SystemSettingSecretSessionName.String(),
- })
- if err != nil {
- return "", err
- }
- if secretSessionNameValue == nil || secretSessionNameValue.Value == "" {
- secretSessionNameValue, err = s.Store.UpsertWorkspaceSetting(ctx, &store.WorkspaceSetting{
- Name: apiv1.SystemSettingSecretSessionName.String(),
- Value: uuid.NewString(),
- })
- if err != nil {
- return "", err
- }
- }
- return secretSessionNameValue.Value, nil
- }
- func grpcRequestSkipper(c echo.Context) bool {
- return strings.HasPrefix(c.Request().URL.Path, "/memos.api.v2.")
- }
- func CORSMiddleware() echo.MiddlewareFunc {
- return func(next echo.HandlerFunc) echo.HandlerFunc {
- return func(c echo.Context) error {
- if grpcRequestSkipper(c) {
- return next(c)
- }
- r := c.Request()
- w := c.Response().Writer
- w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin"))
- w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, PATCH, OPTIONS")
- w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
- w.Header().Set("Access-Control-Allow-Credentials", "true")
- // If it's preflight request, return immediately.
- if r.Method == "OPTIONS" {
- w.WriteHeader(http.StatusOK)
- return nil
- }
- return next(c)
- }
- }
- }
|