eh-downloader/internal/core/bridge.go

package core

// 本地桥口：让油猴脚本把画廊地址推给本工具
//
//	POST /add        { gid, token, url, title }  +  X-EHDL-Token 头
//	POST /add-batch  { items: [...] }            +  X-EHDL-Token 头
//	GET  /ping                                     -> { ok: true }
//	GET  /info                                     -> 端口/版本等信息
//	GET  /tasks                                    -> 当前任务列表
//	GET  /pushed                                   -> 推送历史（任务删了也还在）
//
// 只监听 127.0.0.1，且同时校验 Origin 与固定 token，防止本机其他网页乱塞任务。

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"eh-downloader-wails/internal/eh"
)

// BridgeToken 固定 token（和油猴脚本里的 BRIDGE_TOKEN 保持一致）
const BridgeToken = "ehdl-local-9f2a4c7e"

var allowedOrigins = map[string]bool{
	"https://e-hentai.org":      true,
	"https://exhentai.org":      true,
	"https://www.e-hentai.org": true,
}

const maxBody = 512 * 1024 // 批量最多 500 个，标题也要带，得给够

// Bridge 是本地 HTTP 桥口
type Bridge struct {
	store   *Store
	pushed  *PushedStore
	log     func(level, msg string)
	onAdded func(t Task, added bool, batch bool)

	mu     sync.Mutex
	server *http.Server
	port   int
}

// NewBridge 建桥口
func NewBridge(store *Store, log func(string, string), onAdded func(Task, bool, bool)) *Bridge {
	return &Bridge{store: store, log: log, onAdded: onAdded, port: 37777}
}

// SetPushedStore 挂上推送历史存储（可选，用于 /pushed 返回历史推送记录）
func (b *Bridge) SetPushedStore(p *PushedStore) {
	b.pushed = p
}

// Start 启动（只绑回环）
func (b *Bridge) Start(port int) {
	if port <= 0 {
		port = 37777
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	b.port = port

	mux := http.NewServeMux()
	mux.HandleFunc("/ping", b.handlePing)
	mux.HandleFunc("/info", b.handleInfo)
	mux.HandleFunc("/tasks", b.handleTasks)
	mux.HandleFunc("/pushed", b.handlePushed)
	mux.HandleFunc("/add", b.handleAdd)
	mux.HandleFunc("/add-batch", b.handleAddBatch)

	srv := &http.Server{
		Handler:           b.cors(mux),
		ReadHeaderTimeout: 10 * time.Second,
	}

	ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
	if err != nil {
		if strings.Contains(err.Error(), "address already in use") {
			b.log("warn", fmt.Sprintf("桥口端口 %d 被占用，油猴推送不可用。可在设置里改端口，或关掉占用它的程序。", port))
		} else {
			b.log("error", fmt.Sprintf("桥口启动失败：%v", err))
		}
		return
	}

	b.server = srv
	go func() {
		if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed {
			b.log("error", fmt.Sprintf("桥口异常：%v", err))
		}
	}()
	b.log("info", fmt.Sprintf("本地桥口已就绪 http://127.0.0.1:%d（油猴推送可用）", port))
}

// Stop 停止桥口
func (b *Bridge) Stop() {
	b.mu.Lock()
	srv := b.server
	b.server = nil
	b.mu.Unlock()
	if srv != nil {
		_ = srv.Close()
	}
}

// Port 当前端口
func (b *Bridge) Port() int {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.port
}

// cors 处理 CORS 预检与响应头
func (b *Bridge) cors(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		origin := r.Header.Get("Origin")
		if origin != "" && allowedOrigins[origin] {
			w.Header().Set("Access-Control-Allow-Origin", origin)
			w.Header().Set("Vary", "Origin")
		}
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, X-EHDL-Token")
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusNoContent)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func writeJSON(w http.ResponseWriter, code int, v interface{}) {
	body, _ := json.Marshal(v)
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(code)
	w.Write(body)
}

// readBody 读 JSON body（限大小）
func readBody(r *http.Request) (map[string]interface{}, int) {
	data, err := io.ReadAll(io.LimitReader(r.Body, maxBody+1))
	if err != nil {
		return nil, http.StatusBadRequest
	}
	if len(data) > maxBody {
		return nil, http.StatusRequestEntityTooLarge
	}
	var m map[string]interface{}
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, http.StatusBadRequest
	}
	return m, 0
}

func checkToken(w http.ResponseWriter, r *http.Request) bool {
	if r.Header.Get("X-EHDL-Token") != BridgeToken {
		writeJSON(w, http.StatusForbidden, map[string]interface{}{"ok": false, "error": "bad token"})
		return false
	}
	return true
}

func (b *Bridge) handlePing(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "name": "EH Downloader"})
}

func (b *Bridge) handleInfo(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ok":      true,
		"port":    b.Port(),
		"version": AppVersion,
	})
}

// handleTasks 返回已入队的画廊清单（gid/token），供油猴脚本标记「已推送」
// 只读，不改状态；token 用固定头校验，防止本机网页随便翻任务列表。
//
// 注意：这里返回的是「当前任务列表」。任务被清除后就没了。
// 油猴脚本标灰要看的是「推送历史」，用 /pushed 那个。
func (b *Bridge) handleTasks(w http.ResponseWriter, r *http.Request) {
	if !checkToken(w, r) {
		b.log("warn", "桥口收到 token 不正确的任务列表请求，已拒绝")
		return
	}
	tasks := b.store.Tasks()
	out := make([]map[string]string, 0, len(tasks))
	for _, t := range tasks {
		out = append(out, map[string]string{
			"gid":    t.GID,
			"token":  t.Token,
			"title":  t.Title,
			"status": t.Status,
		})
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ok":    true,
		"count": len(out),
		"tasks": out,
	})
}

// handlePushed 返回「推送历史」：所有推送过的画廊，任务删了也还在。
// 油猴脚本靠它把标题标灰 —— 这样点「清除已完成」之后，灰标依然保留。
func (b *Bridge) handlePushed(w http.ResponseWriter, r *http.Request) {
	if !checkToken(w, r) {
		b.log("warn", "桥口收到 token 不正确的推送历史请求，已拒绝")
		return
	}
	entries := []PushedEntry{}
	if b.pushed != nil {
		entries = b.pushed.Entries()
	}
	out := make([]map[string]interface{}, 0, len(entries))
	for _, e := range entries {
		out = append(out, map[string]interface{}{
			"gid":       e.GID,
			"token":     e.Token,
			"title":     e.Title,
			"pushedAt":  e.PushedAt,
			"completed": e.Completed,
		})
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ok":    true,
		"count": len(out),
		"items": out,
	})
}

type batchItem struct {
	GID   string `json:"gid"`
	Token string `json:"token"`
	URL   string `json:"url"`
	Title string `json:"title"`
}

func (b *Bridge) handleAddBatch(w http.ResponseWriter, r *http.Request) {
	if !checkToken(w, r) {
		b.log("warn", "桥口收到 token 不正确的批量请求，已拒绝")
		return
	}
	body, code := readBody(r)
	if code != 0 {
		writeJSON(w, code, map[string]interface{}{"ok": false, "error": "bad request"})
		return
	}

	raw, _ := body["items"].([]interface{})
	if len(raw) == 0 {
		writeJSON(w, http.StatusBadRequest, map[string]interface{}{"ok": false, "error": "items 为空"})
		return
	}

	seen := map[string]bool{}
	var fresh []eh.GalleryRef
	existed := 0

	for _, it := range raw {
		if len(fresh) >= 500 {
			break
		}
		m, ok := it.(map[string]interface{})
		if !ok {
			continue
		}
		item := batchItem{
			GID:   str(m["gid"]),
			Token: str(m["token"]),
			URL:   str(m["url"]),
			Title: str(m["title"]),
		}
		if item.GID == "" || item.Token == "" {
			continue
		}
		key := item.GID + "/" + item.Token
		if seen[key] {
			continue
		}
		seen[key] = true

		if b.store.FindTask(item.GID, item.Token) != nil {
			existed++
			continue
		}
		fresh = append(fresh, eh.GalleryRef{GID: item.GID, Token: item.Token, Raw: item.URL, Title: item.Title})
	}

	added := b.store.AddTasks(fresh)
	for _, ref := range fresh {
		if ref.Title != "" {
			b.store.SetTitle(ref.GID, ref.Token, ref.Title)
		}
	}
	b.store.ScheduleFlush(0)

	// 记进推送历史（任务被清掉也不影响标灰）
	if b.pushed != nil {
		entries := make([]PushedEntry, 0, len(fresh))
		for _, ref := range fresh {
			entries = append(entries, PushedEntry{GID: ref.GID, Token: ref.Token, Title: ref.Title})
		}
		b.pushed.MarkBatch(entries)
	}

	for _, ref := range fresh {
		if t := b.store.FindTask(ref.GID, ref.Token); t != nil && b.onAdded != nil {
			b.onAdded(*t, true, true)
		}
	}

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ok":       true,
		"added":    added,
		"existed":  existed,
		"received": len(raw),
		"invalid":  len(raw) - len(seen),
	})
}

func (b *Bridge) handleAdd(w http.ResponseWriter, r *http.Request) {
	if !checkToken(w, r) {
		b.log("warn", "桥口收到 token 不正确的请求，已拒绝")
		return
	}
	body, code := readBody(r)
	if code != 0 {
		writeJSON(w, code, map[string]interface{}{"ok": false, "error": "bad request"})
		return
	}

	gid := strings.TrimSpace(str(body["gid"]))
	token := strings.TrimSpace(str(body["token"]))
	if gid == "" || token == "" {
		writeJSON(w, http.StatusBadRequest, map[string]interface{}{"ok": false, "error": "缺少 gid 或 token"})
		return
	}

	exists := b.store.FindTask(gid, token) != nil
	added := b.store.AddTasks([]eh.GalleryRef{{GID: gid, Token: token, Raw: str(body["url"])}})
	if t := str(body["title"]); t != "" {
		b.store.SetTitle(gid, token, t)
	}
	b.store.ScheduleFlush(0)

	// 记进推送历史（任务被清掉也不影响标灰）
	if b.pushed != nil {
		b.pushed.Mark(gid, token, str(body["title"]))
	}

	if t := b.store.FindTask(gid, token); t != nil && b.onAdded != nil {
		b.onAdded(*t, added > 0, false)
	}

	queued := false
	if t := b.store.FindTask(gid, token); t != nil {
		queued = t.Status == "queued"
	}

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ok":      true,
		"added":   added > 0,
		"existed": exists,
		"gid":     gid,
		"queued":  queued,
	})
}

// str 安全取字符串
func str(v interface{}) string {
	s, _ := v.(string)
	return s
}