optimize styles and release v4.0.4

api/service/captcha_service.go

@@ -9,11 +9,11 @@ import (
 )
 
 type CaptchaService struct {
-	config types.ChatPlusApiConfig
+	config types.ApiConfig
 	client *req.Client
 }
 
-func NewCaptchaService(config types.ChatPlusApiConfig) *CaptchaService {
+func NewCaptchaService(config types.ApiConfig) *CaptchaService {
 	return &CaptchaService{
 		config: config,
 		client: req.C().SetTimeout(10 * time.Second),

api/service/license_service.go

@@ -0,0 +1,108 @@
+package service
+
+import (
+	"chatplus/core"
+	"chatplus/core/types"
+	"chatplus/store"
+	"errors"
+	"fmt"
+	"github.com/imroc/req/v3"
+	"github.com/shirou/gopsutil/host"
+	"strings"
+	"time"
+)
+
+type LicenseService struct {
+	config types.ApiConfig
+	levelDB *store.LevelDB
+	license types.License
+	machineId string
+}
+
+func NewLicenseService(server *core.AppServer, levelDB *store.LevelDB) * LicenseService {
+	var license types.License
+	var machineId string
+	_ = levelDB.Get(types.LicenseKey, &license)
+	info, err := host.Info()
+	if err == nil {
+		machineId = info.HostID
+	}
+	return &LicenseService{
+		config: server.Config.ApiConfig,
+		levelDB: levelDB,
+		license: license,
+		machineId: machineId,
+	}
+}
+
+// ActiveLicense 激活 License
+func (s *LicenseService) ActiveLicense(license string, machineId string) error {
+	var res struct {
+		Code    types.BizCode `json:"code"`
+		Message string        `json:"message"`
+		Data    struct {
+			Name      string `json:"name"`
+			License   string `json:"license"`
+			Mid       string `json:"mid"`
+			ExpiredAt int64  `json:"expired_at"`
+			UserNum   int    `json:"user_num"`
+		}
+	}
+	apiURL := fmt.Sprintf("%s/%s", s.config.ApiURL, "api/license/active")
+	response, err := req.C().R().
+		SetBody(map[string]string{"license": license, "machine_id": machineId}).
+		SetSuccessResult(&res).Post(apiURL)
+	if err != nil {
+		return fmt.Errorf("发送激活请求失败: %v", err)
+	}
+
+	if response.IsErrorState() {
+		return fmt.Errorf( "发送激活请求失败：%v", response.Status)
+	}
+
+	if res.Code != types.Success {
+		return fmt.Errorf( "激活失败：%v", res.Message)
+	}
+
+	err = s.levelDB.Put(types.LicenseKey, types.License{
+		Key:       license,
+		MachineId: machineId,
+		UserNum:   res.Data.UserNum,
+		ExpiredAt: res.Data.ExpiredAt,
+		IsActive: true,
+	})
+	if err != nil {
+		return fmt.Errorf("保存许可证书失败：%v", err)
+	}
+
+	return nil
+}
+
+// GetLicense 获取许可信息
+func (s *LicenseService) GetLicense() types.License {
+	return s.license
+}
+
+// IsValidApiURL 判断是否合法的中转 URL
+func (s *LicenseService) IsValidApiURL(uri string) error {
+	// 获得许可授权的直接放行
+	if s.license.IsActive {
+		if s.license.MachineId != s.machineId {
+			return errors.New("系统使用了盗版的许可证书")
+		}
+
+		if time.Now().Unix() > s.license.ExpiredAt {
+			return errors.New("系统许可证书已经过期")
+		}
+		return nil
+	}
+
+	if !strings.HasPrefix(uri, "https://gpt.bemore.lol") &&
+		!strings.HasPrefix(uri, "https://api.openai.com") &&
+		!strings.HasPrefix(uri, "http://cdn.chat-plus.net") &&
+		!strings.HasPrefix(uri, "https://api.chat-plus.net") {
+		return fmt.Errorf("当前 API 地址 %s 不在白名单列表当中。",uri)
+	}
+
+	return nil
+}

api/service/mj/pool.go

@@ -3,6 +3,7 @@ package mj
 import (
 	"chatplus/core/types"
 	logger2 "chatplus/logger"
+	"chatplus/service"
 	"chatplus/service/oss"
 	"chatplus/service/sd"
 	"chatplus/store"
@@ -26,7 +27,7 @@ type ServicePool struct {
 
 var logger = logger2.GetLogger()
 
-func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, appConfig *types.AppConfig) *ServicePool {
+func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, appConfig *types.AppConfig, licenseService *service.LicenseService) *ServicePool {
 	services := make([]*Service, 0)
 	taskQueue := store.NewRedisQueue("MidJourney_Task_Queue", redisCli)
 	notifyQueue := store.NewRedisQueue("MidJourney_Notify_Queue", redisCli)
@@ -35,13 +36,19 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
 		if config.Enabled == false {
 			continue
 		}
+		err := licenseService.IsValidApiURL(config.ApiURL)
+		if err != nil {
+			logger.Error(err)
+			continue
+		}
+		
 		cli := NewPlusClient(config)
 		name := fmt.Sprintf("mj-plus-service-%d", k)
-		service := NewService(name, taskQueue, notifyQueue, db, cli)
+		plusService := NewService(name, taskQueue, notifyQueue, db, cli)
 		go func() {
-			service.Run()
+			plusService.Run()
 		}()
-		services = append(services, service)
+		services = append(services, plusService)
 	}
 
 	for k, config := range appConfig.MjProxyConfigs {
@@ -50,11 +57,11 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
 		}
 		cli := NewProxyClient(config)
 		name := fmt.Sprintf("mj-proxy-service-%d", k)
-		service := NewService(name, taskQueue, notifyQueue, db, cli)
+		proxyService := NewService(name, taskQueue, notifyQueue, db, cli)
 		go func() {
-			service.Run()
+			proxyService.Run()
 		}()
-		services = append(services, service)
+		services = append(services, proxyService)
 	}
 
 	return &ServicePool{