fork-it/gpt-load
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 导出密钥功能

Browse Source
This commit is contained in:
tbphp
2025-07-19 17:57:19 +08:00
parent 0d9cd5b031
commit b9edd91723
12 changed files with 426 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
internal/container/container.go
View File

@@ -51,6 +51,9 @@ func BuildContainer() (*dig.Container, error) {
if err := container.Provide(services.NewKeyService); err != nil {
return nil, err
}
if err := container.Provide(services.NewLogService); err != nil {
return nil, err
}
if err := container.Provide(services.NewLogCleanupService); err != nil {
return nil, err
}

3
internal/handler/handler.go
View File

@@ -23,6 +23,7 @@ type Server struct {
KeyManualValidationService *services.KeyManualValidationService
TaskService *services.TaskService
KeyService *services.KeyService
LogService *services.LogService
CommonHandler *CommonHandler
}
@@ -36,6 +37,7 @@ type NewServerParams struct {
KeyManualValidationService *services.KeyManualValidationService
TaskService *services.TaskService
KeyService *services.KeyService
LogService *services.LogService
CommonHandler *CommonHandler
}
@@ -49,6 +51,7 @@ func NewServer(params NewServerParams) *Server {
KeyManualValidationService: params.KeyManualValidationService,
TaskService: params.TaskService,
KeyService: params.KeyService,
LogService: params.LogService,
CommonHandler: params.CommonHandler,
}
}

36
internal/handler/key_handler.go
View File

@@ -5,6 +5,7 @@ import (
app_errors "gpt-load/internal/errors"
"gpt-load/internal/models"
"gpt-load/internal/response"
"log"
"strconv"
"strings"
@@ -301,3 +302,38 @@ func (s *Server) ClearAllInvalidKeys(c *gin.Context) {
response.Success(c, gin.H{"message": fmt.Sprintf("%d invalid keys cleared.", rowsAffected)})
}
// ExportKeys handles exporting keys to a text file.
func (s *Server) ExportKeys(c *gin.Context) {
groupID, err := validateGroupIDFromQuery(c)
if err != nil {
response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, err.Error()))
return
}
statusFilter := c.Query("status")
if statusFilter == "" {
statusFilter = "all"
}
switch statusFilter {
case "all", models.KeyStatusActive, models.KeyStatusInvalid:
default:
response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "Invalid status filter"))
return
}
group, ok := s.findGroupByID(c, groupID)
if !ok {
return
}
filename := fmt.Sprintf("keys-%s-%s.txt", group.Name, statusFilter)
c.Header("Content-Disposition", "attachment; filename="+filename)
c.Header("Content-Type", "text/plain; charset=utf-8")
err = s.KeyService.StreamKeysToWriter(groupID, statusFilter, c.Writer)
if err != nil {
log.Printf("Failed to stream keys: %v", err)
}
}

62
internal/handler/log_handler.go
View File

@@ -1,13 +1,12 @@
package handler
import (
"strconv"
"time"
"gpt-load/internal/db"
"fmt"
app_errors "gpt-load/internal/errors"
"gpt-load/internal/models"
"gpt-load/internal/response"
"log"
"time"
"github.com/gin-gonic/gin"
)
@@ -17,43 +16,9 @@ type LogResponse struct {
models.RequestLog
}
// GetLogs Get request logs
func GetLogs(c *gin.Context) {
query := db.DB.Model(&models.RequestLog{})
if groupName := c.Query("group_name"); groupName != "" {
query = query.Where("group_name LIKE ?", "%"+groupName+"%")
}
if keyValue := c.Query("key_value"); keyValue != "" {
likePattern := "%" + keyValue[1:len(keyValue)-1] + "%"
query = query.Where("key_value LIKE ?", likePattern)
}
if isSuccessStr := c.Query("is_success"); isSuccessStr != "" {
if isSuccess, err := strconv.ParseBool(isSuccessStr); err == nil {
query = query.Where("is_success = ?", isSuccess)
}
}
if statusCodeStr := c.Query("status_code"); statusCodeStr != "" {
if statusCode, err := strconv.Atoi(statusCodeStr); err == nil {
query = query.Where("status_code = ?", statusCode)
}
}
if sourceIP := c.Query("source_ip"); sourceIP != "" {
query = query.Where("source_ip = ?", sourceIP)
}
if errorContains := c.Query("error_contains"); errorContains != "" {
query = query.Where("error_message LIKE ?", "%"+errorContains+"%")
}
if startTimeStr := c.Query("start_time"); startTimeStr != "" {
if startTime, err := time.Parse(time.RFC3339, startTimeStr); err == nil {
query = query.Where("timestamp >= ?", startTime)
}
}
if endTimeStr := c.Query("end_time"); endTimeStr != "" {
if endTime, err := time.Parse(time.RFC3339, endTimeStr); err == nil {
query = query.Where("timestamp <= ?", endTime)
}
}
// GetLogs handles fetching request logs with filtering and pagination.
func (s *Server) GetLogs(c *gin.Context) {
query := s.LogService.GetLogsQuery(c)
var logs []models.RequestLog
query = query.Order("timestamp desc")
@@ -66,3 +31,18 @@ func GetLogs(c *gin.Context) {
pagination.Items = logs
response.Success(c, pagination)
}
// ExportLogs handles exporting filtered log keys to a CSV file.
func (s *Server) ExportLogs(c *gin.Context) {
filename := fmt.Sprintf("log_keys_export_%s.csv", time.Now().Format("20060102150405"))
c.Header("Content-Disposition", "attachment; filename="+filename)
c.Header("Content-Type", "text/csv; charset=utf-8")
// Stream the response
err := s.LogService.StreamLogKeysToCSV(c, c.Writer)
if err != nil {
log.Printf("Failed to stream log keys to CSV: %v", err)
c.JSON(500, gin.H{"error": "Failed to export logs"})
return
}
}

10
internal/middleware/middleware.go
View File

@@ -135,7 +135,7 @@ func Auth(
if strings.HasPrefix(path, "/api") {
// Handle backend API authentication
key = extractBearerKey(c)
key = extractApiKey(c)
} else if strings.HasPrefix(path, "/proxy/") {
// Handle proxy authentication
key, err = extractProxyKey(c, groupManager, channelFactory)
@@ -228,7 +228,7 @@ func isMonitoringEndpoint(path string) bool {
}
// extractBearerKey extracts a key from the "Authorization: Bearer <key>" header.
func extractBearerKey(c *gin.Context) string {
func extractApiKey(c *gin.Context) string {
authHeader := c.GetHeader("Authorization")
if authHeader != "" {
const bearerPrefix = "Bearer "
@@ -236,6 +236,12 @@ func extractBearerKey(c *gin.Context) string {
return authHeader[len(bearerPrefix):]
}
}
authKey := c.Query("auth_key")
if authKey != "" {
return authKey
}
return ""
}

7
internal/router/router.go
View File

@@ -121,6 +121,7 @@ func registerProtectedAPIRoutes(api *gin.RouterGroup, serverHandler *handler.Ser
keys := api.Group("/keys")
{
keys.GET("", serverHandler.ListKeysInGroup)
keys.GET("/export", serverHandler.ExportKeys)
keys.POST("/add-multiple", serverHandler.AddMultipleKeys)
keys.POST("/delete-multiple", serverHandler.DeleteMultipleKeys)
keys.POST("/restore-multiple", serverHandler.RestoreMultipleKeys)
@@ -141,7 +142,11 @@ func registerProtectedAPIRoutes(api *gin.RouterGroup, serverHandler *handler.Ser
}
// 日志
api.GET("/logs", handler.GetLogs)
logs := api.Group("/logs")
{
logs.GET("", serverHandler.GetLogs)
logs.GET("/export", serverHandler.ExportLogs)
}
// 设置
settings := api.Group("/settings")

26
internal/services/key_service.go
View File

@@ -5,6 +5,7 @@ import (
"fmt"
"gpt-load/internal/keypool"
"gpt-load/internal/models"
"io"
"regexp"
"strings"
@@ -306,3 +307,28 @@ func (s *KeyService) TestMultipleKeys(group *models.Group, keysText string) ([]k
return allResults, nil
}
// StreamKeysToWriter fetches keys from the database in batches and writes them to the provided writer.
func (s *KeyService) StreamKeysToWriter(groupID uint, statusFilter string, writer io.Writer) error {
query := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Select("id, key_value")
switch statusFilter {
case models.KeyStatusActive, models.KeyStatusInvalid:
query = query.Where("status = ?", statusFilter)
case "all":
default:
return fmt.Errorf("invalid status filter: %s", statusFilter)
}
var keys []models.APIKey
err := query.FindInBatches(&keys, chunkSize, func(tx *gorm.DB, batch int) error {
for _, key := range keys {
if _, err := writer.Write([]byte(key.KeyValue + "\n")); err != nil {
return err
}
}
return nil
}).Error
return err
}

121
internal/services/log_service.go Normal file
View File

@@ -0,0 +1,121 @@
package services
import (
"encoding/csv"
"fmt"
"gpt-load/internal/models"
"io"
"strconv"
"time"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
)
// ExportableLogKey defines the structure for the data to be exported to CSV.
type ExportableLogKey struct {
KeyValue string `gorm:"column:key_value"`
GroupName string `gorm:"column:group_name"`
StatusCode int `gorm:"column:status_code"`
}
// LogService provides services related to request logs.
type LogService struct {
DB *gorm.DB
}
// NewLogService creates a new LogService.
func NewLogService(db *gorm.DB) *LogService {
return &LogService{DB: db}
}
// logFiltersScope returns a GORM scope function that applies filters from the Gin context.
func logFiltersScope(c *gin.Context) func(db *gorm.DB) *gorm.DB {
return func(db *gorm.DB) *gorm.DB {
if groupName := c.Query("group_name"); groupName != "" {
db = db.Where("group_name LIKE ?", "%"+groupName+"%")
}
if keyValue := c.Query("key_value"); keyValue != "" {
// 安全地处理 keyValue，避免越界错误
var likePattern string
if len(keyValue) > 2 {
likePattern = "%" + keyValue[1:len(keyValue)-1] + "%"
} else {
likePattern = "%" + keyValue + "%"
}
db = db.Where("key_value LIKE ?", likePattern)
}
if isSuccessStr := c.Query("is_success"); isSuccessStr != "" {
if isSuccess, err := strconv.ParseBool(isSuccessStr); err == nil {
db = db.Where("is_success = ?", isSuccess)
}
}
if statusCodeStr := c.Query("status_code"); statusCodeStr != "" {
if statusCode, err := strconv.Atoi(statusCodeStr); err == nil {
db = db.Where("status_code = ?", statusCode)
}
}
if sourceIP := c.Query("source_ip"); sourceIP != "" {
db = db.Where("source_ip = ?", sourceIP)
}
if errorContains := c.Query("error_contains"); errorContains != "" {
db = db.Where("error_message LIKE ?", "%"+errorContains+"%")
}
if startTimeStr := c.Query("start_time"); startTimeStr != "" {
if startTime, err := time.Parse(time.RFC3339, startTimeStr); err == nil {
db = db.Where("timestamp >= ?", startTime)
}
}
if endTimeStr := c.Query("end_time"); endTimeStr != "" {
if endTime, err := time.Parse(time.RFC3339, endTimeStr); err == nil {
db = db.Where("timestamp <= ?", endTime)
}
}
return db
}
}
// GetLogsQuery returns a GORM query for fetching logs with filters.
func (s *LogService) GetLogsQuery(c *gin.Context) *gorm.DB {
return s.DB.Model(&models.RequestLog{}).Scopes(logFiltersScope(c))
}
// StreamLogKeysToCSV fetches unique keys from logs based on filters and streams them as a CSV.
func (s *LogService) StreamLogKeysToCSV(c *gin.Context, writer io.Writer) error {
// Create a CSV writer
csvWriter := csv.NewWriter(writer)
defer csvWriter.Flush()
// Write CSV header
header := []string{"key_value", "group_name", "status_code"}
if err := csvWriter.Write(header); err != nil {
return fmt.Errorf("failed to write CSV header: %w", err)
}
// 直接查询所有结果，不使用分批，避免复杂的SQL问题
var results []ExportableLogKey
err := s.DB.Model(&models.RequestLog{}).
Scopes(logFiltersScope(c)).
Select("key_value, group_name, status_code").
Group("key_value, group_name, status_code").
Order("key_value").
Find(&results).Error
if err != nil {
return fmt.Errorf("failed to fetch log keys: %w", err)
}
// 写入CSV数据
for _, record := range results {
csvRecord := []string{
record.KeyValue,
record.GroupName,
strconv.Itoa(record.StatusCode),
}
if err := csvWriter.Write(csvRecord); err != nil {
return fmt.Errorf("failed to write CSV record: %w", err)
}
}
return nil
}