heicode-win/heicode/controller/heicode_oauth.go

package controller

import (
	"errors"
	"fmt"
	"html"
	"net/http"
	"net/url"
	"strings"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/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=<S>&redirect_uri=<U>&provider_id=<id>
//	2. If the user is not signed in to heicode, 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:
//	     <redirect_uri>?state=<S>&token=sk-XXXX&provider_id=<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 := "/sign-in?redirect=" + url.QueryEscape(authorizeURL)

	body := fmt.Sprintf(`<!doctype html>
<html lang="zh-CN">
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width,initial-scale=1" />
<title>Heicode 登录授权</title>
<style>
  :root { color-scheme: light dark; }
  html,body { margin:0; padding:0; }
  body { font-family: system-ui, -apple-system, "PingFang SC", sans-serif; min-height:100vh; display:flex; align-items:center; justify-content:center; background:#f7f7f8; color:#111; }
  .card { background:#fff; padding:32px 36px; border-radius:14px; box-shadow:0 8px 24px rgba(0,0,0,.08); max-width:440px; width:100%%; box-sizing:border-box; }
  h1 { font-size:18px; margin:0 0 8px; }
  p { color:#555; line-height:1.6; margin:8px 0; font-size:14px; }
  a.btn { display:inline-block; margin-top:16px; padding:10px 18px; background:#111; color:#fff; border-radius:8px; text-decoration:none; font-size:14px; }
  a.btn:hover { background:#000; }
  .muted { color:#888; font-size:12px; margin-top:18px; }
  code { background:#f1f1f3; padding:2px 6px; border-radius:4px; font-size:12px; }
</style>
</head>
<body>
<div class="card">
  <h1>请先登录 Heicode 控制台</h1>
  <p>Heicode 桌面端正在等待你完成授权。请在新打开的页面中登录到 <strong>%s</strong>，登录成功后本页会自动跳转回客户端。</p>
  <a id="goLogin" class="btn" href="%s" target="_blank" rel="noopener">打开登录页</a>
  <p class="muted">如果浏览器没有自动跳转，请手动在登录后回到此页面，或刷新本页。</p>
  <p class="muted">租户：<code>%s</code></p>
</div>
<script>
(function(){
  var authorizeUrl = %q;
  var pollMs = 1500;
  function check() {
    fetch('/heicode/oauth/session', { credentials: 'include' })
      .then(function (r) { return r.json(); })
      .then(function (data) {
        if (data && data.logged_in) {
          window.location.replace(authorizeUrl);
        } else {
          setTimeout(check, pollMs);
        }
      })
      .catch(function () { setTimeout(check, pollMs); });
  }
  setTimeout(check, pollMs);
})();
</script>
</body>
</html>`,
		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)
}