package controller import ( "errors" "fmt" "html" "net/http" "net/url" "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" "gorm.io/gorm" ) // HeiCode-specific browser-login bridge. // // Flow: // // 1. HeiCode desktop opens // GET /heicode/oauth/authorize?state=&redirect_uri=&provider_id= // 2. If the user is not signed in to new-api, we render a small HTML page // that links to /login and polls /heicode/oauth/session every 1.5s // until the user has authenticated. Then it auto-redirects back. // 3. Once the session is valid, we find or create a token named "HeiCode" // for the current user and 302-redirect back to redirect_uri with the // token in the query string: // ?state=&token=sk-XXXX&provider_id= // // The redirect_uri is restricted to loopback addresses for safety. const heicodeTokenName = "HeiCode" func isLoopbackHost(host string) bool { host = strings.ToLower(strings.TrimSpace(host)) return host == "127.0.0.1" || host == "localhost" || host == "::1" } // HeicodeOAuthAuthorize handles GET /heicode/oauth/authorize. func HeicodeOAuthAuthorize(c *gin.Context) { state := strings.TrimSpace(c.Query("state")) redirectURI := strings.TrimSpace(c.Query("redirect_uri")) providerID := strings.TrimSpace(c.Query("provider_id")) if state == "" || redirectURI == "" { c.String(http.StatusBadRequest, "missing state or redirect_uri") return } parsed, err := url.Parse(redirectURI) if err != nil || parsed.Scheme == "" || parsed.Host == "" { c.String(http.StatusBadRequest, "invalid redirect_uri") return } if !isLoopbackHost(parsed.Hostname()) { c.String(http.StatusBadRequest, "redirect_uri must be a loopback address") return } session := sessions.Default(c) rawID := session.Get("id") if rawID == nil { renderHeicodeLoginRequired(c, state, redirectURI, providerID) return } userID, ok := rawID.(int) if !ok || userID <= 0 { renderHeicodeLoginRequired(c, state, redirectURI, providerID) return } tokenKey, err := getOrCreateHeicodeToken(userID) if err != nil { c.String(http.StatusInternalServerError, "failed to issue HeiCode token: "+err.Error()) return } q := parsed.Query() q.Set("state", state) q.Set("token", "sk-"+tokenKey) if providerID != "" { q.Set("provider_id", providerID) } parsed.RawQuery = q.Encode() c.Redirect(http.StatusFound, parsed.String()) } // HeicodeOAuthSession handles GET /heicode/oauth/session — used by the // "please log in" fallback page to poll session readiness from JS. func HeicodeOAuthSession(c *gin.Context) { session := sessions.Default(c) id := session.Get("id") loggedIn := false var userID int if v, ok := id.(int); ok && v > 0 { loggedIn = true userID = v } c.JSON(http.StatusOK, gin.H{ "logged_in": loggedIn, "user_id": userID, }) } func getOrCreateHeicodeToken(userID int) (string, error) { var existing model.Token err := model.DB.Where("user_id = ? AND name = ?", userID, heicodeTokenName). Order("id desc"). First(&existing).Error if err == nil { if existing.Status != common.TokenStatusEnabled { existing.Status = common.TokenStatusEnabled if updateErr := existing.Update(); updateErr != nil { return "", updateErr } } return existing.GetFullKey(), nil } if !errors.Is(err, gorm.ErrRecordNotFound) { return "", err } rawKey, err := common.GenerateKey() if err != nil { return "", err } now := common.GetTimestamp() tok := model.Token{ UserId: userID, Name: heicodeTokenName, Key: rawKey, Status: common.TokenStatusEnabled, CreatedTime: now, AccessedTime: now, ExpiredTime: -1, UnlimitedQuota: true, } if err := tok.Insert(); err != nil { return "", err } return tok.GetFullKey(), nil } func renderHeicodeLoginRequired(c *gin.Context, state, redirectURI, providerID string) { authorizeURL := fmt.Sprintf( "/heicode/oauth/authorize?state=%s&redirect_uri=%s&provider_id=%s", url.QueryEscape(state), url.QueryEscape(redirectURI), url.QueryEscape(providerID), ) loginURL := "/login" body := fmt.Sprintf(` HeiCode 登录授权

请先登录新 API 控制台

HeiCode 客户端正在等待你完成授权。请在新打开的页面中登录到 %s,登录成功后本页会自动跳转回 HeiCode。

打开登录页

如果浏览器没有自动跳转,请手动在登录后回到此页面,或刷新本页。

提供方:%s

`, html.EscapeString(providerID), html.EscapeString(loginURL), html.EscapeString(providerID), authorizeURL, ) c.Header("Content-Type", "text/html; charset=utf-8") c.Status(http.StatusOK) _, _ = c.Writer.WriteString(body) }