reface to openteam

cmd/openteam/main.go

@@ -2,13 +2,16 @@ package main
 
 import (
 	"context"
+	"fmt"
 	"log"
 	"net/http"
+	"opencatd-open/middleware"
+	"opencatd-open/pkg/config"
 	"opencatd-open/pkg/store"
-	"opencatd-open/team/dashboard"
 	"opencatd-open/wire"
 	"os"
 	"os/signal"
+	"sync"
 	"syscall"
 	"time"
 
@@ -16,26 +19,47 @@ import (
 )
 
 func main() {
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
 
-	_, err := store.InitDB()
+	cfg, err := config.LoadConfig()
 	if err != nil {
 		panic(err)
 	}
 
-	team, err := wire.InitTeamHandler(ctx, store.DB)
+	db, err := store.InitDB(cfg)
+	if err != nil {
+		panic(err)
+	}
+	sqlDB, err := db.DB()
+	if err != nil {
+		log.Fatalf("Failed to get underlying *sql.DB: %v", err)
+	}
+
+	team, err := wire.InitTeamHandler(ctx, cfg, db)
+	if err != nil {
+		panic(err)
+	}
+
+	api, err := wire.InitAPIHandler(ctx, cfg, db)
+	if err != nil {
+		panic(err)
+	}
+
+	proxy, err := wire.InitProxyHandler(ctx, cfg, db, &wg)
 	if err != nil {
 		panic(err)
 	}
 
 	r := gin.Default()
+	r.Use(middleware.CORS())
 	teamGroup := r.Group("/1")
 	teamGroup.Use(team.AuthMiddleware())
 	{
 		teamGroup.POST("/users/init", team.InitAdmin)
 		// 获取当前用户信息
 		teamGroup.GET("/me", team.Me)
-		//// team.GET("/me/usages", team.HandleMeUsage)
+		// team.GET("/me/usages", team.HandleMeUsage)
 
 		teamGroup.POST("/keys", team.CreateKey)
 		teamGroup.GET("/keys", team.ListKeys)
@@ -50,9 +74,59 @@ func main() {
 		teamGroup.GET("/1/usages", team.ListUsages)
 	}
 
-	api := r.Group("/api")
+	public := r.Group("/api/auth")
 	{
-		api.POST("/login", dashboard.HandleLogin)
+		public.GET("/passkey/begin", api.PasskeyAuthBegin)
+		public.POST("/passkey/finish", api.PasskeyAuthFinish)
+
+		public.POST("/register", api.Register)
+		public.POST("/login", api.Login)
+	}
+
+	apiGroup := r.Group("/api", middleware.Auth)
+	{
+		apiGroup.GET("/profile", api.Profile)
+		apiGroup.POST("/profile/update", api.UpdateProfile)
+		apiGroup.POST("/profile/update/password", api.UpdatePassword)
+		// 绑定PassKey
+		apiGroup.GET("/profile/passkey", api.PasskeyCreateBegin)
+		apiGroup.POST("/profile/passkey", api.PasskeyCreateFinish)
+		apiGroup.GET("/profile/passkeys", api.ListPasskey)
+		apiGroup.DELETE("/profile/passkeys/:id", api.DeletePasskey)
+
+		userGroup := apiGroup.Group("/users")
+		{
+			userGroup.POST("", api.CreateUser)
+			userGroup.GET("", api.ListUser)
+			userGroup.GET("/:id", api.GetUser)
+			userGroup.PUT("/:id", api.EditUser)
+			userGroup.DELETE("/:id", api.DeleteUser)
+			userGroup.POST("/batch/:option", api.UserOption)
+		}
+
+		tokenGroup := apiGroup.Group("/tokens")
+		tokenGroup.POST("", api.CreateToken)
+		tokenGroup.GET("", api.ListToken)
+		// tokenGroup.GET("/:id", api.GetToken)
+		tokenGroup.POST("/reset/:id", api.ResetToken)
+		tokenGroup.PUT("/:id", api.UpdateToken)
+		tokenGroup.DELETE("/:id", api.DeleteToken)
+		// tokenGroup.POST("/batch/:option", api.TokenOption)
+
+		apiGroup.POST("keys", api.CreateApiKey)
+		apiGroup.GET("keys", api.ListApiKey)
+		apiGroup.GET("keys/:id", api.GetApiKey)
+		apiGroup.PUT("keys/:id", api.UpdateApiKey)
+		apiGroup.DELETE("keys/:id", api.DeleteApiKey)
+		apiGroup.POST("keys/batch/:option", api.ApiKeyOption)
+	}
+
+	v1 := r.Group("/v1")
+	v1.Use(middleware.AuthLLM(store.DB))
+	{
+		// v1.POST("/v2/*proxypath", router.HandleProxy)
+		v1.POST("/v1/*proxypath", proxy.HandleProxy)
+		// v1.GET("/models", dashboard.HandleModels)
 	}
 
 	srv := &http.Server{
@@ -74,19 +148,31 @@ func main() {
 	// kill -9 is syscall.SIGKILL but can't be catch
 	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
 	<-quit
-	log.Println("Shutdown Server ...")
+	fmt.Println("\nShutdown Server ...")
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer shutdownCancel()
 
-	defer cancel()
-	if err := srv.Shutdown(ctx); err != nil {
-		log.Fatal("Server Shutdown:", err)
+	if err := srv.Shutdown(shutdownCtx); err != nil {
+		log.Fatalln("Server Shutdown:", err)
 	}
-	db, _ := store.DB.DB()
-	db.Close()
-	// catching ctx.Done(). timeout of 1 seconds.
+
+	cancel()
+
+	sqlDB.Close()
+
+	waitChan := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(waitChan)
+	}()
+
 	select {
-	case <-ctx.Done():
-		log.Println("timeout of 5 seconds.")
+	case <-waitChan:
+		fmt.Println("All goroutines have finished")
+	case <-shutdownCtx.Done():
+		fmt.Println("⚠️ Shutdown timeout")
 	}
-	log.Println("Server exiting")
+
+	fmt.Println("Server exited")
 
 }