heicode-mananger/heicode/controller/heicode_agent_session.go

package controller

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/i18n"
	"github.com/heicode/manager/model"

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

// HeicodeAgentSessionRequest accepts tokens obtained only from Agent identity platform.
// Manager verifies them server-side and issues the browser session cookie (same as password login).
type HeicodeAgentSessionRequest struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
}

type agentMeEnvelope struct {
	Success bool   `json:"success"`
	Message string `json:"message"`
	Detail  string `json:"detail"`
	Data    struct {
		ID        string `json:"id"`
		Email     string `json:"email"`
		Name      string `json:"name"`
		Role      string `json:"role"`
		ChannelID string `json:"channelId"`
		Status    string `json:"status"`
	} `json:"data"`
}

type agentRefreshEnvelope struct {
	Success bool   `json:"success"`
	Detail  string `json:"detail"`
	Message string `json:"message"`
	Data    struct {
		Token        string `json:"token"`
		RefreshToken string `json:"refreshToken"`
	} `json:"data"`
}

func defaultHeicodeAuthBaseURL() string {
	if v := strings.TrimSpace(os.Getenv("HEICODE_AUTH_BASE_URL")); v != "" {
		return strings.TrimRight(v, "/")
	}
	// 与 docs/integration/Heicode-登录接口对接文档.md §2.1 生产 Base URL 一致（前端 VITE 默认同源）。
	return "https://apimtaiji.azure-api.net/api/mcp"
}

func jitUsernameFromEmail(email string) string {
	e := strings.TrimSpace(email)
	if len(e) <= model.UserNameMaxLength {
		return e
	}
	sum := sha256.Sum256([]byte(e))
	return "ag_" + hex.EncodeToString(sum[:])[:16]
}

// parseEmailList returns a set of normalized lowercase emails from an env value.
func parseEmailList(raw string) map[string]struct{} {
	out := map[string]struct{}{}
	for _, e := range strings.Split(raw, ",") {
		e = strings.ToLower(strings.TrimSpace(e))
		if e != "" {
			out[e] = struct{}{}
		}
	}
	return out
}

// roleFromAgentWithEmail decides the local role for a JIT-synced Agent user.
//
// 安全策略：管理员权限只能通过本地配置（环境变量白名单）显式授予，
// **不信任** Agent 平台返回的 role 字段。这样防止外部身份平台
// 的角色被直接映射到 Manager 的高权限角色。
//
//   - 邮箱命中 HEICODE_ROOT_EMAILS  -> RoleRootUser
//   - 邮箱命中 HEICODE_ADMIN_EMAILS -> RoleAdminUser
//   - 其他任何情况                  -> RoleCommonUser（默认普通用户）
//
// 第二参数 `role` 当前未使用，保留是为了未来扩展（例如在策略中允许
// 信任部分上游 role），不破坏调用点签名。
func roleFromAgentWithEmail(_ string, email string) int {
	emailKey := strings.ToLower(strings.TrimSpace(email))
	rootEmails := parseEmailList(os.Getenv("HEICODE_ROOT_EMAILS"))
	if _, ok := rootEmails[emailKey]; ok {
		return common.RoleRootUser
	}
	adminEmails := parseEmailList(os.Getenv("HEICODE_ADMIN_EMAILS"))
	if _, ok := adminEmails[emailKey]; ok {
		return common.RoleAdminUser
	}
	return common.RoleCommonUser
}

func statusFromAgent(status string) int {
	if strings.EqualFold(strings.TrimSpace(status), "active") {
		return common.UserStatusEnabled
	}
	return common.UserStatusDisabled
}

func agentHTTPClient() *http.Client {
	return &http.Client{Timeout: 15 * time.Second}
}

func fetchAgentMe(baseURL, accessToken string) (agentMeEnvelope, int, error) {
	var out agentMeEnvelope
	req, err := http.NewRequest(http.MethodGet, baseURL+"/api/auth/me", nil)
	if err != nil {
		return out, 0, err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	req.Header.Set("X-Request-Id", common.GetUUID())
	res, err := agentHTTPClient().Do(req)
	if err != nil {
		return out, 0, err
	}
	defer res.Body.Close()
	body, err := io.ReadAll(io.LimitReader(res.Body, 1<<20))
	if err != nil {
		return out, res.StatusCode, err
	}
	if err := common.Unmarshal(body, &out); err != nil {
		return out, res.StatusCode, fmt.Errorf("invalid response from Agent /me: %w", err)
	}
	return out, res.StatusCode, nil
}

func fetchAgentRefresh(baseURL, refreshToken string) (access string, refresh string, err error) {
	req, err := http.NewRequest(http.MethodPost, baseURL+"/api/auth/refresh", nil)
	if err != nil {
		return "", "", err
	}
	req.Header.Set("Authorization", "Bearer "+refreshToken)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Request-Id", common.GetUUID())
	res, err := agentHTTPClient().Do(req)
	if err != nil {
		return "", "", err
	}
	defer res.Body.Close()
	body, err := io.ReadAll(io.LimitReader(res.Body, 1<<20))
	if err != nil {
		return "", "", err
	}
	var env agentRefreshEnvelope
	if err := common.Unmarshal(body, &env); err != nil {
		return "", "", fmt.Errorf("invalid response from Agent /refresh: %w", err)
	}
	if !env.Success || env.Data.Token == "" {
		msg := env.Message
		if msg == "" {
			msg = env.Detail
		}
		if msg == "" {
			msg = "refresh failed"
		}
		return "", "", errors.New(msg)
	}
	return env.Data.Token, env.Data.RefreshToken, nil
}

// markBillingProviderNewapi notifies mcp-server that this user's billing
// provider should be set to "newapi" (per Heicode 决策 7.7.1 ②).
//
// Fire-and-forget: 失败不影响登录主流程；mcp-server 那边也允许重复调用
// （内部端点对相同 email 幂等）。详见 Heicode-对接进度与待办.md §7.8.1。
//
// Token 通过 K8s secret / docker compose env `MCP_SERVER_INTERNAL_TOKEN`
// 注入；未配置时跳过（开发环境兼容）。
func markBillingProviderNewapi(email string) {
	tok := strings.TrimSpace(os.Getenv("MCP_SERVER_INTERNAL_TOKEN"))
	if tok == "" {
		return
	}
	email = strings.TrimSpace(email)
	if email == "" {
		return
	}
	go func(email, tok string) {
		body, err := common.Marshal(map[string]string{
			"email":            email,
			"billing_provider": "newapi",
		})
		if err != nil {
			common.SysLog("markBillingProviderNewapi: marshal failed: " + err.Error())
			return
		}
		baseURL := defaultHeicodeAuthBaseURL()
		req, err := http.NewRequest(
			http.MethodPut,
			baseURL+"/api/auth/internal/billing-provider",
			strings.NewReader(string(body)),
		)
		if err != nil {
			common.SysLog("markBillingProviderNewapi: req build failed: " + err.Error())
			return
		}
		req.Header.Set("Authorization", "Bearer "+tok)
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("X-Request-Id", common.GetUUID())
		res, err := agentHTTPClient().Do(req)
		if err != nil {
			common.SysLog("markBillingProviderNewapi: send failed: " + err.Error())
			return
		}
		defer res.Body.Close()
		if res.StatusCode >= 400 {
			b, _ := io.ReadAll(io.LimitReader(res.Body, 1<<14))
			common.SysLog(fmt.Sprintf(
				"markBillingProviderNewapi: %d %s — %s",
				res.StatusCode, email, strings.TrimSpace(string(b)),
			))
		}
	}(email, tok)
}

func syncLocalUserFromAgent(me agentMeEnvelope) (*model.User, error) {
	email := strings.TrimSpace(me.Data.Email)
	if email == "" {
		return nil, errors.New("Agent account has no email")
	}

	var user model.User
	err := model.DB.Where("email = ?", email).First(&user).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		display := strings.TrimSpace(me.Data.Name)
		if display == "" {
			display = strings.Split(email, "@")[0]
		}
		// Local model-access bucket — never seed with the Agent channelId
		// (no abilities row matches a random UUID, so the new user would land
		// with zero models on first /v1/models call). Admins control group
		// from the NewAPI dashboard after JIT-create. Fix companion to
		// 578a68f which only patched the every-login overwrite path.
		group := "default"
		_ = me.Data.ChannelID
		nu := model.User{
			Username:    jitUsernameFromEmail(email),
			Password:    common.GetRandomString(32),
			DisplayName: display,
			Email:       email,
			Role:        roleFromAgentWithEmail(me.Data.Role, email),
			Status:      statusFromAgent(me.Data.Status),
			Group:       group,
		}
		if nu.Status != common.UserStatusEnabled {
			return nil, errors.New("Agent account is not active")
		}
		if err := nu.Insert(0); err != nil {
			// Possible race: duplicate email/username — reload.
			if err2 := model.DB.Where("email = ?", email).First(&user).Error; err2 != nil {
				return nil, err
			}
		} else {
			if err := model.DB.Where("email = ?", email).First(&user).Error; err != nil {
				return nil, err
			}
		}
	} else if err != nil {
		return nil, err
	}

	if user.Status != common.UserStatusEnabled {
		return nil, errors.New("local account is disabled")
	}

	changed := false
	if name := strings.TrimSpace(me.Data.Name); name != "" && user.DisplayName != name {
		user.DisplayName = name
		changed = true
	}
	// Don't overwrite the existing user's group with the Agent channelId on
	// every login: NewAPI's `users.group` is the **local model-access bucket**
	// (must match a row in the `abilities` / `channels` group column to expose
	// any models). The Agent channelId is a cross-platform identity that
	// rarely matches a NewAPI-side group, so overwriting strands the user
	// with zero models. mcp-server side already tracks channelId separately
	// (see markBillingProviderNewapi), so we don't need it duplicated here.
	//
	// Only seed the group on the JIT-create path above (when the user record
	// is new and has no admin-set group yet). After that, NewAPI admins own
	// the group via the dashboard.
	_ = me.Data.ChannelID
	// Promote role from Agent / email whitelist on every login (never demote).
	desiredRole := roleFromAgentWithEmail(me.Data.Role, email)
	if desiredRole > user.Role {
		user.Role = desiredRole
		changed = true
	}
	if changed {
		if err := user.Update(false); err != nil {
			return nil, err
		}
	}

	// Notify mcp-server billing_provider=newapi (fire-and-forget; goroutine).
	markBillingProviderNewapi(user.Email)

	return &user, nil
}

// HeicodeAgentSessionLogin establishes Manager session after Agent identity verified via token(s).
func HeicodeAgentSessionLogin(c *gin.Context) {
	var req HeicodeAgentSessionRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid params"})
		return
	}
	access := strings.TrimSpace(req.AccessToken)
	refresh := strings.TrimSpace(req.RefreshToken)
	if access == "" {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": "access_token required"})
		return
	}

	baseURL := defaultHeicodeAuthBaseURL()
	me, status, err := fetchAgentMe(baseURL, access)
	newAccess := ""
	newRefresh := ""

	if (err != nil || status == http.StatusUnauthorized || !me.Success) && refresh != "" {
		na, nr, refErr := fetchAgentRefresh(baseURL, refresh)
		if refErr != nil {
			c.JSON(http.StatusOK, gin.H{"success": false, "message": refErr.Error()})
			return
		}
		newAccess = na
		newRefresh = nr
		me, _, err = fetchAgentMe(baseURL, newAccess)
	}

	if err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
		return
	}
	if !me.Success {
		msg := me.Message
		if msg == "" {
			msg = me.Detail
		}
		if msg == "" {
			msg = "unable to verify identity with Agent"
		}
		c.JSON(http.StatusOK, gin.H{"success": false, "message": msg})
		return
	}

	user, err := syncLocalUserFromAgent(me)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()})
		return
	}

	if model.IsTwoFAEnabled(user.Id) {
		session := sessions.Default(c)
		session.Set("pending_username", user.Username)
		session.Set("pending_user_id", user.Id)
		if err := session.Save(); err != nil {
			common.ApiErrorI18n(c, i18n.MsgUserSessionSaveFailed)
			return
		}
		c.JSON(http.StatusOK, gin.H{
			"message": i18n.T(c, i18n.MsgUserRequire2FA),
			"success": true,
			"data": map[string]interface{}{
				"require_2fa": true,
			},
		})
		return
	}

	model.UpdateUserLastLoginAt(user.Id)
	session := sessions.Default(c)
	session.Set("id", user.Id)
	session.Set("username", user.Username)
	session.Set("role", user.Role)
	session.Set("status", user.Status)
	session.Set("group", user.Group)
	if err := session.Save(); err != nil {
		common.ApiErrorI18n(c, i18n.MsgUserSessionSaveFailed)
		return
	}

	model.EnsureUserRelayToken(user.Id, user.Username)

	data := gin.H{
		"id":           user.Id,
		"username":     user.Username,
		"display_name": user.DisplayName,
		"role":         user.Role,
		"status":       user.Status,
		"group":        user.Group,
	}
	if newAccess != "" {
		data["heicode_token"] = newAccess
		if newRefresh != "" {
			data["heicode_refresh_token"] = newRefresh
		}
	}

	c.JSON(http.StatusOK, gin.H{
		"message": "",
		"success": true,
		"data":    data,
	})
}