beta
This commit is contained in:
“xHuPo”
parent
a45ddf13d5
commit
bcd986e3f7
46 changed files with 6166 additions and 454 deletions
230
cmd/root.go
230
cmd/root.go
|
|
@ -2,151 +2,139 @@ package cmd
|
|||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
|
||||
"otpm/config"
|
||||
"otpm/database"
|
||||
"otpm/handlers"
|
||||
"otpm/utils"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/jmoiron/sqlx"
|
||||
"github.com/julienschmidt/httprouter"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/viper"
|
||||
"otpm/models"
|
||||
"otpm/server"
|
||||
"otpm/services"
|
||||
)
|
||||
|
||||
var rootCmd = &cobra.Command{
|
||||
Use: "otpm",
|
||||
Short: "otp backend for microapp on wechat",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
startApp()
|
||||
},
|
||||
}
|
||||
|
||||
func Execute() {
|
||||
if err := rootCmd.Execute(); err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func init() {
|
||||
cobra.OnInitialize(initConfig)
|
||||
// Set config file with multi-environment support
|
||||
viper.SetConfigName("config")
|
||||
viper.SetConfigType("yaml")
|
||||
viper.AddConfigPath(".")
|
||||
|
||||
rootCmd.PersistentFlags().StringP("config", "c", "", "config file (default is $HOME/config.yaml)")
|
||||
rootCmd.PersistentFlags().StringP("driver", "d", "sqlite3", "database driver (sqlite3, postgres, mysql)")
|
||||
rootCmd.PersistentFlags().StringP("dsn", "s", "", "database connection string")
|
||||
rootCmd.PersistentFlags().StringP("port", "p", "8080", "port to listen on")
|
||||
|
||||
viper.BindPFlag("database.driver", rootCmd.PersistentFlags().Lookup("driver"))
|
||||
viper.BindPFlag("database.dsn", rootCmd.PersistentFlags().Lookup("dsn"))
|
||||
viper.BindPFlag("port", rootCmd.PersistentFlags().Lookup("port"))
|
||||
}
|
||||
|
||||
func initConfig() {
|
||||
if cfgFile := viper.GetString("config"); cfgFile != "" {
|
||||
viper.SetConfigFile(cfgFile)
|
||||
} else {
|
||||
viper.AddConfigPath(".")
|
||||
viper.SetConfigName("config")
|
||||
viper.SetConfigType("yaml")
|
||||
// Set environment specific config (e.g. config.production.yaml)
|
||||
env := os.Getenv("OTPM_ENV")
|
||||
if env != "" {
|
||||
viper.SetConfigName(fmt.Sprintf("config.%s", env))
|
||||
}
|
||||
|
||||
if err := viper.ReadInConfig(); err != nil {
|
||||
log.Fatalf("Error reading config file: %v", err)
|
||||
}
|
||||
// Set default values
|
||||
viper.SetDefault("server.port", "8080")
|
||||
viper.SetDefault("server.timeout.read", "15s")
|
||||
viper.SetDefault("server.timeout.write", "15s")
|
||||
viper.SetDefault("server.timeout.idle", "60s")
|
||||
viper.SetDefault("database.max_open_conns", 25)
|
||||
viper.SetDefault("database.max_idle_conns", 5)
|
||||
viper.SetDefault("database.conn_max_lifetime", "5m")
|
||||
|
||||
// Set environment variable prefix
|
||||
viper.SetEnvPrefix("OTPM")
|
||||
viper.AutomaticEnv()
|
||||
|
||||
// Bind environment variables
|
||||
viper.BindEnv("database.url", "OTPM_DB_URL")
|
||||
viper.BindEnv("database.password", "OTPM_DB_PASSWORD")
|
||||
}
|
||||
|
||||
type App struct {
|
||||
db *sqlx.DB
|
||||
router http.Handler
|
||||
port int
|
||||
}
|
||||
|
||||
func NewApp() (*App, error) {
|
||||
db, err := connectDB()
|
||||
// Execute is the entry point for the application
|
||||
func Execute() error {
|
||||
// Load configuration
|
||||
cfg, err := config.LoadConfig()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return fmt.Errorf("failed to load config: %w", err)
|
||||
}
|
||||
|
||||
if err := runMigrations(db); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
router := setupRouter(db)
|
||||
|
||||
port := viper.GetInt("port")
|
||||
|
||||
return &App{
|
||||
db: db,
|
||||
router: router,
|
||||
port: port,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func connectDB() (*sqlx.DB, error) {
|
||||
db, err := database.InitDB()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to connect to the database: %v", err)
|
||||
}
|
||||
return db, nil
|
||||
}
|
||||
|
||||
func runMigrations(db *sqlx.DB) error {
|
||||
if err := database.MigrateDB(db); err != nil {
|
||||
log.Fatalf("Error migrating the database: %v", err)
|
||||
return fmt.Errorf("error migrating the database: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func setupRouter(db *sqlx.DB) http.Handler {
|
||||
handler := &handlers.Handler{DB: db}
|
||||
|
||||
router := httprouter.New()
|
||||
router.POST("/login", utils.AdaptHandler(handler.Login))
|
||||
router.POST("/refresh", utils.AdaptHandler(utils.AuthMiddleware(handler.RefreshToken)))
|
||||
router.POST("/set", utils.AdaptHandler(utils.AuthMiddleware(handler.UpdateOrCreateOtp)))
|
||||
router.GET("/get", utils.AdaptHandler(utils.AuthMiddleware(handler.GetOtp)))
|
||||
|
||||
return router
|
||||
}
|
||||
|
||||
func (a *App) Start() error {
|
||||
server := &http.Server{Addr: fmt.Sprintf(":%d", a.port), Handler: a.router}
|
||||
go func() {
|
||||
if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
log.Fatalf("Failed to start server: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
quit := make(chan os.Signal, 1)
|
||||
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
||||
<-quit
|
||||
log.Println("Shutting down server...")
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
// Create context with cancellation
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
return server.Shutdown(ctx)
|
||||
}
|
||||
// Setup signal handling
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
||||
go func() {
|
||||
sig := <-sigChan
|
||||
log.Printf("Received signal: %v", sig)
|
||||
cancel()
|
||||
}()
|
||||
|
||||
func startApp() {
|
||||
app, err := NewApp()
|
||||
// Initialize database
|
||||
db, err := database.New(&cfg.Database)
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to initialize application: %v", err)
|
||||
return fmt.Errorf("failed to initialize database: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := app.db.Close(); err != nil {
|
||||
log.Printf("Failed to close database connection: %v", err)
|
||||
defer db.Close()
|
||||
|
||||
// Run database migrations
|
||||
if err := database.MigrateWithContext(ctx, db.DB, cfg.Database.SkipMigration); err != nil {
|
||||
return fmt.Errorf("failed to run migrations: %w", err)
|
||||
}
|
||||
|
||||
// Initialize repositories
|
||||
userRepo := models.NewUserRepository(db.DB)
|
||||
otpRepo := models.NewOTPRepository(db.DB)
|
||||
|
||||
// Initialize services
|
||||
authService := services.NewAuthService(cfg, userRepo)
|
||||
otpService := services.NewOTPService(otpRepo)
|
||||
|
||||
// Initialize handlers
|
||||
authHandler := handlers.NewAuthHandler(authService)
|
||||
otpHandler := handlers.NewOTPHandler(otpService)
|
||||
|
||||
// Create and configure server
|
||||
srv := server.New(cfg)
|
||||
|
||||
// Register health check endpoint
|
||||
srv.RegisterHealthCheck()
|
||||
|
||||
// Register public routes with type conversion
|
||||
authRoutes := make(map[string]http.Handler)
|
||||
for path, handler := range authHandler.Routes() {
|
||||
authRoutes[path] = http.HandlerFunc(handler)
|
||||
}
|
||||
srv.RegisterRoutes(authRoutes)
|
||||
|
||||
// Register authenticated routes with type conversion
|
||||
otpRoutes := make(map[string]http.Handler)
|
||||
for path, handler := range otpHandler.Routes() {
|
||||
otpRoutes[path] = http.HandlerFunc(handler)
|
||||
}
|
||||
srv.RegisterAuthRoutes(otpRoutes)
|
||||
|
||||
// Start server in goroutine
|
||||
serverErr := make(chan error, 1)
|
||||
go func() {
|
||||
log.Printf("Starting server on %s:%d", cfg.Server.Host, cfg.Server.Port)
|
||||
if err := srv.Start(); err != nil {
|
||||
serverErr <- fmt.Errorf("server error: %w", err)
|
||||
}
|
||||
}()
|
||||
if err := app.Start(); err != nil {
|
||||
log.Fatalf("Failed to start application: %v", err)
|
||||
|
||||
// Wait for shutdown signal or server error
|
||||
select {
|
||||
case err := <-serverErr:
|
||||
return err
|
||||
case <-ctx.Done():
|
||||
// Graceful shutdown with timeout
|
||||
log.Println("Shutting down server...")
|
||||
if err := srv.Shutdown(); err != nil {
|
||||
return fmt.Errorf("server shutdown error: %w", err)
|
||||
}
|
||||
log.Println("Server stopped gracefully")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue