aigrammar/main.go

package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/golang-jwt/jwt"
	echojwt "github.com/labstack/echo-jwt/v4"
	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/natefinch/lumberjack"
	"github.com/sirupsen/logrus"
	"go.uber.org/zap"
)

// 私有变量，只能在main包内访问
var jwtSigningKey []byte

func main() {
	//检查时区配置，需要包含东八区
	_, err := time.LoadLocation(KEY_LOCAL_TIMEZONE)
	if err != nil {
		fmt.Println("Error loading location:", err)
		return
	}

	// 获取配置
	configManager, err := GetConfigManager()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load config: %v\n", err)
		os.Exit(1) // Exit the program with an error code
	}

	logconfig := configManager.GetLogConfig()
	initLogger(logconfig.LogFile, logconfig.MaxSize, logconfig.MaxBackups, logconfig.MaxAge, logconfig.Compress, logconfig.Level) // 初始化全局日志
	defer logger.Sync()                                                                                                           // 刷到磁盘

	// 初始化数据库，并创建实例
	dbManager, errdb := GetDBManager()
	if errdb != nil {
		logger.Fatal("DBManager init error", zap.Error(errdb)) // 记录错误信息
		os.Exit(1)                                             // Exit the program with an error code
	} else {
		logger.Info("DBManager init successfully. Mysql ", zap.String("mysql_conn", configManager.GetDatabaseConfig().MysqlConn))
	}
	defer dbManager.CloseDB()

	// JWT密钥，写到配置文件中
	baseConfig := configManager.GetBaseConfig()
	jwtSigningKey = []byte(baseConfig.JwtSecret) // 设置JWT密钥

	e := echo.New()

	// 处理日志，格式可定义，日志输出到文件，且文件自动分割
	logger := logrus.New()
	logger.SetFormatter(&logrus.JSONFormatter{})

	// Configure Lumberjack for log rotation
	logOutput := &lumberjack.Logger{
		Filename:   logconfig.EchoLogFile,
		MaxSize:    logconfig.MaxSize, // megabytes
		MaxBackups: logconfig.MaxBackups,
		MaxAge:     logconfig.MaxAge,   // days
		Compress:   logconfig.Compress, // disabled by default
	}

	// Set output of logger to both stdout and Lumberjack
	multiWriter := io.MultiWriter(os.Stdout, logOutput)
	logger.SetOutput(multiWriter)

	// Custom log format middleware
	e.Use(middleware.RequestLoggerWithConfig(middleware.RequestLoggerConfig{
		LogStatus:    true,
		LogURI:       true,
		LogMethod:    true,
		LogRemoteIP:  true,
		LogUserAgent: true,
		LogValuesFunc: func(c echo.Context, v middleware.RequestLoggerValues) error {
			logger.WithFields(logrus.Fields{
				"status":     v.Status,
				"uri":        v.URI,
				"method":     v.Method,
				"remote_ip":  c.RealIP(),
				"user_agent": v.UserAgent,
			}).Info("request log")
			return nil
		},
	}))

	// 全局中间件
	//e.Use(middleware.Logger())
	e.Use(middleware.Recover())
	e.Use(unifiedResponseHandler) // Register the unified response handler
	e.HTTPErrorHandler = customHTTPErrorHandler

	// 定义一组无需jwt鉴权的处理功能
	s := e.Group("/pub")
	s.GET("/ping", PingHander)
	s.POST("/iap/callback", IapCallbackHandler)

	// 处理应用商店的回调
	p := e.Group("/")
	// 使用 echo-jwt 替换 deprecated JWT middleware
	p.Use(echojwt.WithConfig(echojwt.Config{
		SigningKey:     baseConfig.JwtSecret,
		ContextKey:     "user",
		ParseTokenFunc: parseToken,
	}))
	p.POST("grammar/translate", TranslateHandler)
	p.POST("grammar/grammar", GrammarHandler)
	p.POST("grammar/words", WordsHandler)
	p.POST("grammar/feedback", TranslateFeedBackHandler)

	p.POST("iap/verify", IapVerify)
	p.POST("user/get", queryUserHandler)

	i := e.Group("/internal")
	i.POST("/iap/history", IapHistory)
	i.GET("/user/rights", UserRightsHandler)
	i.GET("/user/rights/reset", ResetUserRightsHandler)

	// 启动服务器
	logger.Fatal("Failed to start server", zap.Error(e.Start(baseConfig.BindAddr)))
	//e.Logger.Fatal(e.Start(baseConfig.BindAddr))
}

// 自定义的JWT Claims结构
type jwtCustomClaims struct {
	DeviceID string `json:"deviceID"`
	GID      int    `json:"gid"`
	Exp1     int64  `json:"exp1"`
	jwt.StandardClaims
}

func parseToken(c echo.Context, auth string) (interface{}, error) {
	logger.Debug("into func")

	token, err := jwt.ParseWithClaims(auth, &jwtCustomClaims{}, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return jwtSigningKey, nil
	})

	if err != nil {
		logger.Fatal("ParaseToken Error.", zap.Error(err))
		return nil, err
	}

	if claims, ok := token.Claims.(*jwtCustomClaims); ok && token.Valid {
		logger.Info("claims: ", zap.Any("claims", claims))
		// 判断token有效期
		if time.Now().Unix() > claims.Exp1 {
			return nil, echo.NewHTTPError(http.StatusUnauthorized, "Token expired")
		}

		// 设置参数
		c.Set(KEY_DEVICEID, claims.DeviceID)
		c.Set(KEY_GID, claims.GID)
		return token, nil
	}

	return nil, fmt.Errorf("invalid token")
}

// 处理翻译请求
func PingHander(c echo.Context) error {
	setResponse(c, map[string]string{"pang": "ok"})
	return nil
}