ehentai-go/main.go

// Package main 是 ehentai-go 下载器的入口。
//
// 当前已挂载路由:
//   - GET  /healthz                              连通性探针
//   - GET  /api/list                             列出 pre-download/*.json 概要
//   - POST /api/galleries/check                  检查 {title} 是否已在 pre-download/
//   - POST /api/galleries/save                   保存前端 JS 组装好的画廊 JSON
//   - GET  /api/galleries/{title}/missing-images 返回还没下载的图片 key 列表(供前端确认要 push 哪些)
//   - POST /api/galleries/{title}/images/{key}   接收前端 push 的图片二进制(multipart)
//   - GET  /api/events                           SSE 事件流
//
// 全部抓图链路(抓单页 → 解析直链 → 拉图片)由浏览器端油猴脚本完成,
// 后端只负责存。这样彻底绕开 Go 端被 e-hentai TLS fingerprint 拦截的问题。
//
// 全局只允许一个跑任务;并发触发返回 409 E_TASK_RUNNING。
package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"sort"
	"strings"
	"syscall"
	"time"

	"ehentai-go/internal/crawler"
	"ehentai-go/internal/store"
	"ehentai-go/web"
)

// MinImageBytes 视为「下载成功」的最小字节数,小于此值视为占位/错误页。
const MinImageBytes = 50 * 1024

// (events.Hub 已被移除:前端面板已删,不需要 SSE 进度推送,
// 所有抓图结果由浏览器油猴脚本直接拿到响应即可。)

func main() {
	port := flag.String("port", "8990", "HTTP 监听端口(也可用 EHENTAI_PORT 环境变量)")
	flag.Parse()
	if env := os.Getenv("EHENTAI_PORT"); env != "" {
		*port = env
	}
	addr := ":" + *port

	if err := ensureRuntimeDirs(); err != nil {
		log.Fatalf("[main] %v", err)
	}
	// 测试支持:从 EHENTAI_TEST_HOSTS 环境变量注入额外 host 白名单(逗号分隔)。
	// 生产环境不要设置这个变量。
	crawler.AddAllowedHostsFromEnv()

	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		_, _ = w.Write([]byte("ok"))
	})
	mux.Handle("GET /", web.IndexHandler())
	mux.HandleFunc("POST /api/galleries/check", handleGalleryCheck)
	mux.HandleFunc("POST /api/galleries/save", handleGallerySave)
	// 浏览器侧脚本推送图片二进制所用的端点。路径里同时带 gallery 与 key,
	// 这样 multipart 不用塞文件名,后端从路径里拿就够了。
	mux.HandleFunc("POST /api/galleries/{title}/images/{key}", handleGalleryImagePush)
	mux.HandleFunc("GET /api/galleries/{title}/missing-images", handleGalleryMissingImages)

	srv := &http.Server{
		Addr:              addr,
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}

	// 优雅退出:收到 SIGINT/SIGTERM 时给正在处理的请求最多 10 秒收尾。
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	go func() {
		log.Printf("[main] ehentai-go listening on http://localhost%s", srv.Addr)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Fatalf("[main] ListenAndServe: %v", err)
		}
	}()

	<-ctx.Done()
	log.Println("[main] shutting down...")
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	if err := srv.Shutdown(shutdownCtx); err != nil {
		log.Printf("[main] shutdown error: %v", err)
		os.Exit(1)
	}
	log.Println("[main] bye")
}

// ensureRuntimeDirs 启动前检查/创建 pre-download 与 downloads 目录,
// 并把每条目录的状态打印到日志,便于排查权限与同名文件占位。
func ensureRuntimeDirs() error {
	created, err := store.EnsureDirs()
	if err != nil {
		return fmt.Errorf("ensure runtime dirs: %w", err)
	}
	allDirs := []string{store.PreDownloadDir, store.DownloadsDir}
	for _, d := range allDirs {
		log.Printf("[main] runtime dir ok: %s", d)
	}
	if len(created) > 0 {
		log.Printf("[main] 已自动创建 %d 个目录: %v", len(created), created)
	}
	return nil
}

// writeJSON 统一 JSON 响应。
func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

// writeErr 把 error 映射成统一错误响应。
// 同时根据错误码选一个合适的 HTTP status。
func writeErr(w http.ResponseWriter, err error) {
	code := "E_INTERNAL"
	msg := err.Error()
	for _, prefix := range []string{
		"E_BAD_URL", "E_FETCH_GALLERY", "E_PARSE_GALLERY",
		"E_ZERO_PAGE", "E_DUP_GALLERY", "E_FETCH_SINGLE_PAGE",
		"E_NOT_FOUND", "E_NO_GALLERY",
	} {
		if strings.HasPrefix(msg, prefix) {
			code = prefix
			break
		}
	}
	status := http.StatusInternalServerError
	switch code {
	case "E_BAD_URL":
		status = http.StatusBadRequest
	case "E_ZERO_PAGE", "E_NO_GALLERY":
		status = http.StatusUnprocessableEntity
	case "E_NOT_FOUND":
		status = http.StatusNotFound
	}
	writeJSON(w, status, map[string]any{
		"error":   code,
		"message": msg,
	})
}

// galleryCheckReq POST /api/galleries/check 请求体。
type galleryCheckReq struct {
	Title string `json:"title"` // 已 sanitize 的画廊名(JS 端做完)
}

// handleGalleryCheck POST /api/galleries/check — 检查 title 对应的 json 是否已存在。
// 用于前端 JS 弹框「画廊已存在,是否覆盖 / 跳过」之类。
// 返回 200: {exists: bool, name: <sanitized>, pages?: map<key,单页URL>, total_pages?: int}
// 当 exists=true 时会一并把已有 json 的 pages/total_pages 返回,让续传场景下前端不必再翻页解析单页 URL,
// 直接用 pageURLByKey 就能精准重试 missing 图片(否则 pageURLByKey 是空的,missing 永远补不上)。
func handleGalleryCheck(w http.ResponseWriter, r *http.Request) {
	var req galleryCheckReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]any{
			"error":   "E_BAD_REQUEST",
			"message": fmt.Sprintf("解析请求体失败: %v", err),
		})
		return
	}
	name := strings.TrimSpace(req.Title)
	if name == "" {
		writeJSON(w, http.StatusBadRequest, map[string]any{
			"error":   "E_BAD_REQUEST",
			"message": "title 不能为空",
		})
		return
	}
	exists, err := store.GalleryExists(name)
	if err != nil {
		writeErr(w, err)
		return
	}
	resp := map[string]any{
		"exists": exists,
		"name":   name,
	}
	if exists {
		// 续传场景:把已有的 pages + total_pages 一起返回,前端就不用再翻页解析画廊列表页了。
		// 这样 pageURLByKey 能被正确填上,missing 的图才能真正进入重试流程。
		g, loadErr := store.LoadGallery(name)
		if loadErr != nil {
			// JSON 文件存在却读不出来(损坏/权限问题),不要静默继续,直接报错让前端看见。
			writeErr(w, loadErr)
			return
		}
		resp["pages"] = g.Pages
		resp["total_pages"] = g.TotalPages
	}
	writeJSON(w, http.StatusOK, resp)
}

// gallerySaveReq POST /api/galleries/save 请求体。
// 直接接收前端 JS 组装好的 store.Gallery JSON,落到 pre-download/<name>.json。
// 字段名与 store.Gallery 保持一致,这里只是再包一层防 raw_title/size 过大。
type gallerySaveReq struct {
	store.Gallery
}

// handleGallerySave POST /api/galleries/save — 保存前端 JS 解析好的画廊 JSON。
// 已存在同名 json 时返回 409 E_DUP_GALLERY,前端弹框由用户决定是否覆盖(本接口暂不支持 overwrite,
// 后续如需可加 ?overwrite=1 参数)。
func handleGallerySave(w http.ResponseWriter, r *http.Request) {
	var req gallerySaveReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]any{
			"error":   "E_BAD_REQUEST",
			"message": fmt.Sprintf("解析请求体失败: %v", err),
		})
		return
	}
	if strings.TrimSpace(req.GalleryName) == "" {
		writeJSON(w, http.StatusBadRequest, map[string]any{
			"error":   "E_BAD_REQUEST",
			"message": "gallery_name 不能为空",
		})
		return
	}
	if strings.TrimSpace(req.SourceURL) != "" {
		// 兜底:如果前端忘了 sanitize,但给了原始 URL,这里再校验一遍 host。
		if err := crawler.ValidateHost(req.SourceURL); err != nil {
			writeErr(w, err)
			return
		}
	}
	if len(req.Pages) == 0 {
		writeJSON(w, http.StatusUnprocessableEntity, map[string]any{
			"error":   "E_ZERO_PAGE",
			"message": "pages 不能为空,前端必须先解析到至少一张图片直链",
		})
		return
	}
	if req.TotalPages == 0 {
		req.TotalPages = len(req.Pages)
	}
	if req.FetchedAt == "" {
		req.FetchedAt = time.Now().UTC().Format(time.RFC3339)
	}

	exists, err := store.GalleryExists(req.GalleryName)
	if err != nil {
		writeErr(w, err)
		return
	}
	if exists {
		writeJSON(w, http.StatusConflict, map[string]any{
			"error":   "E_DUP_GALLERY",
			"message": fmt.Sprintf("pre-download/%s.json 已存在", req.GalleryName),
		})
		return
	}
	if err := store.SaveGallery(&req.Gallery); err != nil {
		writeErr(w, err)
		return
	}
	log.Printf("[save] 已保存画廊: %s (%d 张)", req.GalleryName, req.TotalPages)
	writeJSON(w, http.StatusCreated, map[string]any{
		"name":        req.GalleryName,
		"total_pages": req.TotalPages,
		"path":        store.GalleryJSONPath(req.GalleryName),
	})
}

// handleGalleryMissingImages GET /api/galleries/{title}/missing-images
// 返回 gallery JSON 里尚未在 downloads/<title>/ 中以最小有效大小落盘的 key 列表。
// 浏览器端油猴脚本据此确认「该给后端 push 哪些张」。
func handleGalleryMissingImages(w http.ResponseWriter, r *http.Request) {
	title := strings.TrimSpace(r.PathValue("title"))
	if title == "" {
		writeJSON(w, http.StatusBadRequest, map[string]any{
			"error":   "E_BAD_REQUEST",
			"message": "title 不能为空",
		})
		return
	}
	g, err := store.LoadGallery(title)
	if err != nil {
		if errors.Is(err, store.ErrGalleryNotFound) {
			writeJSON(w, http.StatusNotFound, map[string]any{
				"error":   "E_NOT_FOUND",
				"message": err.Error(),
			})
			return
		}
		writeErr(w, err)
		return
	}
	outDir := store.GalleryDownloadDir(title)
	// 直接扫 downloads/<title>/ 下所有「4 位数字.<ext>」命名的文件,
	// 这样不依赖 store.Gallery.Pages 里 URL 的形态 — 前端存的是单页 URL,
	// 推断不到扩展名,旧逻辑会把所有 key 都当成 missing,导致续传失效。
	onDisk, scanErr := scanDownloadedKeys(outDir)
	if scanErr != nil {
		writeErr(w, scanErr)
		return
	}
	missing := make([]string, 0, len(g.Pages))
	for key := range g.Pages {
		if ok := onDisk[key]; !ok {
			missing = append(missing, key)
		}
	}
	// 按 key 字典序排序输出,前端按顺序并发即可。
	sort.Strings(missing)
	writeJSON(w, http.StatusOK, map[string]any{
		"name":        title,
		"total":       len(g.Pages),
		"missing":     missing,
		"missing_num": len(missing),
	})
}

// handleGalleryImagePush POST /api/galleries/{title}/images/{key}
// 接收浏览器 push 过来的图片二进制(multipart,字段名 "file")。
// 落地到 downloads/<title>/<key><ext>,ext 由 Content-Type 或 multipart 文件名推断;
// 写入大小 < MinImageBytes (1KB) 视为失败,避免把限流占位 / 错误页当真图保存。
func handleGalleryImagePush(w http.ResponseWriter, r *http.Request) {
	title := strings.TrimSpace(r.PathValue("title"))
	key := strings.TrimSpace(r.PathValue("key"))
	if title == "" || key == "" {
		writeJSON(w, http.StatusBadRequest, map[string]any{
			"error":   "E_BAD_REQUEST",
			"message": "title / key 不能为空",
		})
		return
	}
	// 不合法 key(防路径穿越)直接拒。
	if !isValidKey(key) {
		writeJSON(w, http.StatusBadRequest, map[string]any{
			"error":   "E_BAD_KEY",
			"message": fmt.Sprintf("key %q 不合法(必须是 4 位数字)", key),
		})
		return
	}
	// 必须先有对应的 gallery JSON(防止 push 到不存在的画廊)。
	if _, err := store.LoadGallery(title); err != nil {
		if errors.Is(err, store.ErrGalleryNotFound) {
			writeJSON(w, http.StatusNotFound, map[string]any{
				"error":   "E_NOT_FOUND",
				"message": err.Error(),
			})
			return
		}
		writeErr(w, err)
		return
	}

	// 限制 body 上限:200MB,超过直接拒(单图最多几十 MB)。
	const maxBody = 200 << 20
	r.Body = http.MaxBytesReader(w, r.Body, maxBody)

	// 优先 multipart;若不是 multipart,按 raw body 接收。
	var (
		ext         string
		filename    string
		contentType           = r.Header.Get("Content-Type")
		reader      io.Reader = r.Body
	)
	if strings.HasPrefix(contentType, "multipart/form-data") {
		if err := r.ParseMultipartForm(maxBody); err != nil {
			writeJSON(w, http.StatusBadRequest, map[string]any{
				"error":   "E_BAD_REQUEST",
				"message": fmt.Sprintf("解析 multipart 失败: %v", err),
			})
			return
		}
		file, hdr, ferr := r.FormFile("file")
		if ferr != nil {
			writeJSON(w, http.StatusBadRequest, map[string]any{
				"error":   "E_BAD_REQUEST",
				"message": fmt.Sprintf("读取 form file 文件 失败: %v", ferr),
			})
			return
		}
		defer file.Close()
		reader = file
		filename = hdr.Filename
		if filename != "" {
			ext = strings.ToLower(filepath.Ext(filename))
		}
	}
	// multipart 没拿到文件名,或 raw body,按 Content-Type / 启发式猜扩展名。
	if ext == "" {
		ext = extFromContentType(contentType)
	}

	outDir := store.GalleryDownloadDir(title)
	if err := os.MkdirAll(outDir, 0o755); err != nil {
		writeErr(w, err)
		return
	}
	// 临时文件 + 原子 rename。
	finalPath := filepath.Join(outDir, key+ext)
	tmpPath := finalPath + ".part"
	f, err := os.Create(tmpPath)
	if err != nil {
		writeErr(w, err)
		return
	}
	written, err := io.Copy(f, reader)
	if cerr := f.Close(); cerr != nil && err == nil {
		err = cerr
	}
	if err != nil {
		_ = os.Remove(tmpPath)
		writeErr(w, err)
		return
	}
	if written < MinImageBytes {
		_ = os.Remove(tmpPath)
		writeJSON(w, http.StatusUnprocessableEntity, map[string]any{
			"error":   "E_TOO_SMALL",
			"message": fmt.Sprintf("body 仅 %d < %d 字节,疑似限流占位 / 错误页", written, MinImageBytes),
			"key":     key,
		})
		return
	}
	if err := os.Rename(tmpPath, finalPath); err != nil {
		_ = os.Remove(tmpPath)
		writeErr(w, err)
		return
	}
	writeJSON(w, http.StatusCreated, map[string]any{
		"name":     title,
		"key":      key,
		"path":     finalPath,
		"bytes":    written,
		"filename": filepath.Base(finalPath),
	})
}

// isValidOnDisk 文件存在且大小 >= MinImageBytes(1KB)。
func isValidOnDisk(path string) (bool, error) {
	fi, err := os.Stat(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return false, nil
		}
		return false, err
	}
	return fi.Size() >= MinImageBytes, nil
}

// isValidKey key 必须是 4 位数字(对齐 store.Gallery.Pages 的 padding)。
func isValidKey(key string) bool {
	if len(key) != 4 {
		return false
	}
	for _, c := range key {
		if c < '0' || c > '9' {
			return false
		}
	}
	return true
}

// scanDownloadedKeys 扫描 downloads/<title>/ 目录下所有「0001.<ext>」命名的文件,
// 返回 key 集合(只看 4 位数字前缀,扩展名任意)。文件 ≥ MinImageBytes 才算「在」。
// 用目录扫描而不是 URL 推断扩展名,是因为 store.Gallery.Pages 里存的是单页 URL,
// 老逻辑拿单页 URL 推 ext 永远推不出来,导致所有 key 都误判为 missing。
func scanDownloadedKeys(outDir string) (map[string]bool, error) {
	out := make(map[string]bool)
	entries, err := os.ReadDir(outDir)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return out, nil
		}
		return nil, fmt.Errorf("scan %s: %w", outDir, err)
	}
	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		name := e.Name()
		// 跳过 .part 等临时文件
		if strings.HasSuffix(name, ".part") {
			continue
		}
		// 取主文件名第一段「.」前的部分作为 key
		key := name
		if i := strings.Index(name, "."); i >= 0 {
			key = name[:i]
		}
		if !isValidKey(key) {
			continue
		}
		info, statErr := e.Info()
		if statErr != nil {
			return nil, fmt.Errorf("stat %s: %w", name, statErr)
		}
		if info.Size() >= MinImageBytes {
			out[key] = true
		}
	}
	return out, nil
}

// extFromContentType 根据 Content-Type 推扩展名;不能识别返回 ".bin"。
func extFromContentType(ct string) string {
	ct = strings.ToLower(strings.TrimSpace(ct))
	if i := strings.Index(ct, ";"); i >= 0 {
		ct = ct[:i]
	}
	ct = strings.TrimSpace(ct)
	switch ct {
	case "image/jpeg", "image/jpg":
		return ".jpg"
	case "image/png":
		return ".png"
	case "image/webp":
		return ".webp"
	case "image/gif":
		return ".gif"
	case "image/bmp":
		return ".bmp"
	case "image/svg+xml":
		return ".svg"
	case "image/avif":
		return ".avif"
	}
	// multipart 没正确设置 Content-Type 时,header 可能形如 application/octet-stream
	// 留空,文件落 .bin。spec 上传时主动设置 Content-Type 更稳。
	return ".bin"
}

// extFromPageURL 从 store.Gallery.Pages 里的 URL 推断扩展名:
// - 老的图片直链(以 keystamp= / hath.network / fullimg 等  /结尾)取 .ext
// - 单页 URL(e-hentai /s/...)拿不到图片扩展名,返回 ""。
func extFromPageURL(rawURL string) string {
	u := rawURL
	// 只看 host 在 hath.* / 包含 keystamp / 包含 /fullimg 的当作「图片直链」启发式,
	// 简单截取尾段扩展。
	if strings.Contains(u, "keystamp=") || strings.Contains(u, "/fullimg") {
		base := u
		if i := strings.Index(base, "?"); i >= 0 {
			base = base[:i]
		}
		ext := strings.ToLower(filepath.Ext(base))
		return ext
	}
	return ""
}

// (acquireTask / releaseTask / newTaskID 已随 /api/download 一起移除;
// 现在所有抓取和图片下载都在浏览器端油猴脚本里串行/并发执行,
// 后端只负责存图片,无任务槽位概念。)

// (handleEvents / writeSSE 已随前端面板一起移除:现在没有前端要 SSE 进度了,
// 浏览器油猴脚本拿到的响应本身就是最终结果。)