kenaito-media-picker/internal/probe/exif.go

package probe

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"time"
)

// ffmpegRotateTimeout 调 ffmpeg 转正图片的超时时间。超大图转起来慢，给足时间。
const ffmpegRotateTimeout = 120 * time.Second

// ReadExifOrientation 读取 JPEG 的 EXIF Orientation（0x0112）标记。
// 返回 1~8（标准 EXIF 方向值）；文件不是 JPEG、没有该标记或读取失败时返回 0。
func ReadExifOrientation(filePath string) int {
	app1, err := readExifAPP1(filePath)
	if err != nil {
		return 0
	}
	return exifOrientation(app1)
}

// NeedsDimensionSwap 判断该方向是否需要"宽高对调"（90°/270° 旋转）。
// EXIF 方向 5/6/7/8 是 90°/270° 旋转，横竖会互换；2/3/4 只是镜像/180°，横竖不变。
func NeedsDimensionSwap(orientation int) bool {
	switch orientation {
	case 5, 6, 7, 8:
		return true
	}
	return false
}

// RotateImage 用 ffmpeg 把 EXIF 旋转"烧进"像素，把图片转正后写回原文件。
//
// 做法：
//  1. 先把原来的 EXIF APP1 段读出来（里面通常有拍摄日期、相机型号等，不能丢）；
//  2. 调 ffmpeg 旋转像素（输出时清掉旋转标记，防止被二次旋转）；
//  3. 把 EXIF 方向标记改成 1（正常方向）后回写进转正后的文件。
//
// ffmpegPath 为空时直接报错，调用方决定怎么兜底。
func RotateImage(ffmpegPath, filePath string) error {
	if ffmpegPath == "" {
		return errors.New("未配置 ffmpeg，无法转正图片")
	}
	app1, _ := readExifAPP1(filePath) // 读不到 EXIF 也不影响转正，忽略即可

	tmp := filePath + ".rotate.tmp"
	ctx, cancel := context.WithTimeout(context.Background(), ffmpegRotateTimeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, ffmpegPath, "-y", "-v", "error",
		"-i", filePath,
		"-f", "image2",
		"-map_metadata", "0",
		"-metadata:s:v", "rotate=0",
		"-q:v", "2",
		tmp)
	if out, err := cmd.CombinedOutput(); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("ffmpeg 转正失败: %v: %s", err, string(out))
	}

	// 把原 EXIF 回写回去（方向标记改成 1，避免查看器再转一次）
	if len(app1) > 0 {
		setOrientationToOne(app1)
		if err := injectExifAPP1(tmp, app1); err != nil {
			fmt.Printf("=== 回写EXIF信息失败(不影响转正结果): %v\n", err)
		}
	}

	// 用转正后的文件替换原文件
	if err := os.Remove(filePath); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("替换原文件失败: %v", err)
	}
	if err := os.Rename(tmp, filePath); err != nil {
		return fmt.Errorf("替换原文件失败: %v", err)
	}
	return nil
}

// exifHeaderReadLimit 读 EXIF 只扫 JPEG 开头的这一小段。
// APP1 段永远在文件最前面（SOI 之后不久），读 64KB 足够定位。
const exifHeaderReadLimit = 64 * 1024

// readExifAPP1 从 JPEG 文件里读出第一个"Exif"的 APP1 段（FF E1 开头的整段）。
func readExifAPP1(filePath string) ([]byte, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	buf := make([]byte, exifHeaderReadLimit)
	n, _ := io.ReadFull(f, buf)
	buf = buf[:n]

	if len(buf) < 2 || buf[0] != 0xFF || buf[1] != 0xD8 {
		return nil, errors.New("不是 JPEG 文件")
	}
	pos := 2
	for pos+4 <= len(buf) {
		if buf[pos] != 0xFF {
			break
		}
		marker := buf[pos+1]
		// 无长度段的标记：跳过
		if marker == 0xD8 || (marker >= 0xD0 && marker <= 0xD7) || marker == 0x01 {
			pos += 2
			continue
		}
		// 到图像数据（SOS）或文件尾了，后面没有 EXIF
		if marker == 0xDA || marker == 0xD9 {
			break
		}
		segLen := int(binary.BigEndian.Uint16(buf[pos+2 : pos+4]))
		if segLen < 2 {
			break
		}
		if marker == 0xE1 && segLen >= 8 && string(buf[pos+4:pos+8]) == "Exif" {
			end := pos + 2 + segLen
			if end > len(buf) {
				end = len(buf)
			}
			seg := make([]byte, end-pos)
			copy(seg, buf[pos:end])
			return seg, nil
		}
		pos += 2 + segLen
	}
	return nil, errors.New("没有 Exif APP1 段")
}

// exifOrientation 从 Exif 的 APP1 段里解析 Orientation（0x0112）。
// 布局：FF E1 [段长] "Exif\0\0" TIFF头 IFD0...
func exifOrientation(app1 []byte) int {
	if len(app1) < 12 {
		return 0
	}
	t := app1[10:] // TIFF 头起始（"Exif\0\0" 占 6 字节）
	if len(t) < 8 {
		return 0
	}
	order := tiffByteOrder(t)
	if order == nil {
		return 0
	}
	if order.Uint16(t[2:4]) != 42 { // TIFF 魔数
		return 0
	}
	ifd := int(order.Uint32(t[4:8]))
	if ifd < 8 || ifd+2 > len(t) {
		return 0
	}
	count := int(order.Uint16(t[ifd : ifd+2]))
	for i := 0; i < count; i++ {
		e := ifd + 2 + i*12
		if e+12 > len(t) {
			break
		}
		if order.Uint16(t[e:e+2]) == 0x0112 {
			return int(order.Uint16(t[e+8 : e+10]))
		}
	}
	return 0
}

// setOrientationToOne 把 EXIF 的 Orientation 标记改成 1（正常方向）。
func setOrientationToOne(app1 []byte) {
	if len(app1) < 12 {
		return
	}
	t := app1[10:]
	if len(t) < 8 {
		return
	}
	order := tiffByteOrder(t)
	if order == nil {
		return
	}
	ifd := int(order.Uint32(t[4:8]))
	if ifd < 8 || ifd+2 > len(t) {
		return
	}
	count := int(order.Uint16(t[ifd : ifd+2]))
	for i := 0; i < count; i++ {
		e := ifd + 2 + i*12
		if e+12 > len(t) {
			break
		}
		if order.Uint16(t[e:e+2]) == 0x0112 {
			t[e+8] = 1
			t[e+9] = 0
			t[e+10] = 0
			t[e+11] = 0
			return
		}
	}
}

// tiffByteOrder 判断 TIFF 头的字节序，返回对应的 binary.ByteOrder（不认识返回 nil）。
func tiffByteOrder(t []byte) binary.ByteOrder {
	switch {
	case len(t) >= 2 && t[0] == 'I' && t[1] == 'I':
		return binary.LittleEndian
	case len(t) >= 2 && t[0] == 'M' && t[1] == 'M':
		return binary.BigEndian
	}
	return nil
}

// injectExifAPP1 把 Exif APP1 段插到 JPEG 的 SOI 之后（覆盖式重写文件）。
// 文件里原本的 Exif APP1 会被去掉，只留新插入的这一份。
func injectExifAPP1(filePath string, app1 []byte) error {
	data, err := os.ReadFile(filePath)
	if err != nil {
		return err
	}
	if len(data) < 2 || data[0] != 0xFF || data[1] != 0xD8 {
		return errors.New("不是 JPEG 文件")
	}
	rest := stripExifAPP1(data[2:])
	out := make([]byte, 0, 2+len(app1)+len(rest))
	out = append(out, data[0], data[1])
	out = append(out, app1...)
	out = append(out, rest...)
	return os.WriteFile(filePath, out, 0644)
}

// stripExifAPP1 从 JPEG 的 SOI 之后的内容里去掉所有"Exif"的 APP1 段，其余原样保留。
func stripExifAPP1(rest []byte) []byte {
	pos := 0
	out := make([]byte, 0, len(rest))
	for pos+4 <= len(rest) {
		if rest[pos] != 0xFF {
			break
		}
		marker := rest[pos+1]
		// SOS（图像数据开始）或 EOI：后面原样保留
		if marker == 0xDA || marker == 0xD9 {
			out = append(out, rest[pos:]...)
			return out
		}
		if marker >= 0xD0 && marker <= 0xD7 || marker == 0x01 {
			out = append(out, rest[pos], rest[pos+1])
			pos += 2
			continue
		}
		segLen := int(binary.BigEndian.Uint16(rest[pos+2 : pos+4]))
		if segLen < 2 {
			out = append(out, rest[pos:]...)
			return out
		}
		isExif := marker == 0xE1 && segLen >= 8 && string(rest[pos+4:pos+8]) == "Exif"
		if !isExif {
			out = append(out, rest[pos:pos+2+segLen]...)
		}
		pos += 2 + segLen
	}
	out = append(out, rest[pos:]...)
	return out
}