wechat-robot/internal/handler/logto.go

package handler

import (
	"net/http"
	"strings"
	"time"

	"gitee.ltd/lxh/wechat-robot/internal/types"
	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/session"
	logtoClient "github.com/logto-io/go/v2/client"
	"github.com/valyala/fasthttp/fasthttpadaptor"

	"gitee.ltd/lxh/wechat-robot/internal/config"
)

// 实现Logto会话存储接口
var store = session.New(session.Config{
	KeyLookup:      "cookie:logto-session",
	CookieSecure:   false, // 开发环境可设成false
	CookieHTTPOnly: true,
	// 可以配置 CookieSecret
})

// getProtocol
// @description: 获取请求协议
// @param c
// @return string
func getProtocol(c *fiber.Ctx) string {
	protocol := c.Protocol()
	for k, v := range c.GetReqHeaders() {
		if strings.ToLower(k) == "referer" {
			refererProtocol := strings.Split(v[0], ":")[0]
			if refererProtocol != protocol {
				protocol = refererProtocol
			}
		}
	}
	return protocol
}

// LogtoLogin 重定向到Logto登录页面
func LogtoLogin(c *fiber.Ctx) error {
	// 加载配置
	cfg, err := config.Load()
	if err != nil {
		return c.Redirect("/error?error=系统错误，无法加载配置")
	}

	// 构建回调URL

	callbackURL := getProtocol(c) + "://" + c.Hostname() + "/auth/logto/callback"

	// 创建Logto客户端配置
	logtoConfig := &logtoClient.LogtoConfig{
		Endpoint:  cfg.Auth.Logto.Endpoint,
		AppId:     cfg.Auth.Logto.AppId,
		AppSecret: cfg.Auth.Logto.AppSecret,
	}
	// 获取会话
	fiberSessionStorage := &types.LogtoSessionStorage{Store: store, Ctx: c}

	// 创建Logto客户端
	client := logtoClient.NewLogtoClient(logtoConfig, fiberSessionStorage)

	// 创建登录选项，确保包含正确的回调URL
	signInOptions := &logtoClient.SignInOptions{
		RedirectUri: callbackURL,
	}

	// 获取登录链接并重定向
	signInUri, err := client.SignIn(signInOptions)
	if err != nil {
		return c.Redirect("/error?error=Logto登录错误: " + err.Error() + "。请在Logto控制台确认回调URI为: " + callbackURL)
	}
	// 去登录
	return c.Redirect(signInUri)
}

// LogtoCallback 处理Logto登录回调
func LogtoCallback(c *fiber.Ctx) error {
	// 加载配置
	cfg, err := config.Load()
	if err != nil {
		return c.Redirect("/error?error=系统错误，无法加载配置")
	}

	// 创建Logto客户端配置
	logtoConfig := &logtoClient.LogtoConfig{
		Endpoint:  cfg.Auth.Logto.Endpoint,
		AppId:     cfg.Auth.Logto.AppId,
		AppSecret: cfg.Auth.Logto.AppSecret,
	}
	// 获取会话
	fiberSessionStorage := &types.LogtoSessionStorage{Store: store, Ctx: c}
	// 创建Logto客户端
	client := logtoClient.NewLogtoClient(logtoConfig, fiberSessionStorage)

	r := &http.Request{}
	if err = fasthttpadaptor.ConvertRequest(c.Context(), r, true); err != nil {
		return c.Redirect("/error?error=转换请求错误: " + err.Error())
	}

	// 处理回调
	err = client.HandleSignInCallback(r)
	if err != nil {
		return c.Redirect("/error?error=Logto回调处理错误: " + err.Error())
	}

	// 设置auth_token cookie，使用logto前缀
	cookie := new(fiber.Cookie)
	cookie.Name = "auth_token"
	cookie.Value = "logto:auth"                     // 设置auth_token带有logto前缀
	cookie.Expires = time.Now().Add(24 * time.Hour) // 24小时过期
	cookie.HTTPOnly = true
	cookie.Path = "/"
	c.Cookie(cookie)

	// 重定向到机器人列表页面
	return c.Redirect("/admin/robots")
}

// LogtoLogout 处理Logto退出登录
func LogtoLogout(c *fiber.Ctx) error {
	// 加载配置
	cfg, err := config.Load()
	if err != nil {
		return c.Redirect("/error?error=系统错误，无法加载配置")
	}

	// 创建Logto客户端配置
	logtoConfig := &logtoClient.LogtoConfig{
		Endpoint:  cfg.Auth.Logto.Endpoint,
		AppId:     cfg.Auth.Logto.AppId,
		AppSecret: cfg.Auth.Logto.AppSecret,
	}
	// 获取会话
	fiberSessionStorage := &types.LogtoSessionStorage{Store: store, Ctx: c}
	// 创建Logto客户端
	client := logtoClient.NewLogtoClient(logtoConfig, fiberSessionStorage)

	// 构建登出后重定向URL
	postLogoutRedirectURL := getProtocol(c) + "://" + c.Hostname()

	// 获取登出链接
	signOutUri, err := client.SignOut(postLogoutRedirectURL)
	if err != nil {
		return c.Redirect("/error?error=Logto登出错误: " + err.Error())
	}

	// 清除本地Cookie
	cookie := new(fiber.Cookie)
	cookie.Name = "auth_token"
	cookie.Value = ""
	cookie.Expires = time.Now().Add(-time.Hour) // 设置为过期
	cookie.HTTPOnly = true
	cookie.Path = "/"
	c.Cookie(cookie)

	// 重定向到Logto登出页面
	return c.Redirect(signOutUri)
}