package server import ( "context" "encoding/json" "fmt" "net" "net/http" "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/common/log" "github.com/usememos/memos/common/util" "github.com/usememos/memos/plugin/telegram" "github.com/usememos/memos/server/profile" "github.com/usememos/memos/store" "go.uber.org/zap" ) type Server struct { e *echo.Echo ID string Secret string Profile *profile.Profile Store *store.Store // API services. apiV2Service *apiv2.APIV2Service // Asynchronous runners. backupRunner *BackupRunner 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. backupRunner: NewBackupRunner(store), telegramBot: telegram.NewBotWithHandler(newTelegramHandler(store)), } e.Use(middleware.LoggerWithConfig(middleware.LoggerConfig{ Format: `{"time":"${time_rfc3339}",` + `"method":"${method}","uri":"${uri}",` + `"status":${status},"error":"${error}"}` + "\n", })) e.Use(middleware.Gzip()) e.Use(middleware.CORS()) e.Use(middleware.SecureWithConfig(middleware.SecureConfig{ Skipper: defaultGetRequestSkipper, XSSProtection: "1; mode=block", ContentTypeNosniff: "nosniff", XFrameOptions: "SAMEORIGIN", HSTSPreloadEnabled: false, })) e.Use(middleware.TimeoutWithConfig(middleware.TimeoutConfig{ Skipper: func(c echo.Context) bool { // this is a hack to skip timeout for openai chat streaming // because streaming require to flush response. But the timeout middleware will break it. return c.Request().URL.Path == "/api/v1/openai/chat-streaming" }, ErrorMessage: "Request timeout", Timeout: 30 * time.Second, })) serverID, err := s.getSystemServerID(ctx) if err != nil { return nil, fmt.Errorf("failed to retrieve system server ID: %w", err) } s.ID = serverID embedFrontend(e) secret := "usememos" if profile.Mode == "prod" { secret, err = s.getSystemSecretSessionName(ctx) if err != nil { return nil, fmt.Errorf("failed to retrieve system secret session name: %w", err) } } s.Secret = secret rootGroup := e.Group("") apiV1Service := apiv1.NewAPIV1Service(s.Secret, profile, store) apiV1Service.Register(rootGroup) s.apiV2Service = apiv2.NewAPIV2Service(s.Secret, profile, store, s.Profile.Port+1) // Register gRPC gateway as api v2. if err := s.apiV2Service.RegisterGateway(ctx, e); err != nil { return nil, fmt.Errorf("failed to register gRPC gateway: %w", err) } return s, nil } func (s *Server) Start(ctx context.Context) error { if err := s.createServerStartActivity(ctx); err != nil { return errors.Wrap(err, "failed to create activity") } go s.telegramBot.Start(ctx) go s.backupRunner.Run(ctx) // Start gRPC server. listen, err := net.Listen("tcp", fmt.Sprintf(":%d", s.Profile.Port+1)) if err != nil { return err } go func() { if err := s.apiV2Service.GetGRPCServer().Serve(listen); err != nil { log.Error("grpc server listen error", zap.Error(err)) } }() return s.e.Start(fmt.Sprintf(":%d", 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.GetDB().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.GetSystemSetting(ctx, &store.FindSystemSetting{ Name: apiv1.SystemSettingServerIDName.String(), }) if err != nil { return "", err } if serverIDSetting == nil || serverIDSetting.Value == "" { serverIDSetting, err = s.Store.UpsertSystemSetting(ctx, &store.SystemSetting{ 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.GetSystemSetting(ctx, &store.FindSystemSetting{ Name: apiv1.SystemSettingSecretSessionName.String(), }) if err != nil { return "", err } if secretSessionNameValue == nil || secretSessionNameValue.Value == "" { secretSessionNameValue, err = s.Store.UpsertSystemSetting(ctx, &store.SystemSetting{ Name: apiv1.SystemSettingSecretSessionName.String(), Value: uuid.NewString(), }) if err != nil { return "", err } } return secretSessionNameValue.Value, nil } func (s *Server) createServerStartActivity(ctx context.Context) error { payload := apiv1.ActivityServerStartPayload{ ServerID: s.ID, Profile: s.Profile, } payloadBytes, err := json.Marshal(payload) if err != nil { return errors.Wrap(err, "failed to marshal activity payload") } activity, err := s.Store.CreateActivity(ctx, &store.Activity{ CreatorID: apiv1.UnknownID, Type: apiv1.ActivityServerStart.String(), Level: apiv1.ActivityInfo.String(), Payload: string(payloadBytes), }) if err != nil || activity == nil { return errors.Wrap(err, "failed to create activity") } return err } func defaultGetRequestSkipper(c echo.Context) bool { return c.Request().Method == http.MethodGet } func defaultAPIRequestSkipper(c echo.Context) bool { path := c.Request().URL.Path return util.HasPrefixes(path, "/api", "/api/v1", "api/v2") }