package server import ( "context" "fmt" "net/http" "strings" "time" "github.com/google/uuid" "github.com/labstack/echo/v4" "github.com/pkg/errors" storepb "github.com/usememos/memos/proto/gen/store" "github.com/usememos/memos/server/profile" "github.com/usememos/memos/server/route/api/auth" apiv2 "github.com/usememos/memos/server/route/api/v2" "github.com/usememos/memos/server/route/frontend" "github.com/usememos/memos/server/route/resource" "github.com/usememos/memos/server/route/rss" 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 } 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, } // Register CORS middleware. e.Use(CORSMiddleware(s.Profile.Origins)) workspaceBasicSetting, err := s.getOrUpsertWorkspaceBasicSetting(ctx) if err != nil { return nil, errors.Wrap(err, "failed to get workspace basic setting") } secret := "usememos" if profile.Mode == "prod" { secret = workspaceBasicSetting.SecretKey } s.ID = workspaceBasicSetting.ServerId s.Secret = secret // Register healthz endpoint. e.GET("/healthz", func(c echo.Context) error { return c.String(http.StatusOK, "Service ready.") }) // Only serve frontend when it's enabled. if profile.Frontend { frontendService := frontend.NewFrontendService(profile, store) frontendService.Serve(ctx, e) } rootGroup := e.Group("") // Register public routes. publicGroup := rootGroup.Group("/o") publicGroup.Use(func(next echo.HandlerFunc) echo.HandlerFunc { return auth.JWTMiddleware(s.Store, next, s.Secret) }) // Create and register resource public routes. resource.NewResourceService(s.Profile, s.Store).RegisterRoutes(publicGroup) // Create and register RSS routes. rss.NewRSSService(s.Profile, s.Store).RegisterRoutes(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) 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) getOrUpsertWorkspaceBasicSetting(ctx context.Context) (*storepb.WorkspaceBasicSetting, error) { workspaceBasicSetting, err := s.Store.GetWorkspaceBasicSetting(ctx) if err != nil { return nil, errors.Wrap(err, "failed to get workspace basic setting") } modified := false if workspaceBasicSetting.ServerId == "" { workspaceBasicSetting.ServerId = uuid.NewString() modified = true } if workspaceBasicSetting.SecretKey == "" { workspaceBasicSetting.SecretKey = uuid.NewString() modified = true } if modified { workspaceSetting, err := s.Store.UpsertWorkspaceSettingV1(ctx, &storepb.WorkspaceSetting{ Key: storepb.WorkspaceSettingKey_WORKSPACE_SETTING_BASIC, Value: &storepb.WorkspaceSetting_BasicSetting{BasicSetting: workspaceBasicSetting}, }) if err != nil { return nil, errors.Wrap(err, "failed to upsert workspace setting") } workspaceBasicSetting = workspaceSetting.GetBasicSetting() } return workspaceBasicSetting, nil } func grpcRequestSkipper(c echo.Context) bool { return strings.HasPrefix(c.Request().URL.Path, "/memos.api.v2.") } func CORSMiddleware(origins []string) 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 requestOrigin := r.Header.Get("Origin") if len(origins) == 0 { w.Header().Set("Access-Control-Allow-Origin", requestOrigin) } else { for _, origin := range origins { if origin == requestOrigin { w.Header().Set("Access-Control-Allow-Origin", origin) break } } } 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) } } }