opencatd-open/cmd/openteam/main.go

package main

import (
	"context"
	"embed"
	"fmt"
	"io/fs"
	"log"
	"net/http"
	"opencatd-open/middleware"
	"opencatd-open/pkg/config"
	"opencatd-open/pkg/store"
	"opencatd-open/wire"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"
)

//go:embed dist/*
var web embed.FS

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup

	cfg, err := config.LoadConfig()
	if err != nil {
		panic(err)
	}

	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)

		teamGroup.POST("/keys", team.CreateKey)
		teamGroup.GET("/keys", team.ListKeys)
		teamGroup.POST("/keys/:id", team.UpdateKey)
		teamGroup.DELETE("/keys/:id", team.DeleteKey)

		teamGroup.POST("/users", team.CreateUser)
		teamGroup.GET("/users", team.ListUsers)
		teamGroup.POST("/users/:id/reset", team.ResetUserToken)
		teamGroup.DELETE("/users/:id", team.DeleteUser)

		teamGroup.GET("/1/usages", team.ListUsages)
	}

	public := r.Group("/api/auth")
	{
		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(db))
	{
		// v1.POST("/v2/*proxypath", router.HandleProxy)
		v1.POST("/*proxypath", proxy.HandleProxy)
		// v1.GET("/models", dashboard.HandleModels)
	}

	idxFS, err := fs.Sub(web, "dist")
	if err != nil {
		panic(err)
	}

	assetsFS, err := fs.Sub(web, "dist/assets")
	if err != nil {
		panic(err)
	}
	r.StaticFS("/assets", http.FS(assetsFS))

	r.NoRoute(func(c *gin.Context) {
		if c.Writer.Status() == http.StatusNotFound {
			c.FileFromFS("/", http.FS(idxFS))
		}
	})

	srv := &http.Server{
		Addr:    fmt.Sprintf(":%d", cfg.Port),
		Handler: r,
	}

	go func() {
		// 服务启动
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("listen: %s\n", err)
		}
	}()

	// 等待中断信号来优雅地关闭服务器
	quit := make(chan os.Signal, 1)
	// kill (no param) default send syscall.SIGTERM
	// kill -2 is syscall.SIGINT
	// kill -9 is syscall.SIGKILL but can't be catch
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	fmt.Println("\nShutdown Server ...")
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer shutdownCancel()

	if err := srv.Shutdown(shutdownCtx); err != nil {
		log.Fatalln("Server Shutdown:", err)
	}

	cancel()

	sqlDB.Close()

	waitChan := make(chan struct{})
	go func() {
		wg.Wait()
		close(waitChan)
	}()

	select {
	case <-waitChan:
		fmt.Println("All goroutines have finished")
	case <-shutdownCtx.Done():
		fmt.Println("⚠️ Shutdown timeout")
	}

	fmt.Println("Server exited")

}

func printFilesAndDirs(fsys fs.FS, prefix string) error {
	return fs.WalkDir(fsys, ".", func(p string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}

		if d.IsDir() {
			fmt.Printf("%s[DIR] %s\n", prefix, p)
		} else {
			info, err := d.Info()
			if err != nil {
				return err
			}
			fmt.Printf("%s[FILE] %s (%d bytes)\n", prefix, p, info.Size())
		}
		return nil
	})
}