heicode-mananger/heicode/middleware/auth.go

package middleware

import (
	"errors"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/constant"
	"github.com/heicode/manager/i18n"
	"github.com/heicode/manager/logger"
	"github.com/heicode/manager/model"
	"github.com/heicode/manager/service"
	"github.com/heicode/manager/setting/operation_setting"
	"github.com/heicode/manager/setting/ratio_setting"
	"github.com/heicode/manager/types"

	"github.com/gin-contrib/sessions"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

func validUserInfo(username string, role int) bool {
	// check username is empty
	if strings.TrimSpace(username) == "" {
		return false
	}
	if !common.IsValidateRole(role) {
		return false
	}
	return true
}

func authHelper(c *gin.Context, minRole int) {
	session := sessions.Default(c)
	username := session.Get("username")
	role := session.Get("role")
	id := session.Get("id")
	status := session.Get("status")
	useAccessToken := false
	if username == nil {
		// Check access token
		accessToken := c.Request.Header.Get("Authorization")
		if accessToken == "" {
			c.JSON(http.StatusUnauthorized, gin.H{
				"success": false,
				"message": common.TranslateMessage(c, i18n.MsgAuthNotLoggedIn),
			})
			c.Abort()
			return
		}
		user, authErr := model.ValidateAccessToken(accessToken)
		if authErr != nil {
			if errors.Is(authErr, model.ErrDatabase) {
				common.SysLog("ValidateAccessToken database error: " + authErr.Error())
				c.JSON(http.StatusInternalServerError, gin.H{
					"success": false,
					"message": common.TranslateMessage(c, i18n.MsgDatabaseError),
				})
			} else {
				c.JSON(http.StatusOK, gin.H{
					"success": false,
					"message": common.TranslateMessage(c, i18n.MsgAuthAccessTokenInvalid),
				})
			}
			c.Abort()
			return
		}
		if user != nil && user.Username != "" {
			if !validUserInfo(user.Username, user.Role) {
				c.JSON(http.StatusOK, gin.H{
					"success": false,
					"message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid),
				})
				c.Abort()
				return
			}
			// Token is valid
			username = user.Username
			role = user.Role
			id = user.Id
			status = user.Status
			useAccessToken = true
		} else {
			c.JSON(http.StatusOK, gin.H{
				"success": false,
				"message": common.TranslateMessage(c, i18n.MsgAuthAccessTokenInvalid),
			})
			c.Abort()
			return
		}
	}
	// get header New-Api-User
	apiUserIdStr := c.Request.Header.Get("New-Api-User")
	if apiUserIdStr == "" {
		c.JSON(http.StatusUnauthorized, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserIdNotProvided),
		})
		c.Abort()
		return
	}
	apiUserId, err := strconv.Atoi(apiUserIdStr)
	if err != nil {
		c.JSON(http.StatusUnauthorized, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserIdFormatError),
		})
		c.Abort()
		return

	}
	idVal, _ := id.(int)
	if idVal != apiUserId {
		c.JSON(http.StatusUnauthorized, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserIdMismatch),
		})
		c.Abort()
		return
	}
	statusVal, ok := status.(int)
	if !ok {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid),
		})
		c.Abort()
		return
	}
	if statusVal == common.UserStatusDisabled {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserBanned),
		})
		c.Abort()
		return
	}
	roleVal, ok := role.(int)
	if !ok {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid),
		})
		c.Abort()
		return
	}
	if roleVal < minRole {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthInsufficientPrivilege),
		})
		c.Abort()
		return
	}
	usernameStr, ok := username.(string)
	if !ok {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid),
		})
		c.Abort()
		return
	}
	if !validUserInfo(usernameStr, roleVal) {
		c.JSON(http.StatusOK, gin.H{
			"success": false,
			"message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid),
		})
		c.Abort()
		return
	}
	// 防止不同newapi版本冲突，导致数据不通用
	c.Header("Auth-Version", "864b7076dbcd0a3c01b5520316720ebf")
	c.Set("username", usernameStr)
	c.Set("role", roleVal)
	c.Set("id", id)
	c.Set("group", session.Get("group"))
	c.Set("user_group", session.Get("group"))
	c.Set("use_access_token", useAccessToken)

	c.Next()
}

func TryUserAuth() func(c *gin.Context) {
	return func(c *gin.Context) {
		session := sessions.Default(c)
		id := session.Get("id")
		if id != nil {
			c.Set("id", id)
		}
		c.Next()
	}
}

func UserAuth() func(c *gin.Context) {
	return func(c *gin.Context) {
		authHelper(c, common.RoleCommonUser)
	}
}

func AdminAuth() func(c *gin.Context) {
	return func(c *gin.Context) {
		authHelper(c, common.RoleAdminUser)
	}
}

func RootAuth() func(c *gin.Context) {
	return func(c *gin.Context) {
		authHelper(c, common.RoleRootUser)
	}
}

func WssAuth(c *gin.Context) {

}

// TokenOrUserAuth allows either session-based user auth or API token auth.
// Used for endpoints that need to be accessible from both the dashboard and API clients.
func TokenOrUserAuth() func(c *gin.Context) {
	return func(c *gin.Context) {
		// Try session auth first (dashboard users)
		session := sessions.Default(c)
		if id := session.Get("id"); id != nil {
			if status, ok := session.Get("status").(int); ok && status == common.UserStatusEnabled {
				c.Set("id", id)
				c.Next()
				return
			}
		}
		// Fall back to token auth (API clients)
		TokenAuth()(c)
	}
}

// TokenAuthReadOnly 宽松版本的令牌认证中间件，用于只读查询接口。
// 只验证令牌 key 是否存在，不检查令牌状态、过期时间和额度。
// 即使令牌已过期、已耗尽或已禁用，也允许访问。
// 仍然检查用户是否被封禁。
func TokenAuthReadOnly() func(c *gin.Context) {
	return func(c *gin.Context) {
		key := c.Request.Header.Get("Authorization")
		if key == "" {
			c.JSON(http.StatusUnauthorized, gin.H{
				"success": false,
				"message": common.TranslateMessage(c, i18n.MsgTokenNotProvided),
			})
			c.Abort()
			return
		}
		if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
			key = strings.TrimSpace(key[7:])
		}
		key = strings.TrimPrefix(key, "sk-")
		parts := strings.Split(key, "-")
		key = parts[0]

		token, err := model.GetTokenByKey(key, false)
		if err != nil {
			if errors.Is(err, gorm.ErrRecordNotFound) {
				c.JSON(http.StatusUnauthorized, gin.H{
					"success": false,
					"message": common.TranslateMessage(c, i18n.MsgTokenInvalid),
				})
			} else {
				common.SysLog("TokenAuthReadOnly GetTokenByKey database error: " + err.Error())
				c.JSON(http.StatusInternalServerError, gin.H{
					"success": false,
					"message": common.TranslateMessage(c, i18n.MsgDatabaseError),
				})
			}
			c.Abort()
			return
		}

		userCache, err := model.GetUserCache(token.UserId)
		if err != nil {
			common.SysLog(fmt.Sprintf("TokenAuthReadOnly GetUserCache error for user %d: %v", token.UserId, err))
			c.JSON(http.StatusInternalServerError, gin.H{
				"success": false,
				"message": common.TranslateMessage(c, i18n.MsgDatabaseError),
			})
			c.Abort()
			return
		}
		if userCache.Status != common.UserStatusEnabled {
			c.JSON(http.StatusForbidden, gin.H{
				"success": false,
				"message": common.TranslateMessage(c, i18n.MsgAuthUserBanned),
			})
			c.Abort()
			return
		}

		c.Set("id", token.UserId)
		c.Set("token_id", token.Id)
		c.Set("token_key", token.Key)
		c.Next()
	}
}

func TokenAuth() func(c *gin.Context) {
	return func(c *gin.Context) {
		// ----- V2 path (Heicode device-signed + encrypted body) -----
		// When the client speaks the V2 protocol (Content-Encoding:
		// heicode-aead-v1), there is no bearer at all. We resolve the
		// token via the X-Heicode-Device-Id header, decrypt the body,
		// verify the Ed25519 signature, and shortcut the rest of the
		// legacy logic. See plan in ~/.claude/plans/peaceful-sprouting-crane.md.
		if IsV2Request(c) {
			if err := DecryptV2RequestBody(c); err != nil {
				common.SysLog("V2 body decrypt failed: " + err.Error() +
					" path=" + c.Request.URL.Path + " ip=" + c.ClientIP())
				// Surface a machine-readable error code + server time so
				// the desktop client can render an actionable message
				// (e.g. "your system clock is X seconds off") instead of
				// the generic "Token invalid" string. See plan v2 §
				// "auth error UX" and HeaderAuthError contract in
				// middleware/device_signature.go.
				SetV2AuthDiagnosticHeaders(c, err)
				abortWithOpenAiMessage(c, http.StatusUnauthorized,
					common.TranslateMessage(c, i18n.MsgTokenInvalid))
				return
			}
			token, err := VerifyV2DeviceSignedRequest(c)
			if err != nil {
				common.SysLog("V2 device-signed auth rejected: " + err.Error() +
					" device_id=" + c.GetHeader(HeaderDeviceID) +
					" ip=" + c.ClientIP())
				SetV2AuthDiagnosticHeaders(c, err)
				abortWithOpenAiMessage(c, http.StatusUnauthorized,
					common.TranslateMessage(c, i18n.MsgTokenInvalid))
				return
			}

			// Same downstream wiring as legacy: load user cache, group
			// resolution, IP allowlist. Body has been replaced with
			// plaintext so the relay handler sees normal JSON.
			if !applyTokenPolicyAndContext(c, token) {
				return // applyTokenPolicyAndContext already aborted
			}
			// Critical: V2 path must populate the same per-token context
			// keys the legacy path does — token_key, token_name,
			// token_quota, model_limits, ContextKeyTokenGroup, etc. —
			// otherwise downstream Distribute / billing logic blow up
			// with "record not found". This bug surfaced as HTTP 403
			// new_api_error on every V2 chat call after device pairing
			// finally started succeeding.
			if err := SetupContextForToken(c, token); err != nil {
				return
			}
			c.Next()
			return
		}

		// ----- Legacy path (bare sk- bearer, optional signature) -----
		// 先检测是否为ws
		if c.Request.Header.Get("Sec-WebSocket-Protocol") != "" {
			// Sec-WebSocket-Protocol: realtime, openai-insecure-api-key.sk-xxx, openai-beta.realtime-v1
			// read sk from Sec-WebSocket-Protocol
			key := c.Request.Header.Get("Sec-WebSocket-Protocol")
			parts := strings.Split(key, ",")
			for _, part := range parts {
				part = strings.TrimSpace(part)
				if strings.HasPrefix(part, "openai-insecure-api-key") {
					key = strings.TrimPrefix(part, "openai-insecure-api-key.")
					break
				}
			}
			c.Request.Header.Set("Authorization", "Bearer "+key)
		}
		// 检查path包含/v1/messages 或 /v1/models
		if strings.Contains(c.Request.URL.Path, "/v1/messages") || strings.Contains(c.Request.URL.Path, "/v1/models") {
			anthropicKey := c.Request.Header.Get("x-api-key")
			if anthropicKey != "" {
				c.Request.Header.Set("Authorization", "Bearer "+anthropicKey)
			}
		}
		// gemini api 从query中获取key
		if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models") ||
			strings.HasPrefix(c.Request.URL.Path, "/v1beta/openai/models") ||
			strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
			skKey := c.Query("key")
			if skKey != "" {
				c.Request.Header.Set("Authorization", "Bearer "+skKey)
			}
			// 从x-goog-api-key header中获取key
			xGoogKey := c.Request.Header.Get("x-goog-api-key")
			if xGoogKey != "" {
				c.Request.Header.Set("Authorization", "Bearer "+xGoogKey)
			}
		}
		key := c.Request.Header.Get("Authorization")
		parts := make([]string, 0)
		if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
			key = strings.TrimSpace(key[7:])
		}
		if key == "" || key == "midjourney-proxy" {
			key = c.Request.Header.Get("mj-api-secret")
			if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
				key = strings.TrimSpace(key[7:])
			}
			key = strings.TrimPrefix(key, "sk-")
			parts = strings.Split(key, "-")
			key = parts[0]
		} else {
			key = strings.TrimPrefix(key, "sk-")
			parts = strings.Split(key, "-")
			key = parts[0]
		}
		token, err := model.ValidateUserToken(key)
		if token != nil {
			id := c.GetInt("id")
			if id == 0 {
				c.Set("id", token.UserId)
			}
		}
		if err != nil {
			if errors.Is(err, model.ErrDatabase) {
				common.SysLog("TokenAuth ValidateUserToken database error: " + err.Error())
				abortWithOpenAiMessage(c, http.StatusInternalServerError,
					common.TranslateMessage(c, i18n.MsgDatabaseError))
			} else {
				abortWithOpenAiMessage(c, http.StatusUnauthorized,
					common.TranslateMessage(c, i18n.MsgTokenInvalid))
			}
			return
		}

		// 30-day legacy deadline: after the configured cutoff, bare-bearer
		// sk- callers on /v1/* are rejected entirely. /api/* management
		// endpoints are exempt (admin needs to keep working through the
		// transition). See plan, phase P3.
		if deadline := operation_setting.GetLegacySkV1DeadlineMs(); deadline > 0 {
			if strings.HasPrefix(c.Request.URL.Path, "/v1/") &&
				time.Now().UnixMilli() > deadline {
				common.SysLog("legacy sk on /v1/* rejected (past deadline) token_id=" +
					fmt.Sprint(token.Id) + " path=" + c.Request.URL.Path)
				abortWithOpenAiMessage(c, http.StatusUnauthorized,
					common.TranslateMessage(c, i18n.MsgTokenInvalid))
				return
			}
		}

		// Device-binding signature check. Runs BEFORE the IP-allowlist and
		// user-status checks because if the signature is wrong we don't
		// need to consult the rest of the policy stack — the bearer alone
		// no longer proves identity. See middleware/device_signature.go.
		// Legacy bare-bearer tokens (no device_pubkey set) pass through
		// transparently until P3 enforcement flips on.
		if sigErr := VerifyDeviceSignatureIfRequired(c, token); sigErr != nil {
			// Use 401 for ALL signature failures so we don't leak the
			// distinction "no signature" vs "bad signature" vs "replayed"
			// to outside observers via different status codes. The
			// SysLog records the actual cause for ops debugging.
			common.SysLog("device-signature reject: " + sigErr.Error() +
				" token_id=" + fmt.Sprint(token.Id) +
				" client_ip=" + c.ClientIP())
			abortWithOpenAiMessage(c, http.StatusUnauthorized,
				common.TranslateMessage(c, i18n.MsgTokenInvalid))
			return
		}

		if !applyTokenPolicyAndContext(c, token) {
			return
		}
		if err := SetupContextForToken(c, token, parts...); err != nil {
			return
		}
		c.Next()
	}
}

// applyTokenPolicyAndContext runs the per-request policy checks that
// previously lived inline in TokenAuth: IP allowlist, user enabled,
// token group authorization. Shared between V1 (bare bearer) and V2
// (device-signed) auth paths so they apply the same downstream gating.
//
// Returns false if any check failed AND aborted the request — caller
// should `return` immediately. Returns true to continue.
func applyTokenPolicyAndContext(c *gin.Context, token *model.Token) bool {
	allowIps := token.GetIpLimits()
	if len(allowIps) > 0 {
		clientIp := c.ClientIP()
		logger.LogDebug(c, "Token has IP restrictions, checking client IP %s", clientIp)
		ip := net.ParseIP(clientIp)
		if ip == nil {
			abortWithOpenAiMessage(c, http.StatusForbidden, "无法解析客户端 IP 地址")
			return false
		}
		if !common.IsIpInCIDRList(ip, allowIps) {
			abortWithOpenAiMessage(c, http.StatusForbidden,
				"您的 IP 不在令牌允许访问的列表中", types.ErrorCodeAccessDenied)
			return false
		}
		logger.LogDebug(c, "Client IP %s passed the token IP restrictions check", clientIp)
	}

	userCache, err := model.GetUserCache(token.UserId)
	if err != nil {
		common.SysLog(fmt.Sprintf("applyTokenPolicy GetUserCache error for user %d: %v", token.UserId, err))
		abortWithOpenAiMessage(c, http.StatusInternalServerError,
			common.TranslateMessage(c, i18n.MsgDatabaseError))
		return false
	}
	if userCache.Status != common.UserStatusEnabled {
		abortWithOpenAiMessage(c, http.StatusForbidden,
			common.TranslateMessage(c, i18n.MsgAuthUserBanned))
		return false
	}
	userCache.WriteContext(c)

	userGroup := userCache.Group
	tokenGroup := token.Group
	if tokenGroup != "" {
		if _, ok := service.GetUserUsableGroups(userGroup)[tokenGroup]; !ok {
			abortWithOpenAiMessage(c, http.StatusForbidden,
				fmt.Sprintf("无权访问 %s 分组", tokenGroup))
			return false
		}
		if !ratio_setting.ContainsGroupRatio(tokenGroup) {
			if tokenGroup != "auto" {
				abortWithOpenAiMessage(c, http.StatusForbidden,
					fmt.Sprintf("分组 %s 已被弃用", tokenGroup))
				return false
			}
		}
		userGroup = tokenGroup
	}
	common.SetContextKey(c, constant.ContextKeyUsingGroup, userGroup)
	return true
}

func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) error {
	if token == nil {
		return fmt.Errorf("token is nil")
	}
	c.Set("id", token.UserId)
	c.Set("token_id", token.Id)
	c.Set("token_key", token.Key)
	c.Set("token_name", token.Name)
	c.Set("token_unlimited_quota", token.UnlimitedQuota)
	if !token.UnlimitedQuota {
		c.Set("token_quota", token.RemainQuota)
	}
	if token.ModelLimitsEnabled {
		c.Set("token_model_limit_enabled", true)
		c.Set("token_model_limit", token.GetModelLimitsMap())
	} else {
		c.Set("token_model_limit_enabled", false)
	}
	common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group)
	common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry)
	if len(parts) > 1 {
		if model.IsAdmin(token.UserId) {
			c.Set("specific_channel_id", parts[1])
		} else {
			c.Header("specific_channel_version", "701e3ae1dc3f7975556d354e0675168d004891c8")
			abortWithOpenAiMessage(c, http.StatusForbidden, "普通用户不支持指定渠道")
			return fmt.Errorf("普通用户不支持指定渠道")
		}
	}
	return nil
}