FranklinWH-Toolbox/service/Untitled-1

// main.go
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gofiber/fiber/v2"
	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
)

var (
	ctx = context.Background()
	rdb *redis.Client

	// 从 env 加载
	clientID     = os.Getenv("CASDOOR_CLIENT_ID")
	clientSecret = os.Getenv("CASDOOR_CLIENT_SECRET")
	authURL      = os.Getenv("CASDOOR_AUTHORIZE_URL") // e.g. https://door.example.com/login/oauth/authorize
	tokenURL     = os.Getenv("CASDOOR_TOKEN_URL")     // e.g. https://door.example.com/api/login/oauth/access_token
	redirectURI  = os.Getenv("APP_REDIRECT_URI")      // e.g. https://api.example.com/api/auth/callback
	frontendURL  = os.Getenv("FRONTEND_URL")          // e.g. https://app.example.com
	// session cookie name
	sessionCookieName = "sid"
	// session TTL (sync with refresh token validity; e.g. 7 days)
	sessionTTL = 7 * 24 * time.Hour
)

type OAuthTokenResp struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	ExpiresIn    int    `json:"expires_in"` // 秒
	// 可能还有其他字段，按实际返回解析
}

/* ------------------- Redis helpers ------------------- */
func initRedis() {
	rdb = redis.NewClient(&redis.Options{
		Addr:     "127.0.0.1:6379", // env 改造
		Password: "",
		DB:       0,
	})
	if _, err := rdb.Ping(ctx).Result(); err != nil {
		log.Fatalf("redis connect error: %v", err)
	}
}

func storeSession(sid string, token OAuthTokenResp) error {
	key := fmt.Sprintf("session:%s", sid)
	exp := time.Now().Add(sessionTTL)
	accessExp := time.Now().Add(time.Duration(token.ExpiresIn) * time.Second).Unix()
	fields := map[string]interface{}{
		"refresh_token": token.RefreshToken,
		"access_token":  token.AccessToken,
		"access_exp":    accessExp,
	}
	if err := rdb.HSet(ctx, key, fields).Err(); err != nil {
		return err
	}
	return rdb.ExpireAt(ctx, key, exp).Err()
}

func getSession(sid string) (map[string]string, error) {
	key := fmt.Sprintf("session:%s", sid)
	res, err := rdb.HGetAll(ctx, key).Result()
	if err != nil {
		return nil, err
	}
	if len(res) == 0 {
		return nil, errors.New("no session")
	}
	return res, nil
}

func deleteSession(sid string) error {
	key := fmt.Sprintf("session:%s", sid)
	return rdb.Del(ctx, key).Err()
}

/* ------------------- Casdoor interaction helpers ------------------- */

// 标准 OAuth2 token exchange
func exchangeCodeForToken(code string) (OAuthTokenResp, error) {
	form := map[string]string{
		"grant_type":    "authorization_code",
		"code":          code,
		"redirect_uri":  redirectURI,
		"client_id":     clientID,
		"client_secret": clientSecret,
	}
	return postTokenForm(form)
}

func refreshWithRefreshToken(refreshToken string) (OAuthTokenResp, error) {
	form := map[string]string{
		"grant_type":    "refresh_token",
		"refresh_token": refreshToken,
		"client_id":     clientID,
		"client_secret": clientSecret,
	}
	return postTokenForm(form)
}

func postTokenForm(form map[string]string) (OAuthTokenResp, error) {
	bodyBytes, _ := json.Marshal(form)
	req, _ := http.NewRequest("POST", tokenURL, bytes.NewBuffer(bodyBytes))
	req.Header.Set("Content-Type", "application/json") // 某些 Casdoor 期望 application/x-www-form-urlencoded：根据 Casdoor 实际调整
	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return OAuthTokenResp{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		return OAuthTokenResp{}, fmt.Errorf("token endpoint returned %d", resp.StatusCode)
	}
	var tr OAuthTokenResp
	if err := json.NewDecoder(resp.Body).Decode(&tr); err != nil {
		return OAuthTokenResp{}, err
	}
	return tr, nil
}

/* ------------------- Handlers ------------------- */

// GET /api/auth/login -> 重定向到 Casdoor authorize URL
func loginRedirect(c *fiber.Ctx) error {
	// state 可用于防 CSRF
	state := uuid.NewString()
	// 保存 state -> 可以放到 Redis 验证，不过此处省略
	u := fmt.Sprintf("%s?client_id=%s&response_type=code&redirect_uri=%s&state=%s", authURL, clientID, redirectURI, state)
	return c.Redirect(u, fiber.StatusFound)
}

// GET /api/auth/callback?code=...&state=...
func callbackHandler(c *fiber.Ctx) error {
	code := c.Query("code")
	if code == "" {
		return c.Status(400).SendString("missing code")
	}

	// 用 code 换 token
	tokenResp, err := exchangeCodeForToken(code)
	if err != nil {
		log.Println("exchange code err:", err)
		return c.Status(500).JSON(fiber.Map{"error": "token exchange failed"})
	}

	// 生成 sid 并存 Redis（只存 refresh_token，access_token 可以缓存）
	sid := uuid.NewString()
	if err := storeSession(sid, tokenResp); err != nil {
		log.Println("store session err:", err)
	}

	// 设置 HttpOnly cookie 保存 sid
	c.Cookie(&fiber.Cookie{
		Name:     sessionCookieName,
		Value:    sid,
		Expires:  time.Now().Add(sessionTTL),
		HTTPOnly: true,
		Secure:   false, // production: true
		SameSite: "Lax",
		Path:     "/",
	})

	// 最后重定向回前端单页 app
	return c.Redirect(frontendURL, fiber.StatusFound)
}

// GET /api/auth/session  -> 前端载入时调用，返回 access_token（若存在）
func sessionHandler(c *fiber.Ctx) error {
	sid := c.Cookies(sessionCookieName)
	if sid == "" {
		return c.Status(200).JSON(fiber.Map{"authenticated": false})
	}
	sess, err := getSession(sid)
	if err != nil {
		// session 不存在或已过期
		// 清除 cookie
		c.Cookie(&fiber.Cookie{
			Name:     sessionCookieName,
			Value:    "",
			Expires:  time.Now().Add(-time.Hour),
			HTTPOnly: true,
			Secure:   false,
			Path:     "/",
		})
		return c.Status(200).JSON(fiber.Map{"authenticated": false})
	}
	// 若 access_token 仍有效则直接返回
	access := sess["access_token"]
	accessExpStr := sess["access_exp"]
	if access != "" && accessExpStr != "" {
		accessExp, _ := strconv.ParseInt(accessExpStr, 10, 64)
		if time.Now().Unix() < accessExp {
			return c.JSON(fiber.Map{"authenticated": true, "access_token": access})
		}
	}

	// access 已过期 -> 试着用 refresh token 刷新（后端直接刷新并返回）
	refreshToken := sess["refresh_token"]
	if refreshToken == "" {
		return c.Status(200).JSON(fiber.Map{"authenticated": false})
	}
	newTok, err := refreshWithRefreshToken(refreshToken)
	if err != nil {
		// refresh 失败 -> session 失效
		_ = deleteSession(sid)
		c.Cookie(&fiber.Cookie{
			Name:     sessionCookieName,
			Value:    "",
			Expires:  time.Now().Add(-time.Hour),
			HTTPOnly: true,
			Secure:   false,
			Path:     "/",
		})
		return c.Status(200).JSON(fiber.Map{"authenticated": false})
	}

	// 更新 session（rotate）
	if err := storeSession(sid, newTok); err != nil {
		log.Println("update session err:", err)
	}
	return c.JSON(fiber.Map{"authenticated": true, "access_token": newTok.AccessToken})
}

// POST /api/auth/refresh -> 前端拦截器触发（带 sid cookie），后端用存在 Redis 的 refresh_token 调 Casdoor 刷新
func refreshHandler(c *fiber.Ctx) error {
	sid := c.Cookies(sessionCookieName)
	if sid == "" {
		return c.Status(401).JSON(fiber.Map{"error": "no session"})
	}
	sess, err := getSession(sid)
	if err != nil {
		return c.Status(401).JSON(fiber.Map{"error": "session not found"})
	}
	rt := sess["refresh_token"]
	if rt == "" {
		return c.Status(401).JSON(fiber.Map{"error": "no refresh token"})
	}

	newTok, err := refreshWithRefreshToken(rt)
	if err != nil {
		// refresh token 也过期或无效
		_ = deleteSession(sid)
		c.Cookie(&fiber.Cookie{
			Name:     sessionCookieName,
			Value:    "",
			Expires:  time.Now().Add(-time.Hour),
			HTTPOnly: true,
			Secure:   false,
			Path:     "/",
		})
		return c.Status(401).JSON(fiber.Map{"error": "refresh failed"})
	}

	// rotate: 更新 Redis 存储（用同 sid 存新 refresh token + new access）
	if err := storeSession(sid, newTok); err != nil {
		log.Println("store session after refresh err:", err)
	}

	// 返回新的 access token 给前端（前端把它放内存）
	return c.JSON(fiber.Map{
		"access_token": newTok.AccessToken,
		"expires_in":   newTok.ExpiresIn,
	})
}

// POST /api/auth/logout
func logoutHandler(c *fiber.Ctx) error {
	sid := c.Cookies(sessionCookieName)
	if sid != "" {
		_ = deleteSession(sid)
	}
	// 清除 cookie
	c.Cookie(&fiber.Cookie{
		Name:     sessionCookieName,
		Value:    "",
		Expires:  time.Now().Add(-time.Hour),
		HTTPOnly: true,
		Secure:   false,
		Path:     "/",
	})
	return c.JSON(fiber.Map{"ok": true})
}

/* ------------------- main ------------------- */
func main() {
	// 检查必需 env
	if clientID == "" || clientSecret == "" || authURL == "" || tokenURL == "" || redirectURI == "" || frontendURL == "" {
		log.Println("Please set CASDOOR_CLIENT_ID, CASDOOR_CLIENT_SECRET, CASDOOR_AUTHORIZE_URL, CASDOOR_TOKEN_URL, APP_REDIRECT_URI, FRONTEND_URL")
		// 继续运行以便开发测试，但最好中断
	}

	initRedis()

	app := fiber.New()

	app.Get("/api/auth/login", loginRedirect)
	app.Get("/api/auth/callback", callbackHandler)
	app.Get("/api/auth/session", sessionHandler)
	app.Post("/api/auth/refresh", refreshHandler)
	app.Post("/api/auth/logout", logoutHandler)

	// example protected endpoint: 前端带 access_token in Authorization header
	app.Get("/api/me", func(c *fiber.Ctx) error {
		auth := c.Get("Authorization")
		if auth == "" {
			return c.Status(401).JSON(fiber.Map{"error": "missing token"})
		}
		// 这里可以验证 Casdoor 的 access token 或调用 Casdoor 的 userinfo endpoint
		// 为示例我们只是回显（生产需校验）
		return c.JSON(fiber.Map{"token": auth})
	})

	log.Fatal(app.Listen(":3000"))
}