wincmd/ntfs_index_ctx.go

package wincmd

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"b612.me/win32api"
	"b612.me/wincmd/ntfs/mft"
	"b612.me/wincmd/ntfs/usn"
)

const defaultWatchMaxBatches = 32

// USNBookmark stores resumable watch state.
type USNBookmark struct {
	Volume       string    `json:"volume"`
	VolumeSerial uint32    `json:"volume_serial"`
	UsnJournalID uint64    `json:"usn_journal_id"`
	BookmarkUSN  uint64    `json:"bookmark_usn"`
	UpdatedAt    time.Time `json:"updated_at"`
}

// BuildVolumeIndexContext builds a unified file index and supports cancellation.
func BuildVolumeIndexContext(ctx context.Context, volume string, opts IndexOptions) (*VolumeIndex, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := checkContext(ctx); err != nil {
		return nil, err
	}
	vol, err := normalizeVolume(volume)
	if err != nil {
		return nil, err
	}
	if !opts.IncludeUSN && !opts.IncludeMFT {
		opts.IncludeUSN = true
	}

	idx := &VolumeIndex{
		Volume:  vol,
		BuiltAt: time.Now(),
		ByID:    make(map[uint64]FileMeta),
		ByPath:  make(map[string]uint64),
	}

	emit := func(meta FileMeta) error {
		if err := checkContext(ctx); err != nil {
			return err
		}
		indexMergeMeta(idx, meta, opts.MaxEntries)
		return nil
	}
	if err := BuildVolumeIndexStream(ctx, vol, opts, emit); err != nil {
		return nil, err
	}

	if opts.IncludeUSN {
		if bookmark, err := currentUSNBookmark(vol); err == nil {
			idx.BookmarkUSN = bookmark
		}
	}
	if opts.Filter != nil {
		applyIndexFilter(idx, opts.Filter)
	}
	return idx, nil
}

// BuildVolumeIndexStream emits file metadata incrementally.
func BuildVolumeIndexStream(ctx context.Context, volume string, opts IndexOptions, emit func(FileMeta) error) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := checkContext(ctx); err != nil {
		return err
	}
	if emit == nil {
		return wrapInputError("nil stream emitter")
	}
	vol, err := normalizeVolume(volume)
	if err != nil {
		return err
	}
	if !opts.IncludeUSN && !opts.IncludeMFT {
		opts.IncludeUSN = true
	}

	if opts.IncludeUSN {
		usnMap, err := usn.ListUsnFile(vol)
		if err != nil {
			return fmt.Errorf("list usn files: %w", err)
		}
		resolver := newUSNMetaResolver(vol, usnMap, opts.IncludeFileStat)
		defer resolver.Close()
		for id, entry := range usnMap {
			if err := checkContext(ctx); err != nil {
				return err
			}
			meta := FileMeta{
				ID:       uint64(id),
				ParentID: uint64(entry.Parent),
				Name:     entry.Name,
				Path:     resolver.Path(id),
				IsDir:    entry.Type == 1,
				Source:   "usn",
			}
			if opts.IncludeFileStat {
				resolver.ApplyStat(id, &meta)
			}
			if opts.Filter != nil && !opts.Filter(meta) {
				continue
			}
			if err := emit(meta); err != nil {
				return err
			}
		}
	}

	if opts.IncludeMFT {
		metas := make([]FileMeta, 0)
		fileMap := make(map[uint64]mft.FileEntry)
		err := mft.WalkRecordsByMFT(vol, func(record mft.Record) error {
			if err := checkContext(ctx); err != nil {
				return err
			}
			file, ok := mft.FileFromRecord(record)
			if !ok {
				return nil
			}
			meta := FileMeta{
				ID:       file.Node,
				ParentID: file.Parent,
				Name:     file.Name,
				IsDir:    file.IsDir,
				Size:     file.Size,
				ModTime:  file.ModTime,
				Source:   "mft",
			}
			metas = append(metas, meta)
			fileMap[file.Node] = mft.FileEntry{Name: file.Name, Parent: file.Parent}
			return nil
		})
		if err != nil {
			return fmt.Errorf("walk mft records: %w", err)
		}
		resolveMFTMetaPaths(vol, fileMap, metas)
		for i := range metas {
			if err := checkContext(ctx); err != nil {
				return err
			}
			if opts.Filter != nil && !opts.Filter(metas[i]) {
				continue
			}
			if err := emit(metas[i]); err != nil {
				return err
			}
		}
	}

	return nil
}

func resolveMFTMetaPaths(volume string, fileMap map[uint64]mft.FileEntry, metas []FileMeta) {
	for i := range metas {
		metas[i].Path = mft.GetFullUsnPath(volume, fileMap, metas[i].ID)
	}
}

// WalkIndex traverses entries from an already built index.
func WalkIndex(idx *VolumeIndex, fn func(FileMeta) error) error {
	if idx == nil {
		return wrapInputError("nil index")
	}
	if fn == nil {
		return wrapInputError("nil callback")
	}
	for _, meta := range idx.ByID {
		if err := fn(meta); err != nil {
			return err
		}
	}
	return nil
}

// ResolveFileByIDContext resolves a file id and supports cancellation.
func ResolveFileByIDContext(ctx context.Context, volume string, id uint64) (FileMeta, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := checkContext(ctx); err != nil {
		return FileMeta{}, err
	}
	return ResolveFileByID(volume, id)
}

// WalkFilesContext streams files from USN with cancellation support.
func WalkFilesContext(ctx context.Context, volume string, filter func(FileMeta) bool, fn func(FileMeta) error) error {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := checkContext(ctx); err != nil {
		return err
	}
	if fn == nil {
		return wrapInputError("nil callback")
	}
	vol, err := normalizeVolume(volume)
	if err != nil {
		return err
	}

	fileMap, err := usn.ListUsnFile(vol)
	if err != nil {
		return err
	}
	resolver := newUSNMetaResolver(vol, fileMap, false)
	defer resolver.Close()

	for id, entry := range fileMap {
		if err := checkContext(ctx); err != nil {
			return err
		}
		meta := FileMeta{
			ID:       uint64(id),
			ParentID: uint64(entry.Parent),
			Name:     entry.Name,
			Path:     resolver.Path(id),
			IsDir:    entry.Type == 1,
			Source:   "usn",
		}
		if filter != nil && !filter(meta) {
			continue
		}
		if err := fn(meta); err != nil {
			return err
		}
	}
	return nil
}

// WatchVolumeChangesContext consumes one or more USN batches from a bookmark and emits normalized events.
func WatchVolumeChangesContext(ctx context.Context, volume string, fromUSN uint64, fn func(ChangeEvent) error) (uint64, error) {
	next, _, _, err := watchVolumeChanges(ctx, volume, fromUSN, defaultWatchMaxBatches, fn)
	return next, err
}

// WatchVolumeChangesWithBookmark watches USN changes and persists bookmark to disk.
// It auto-rescans from current head when bookmark is stale due to journal rollover/recreate.
func WatchVolumeChangesWithBookmark(ctx context.Context, volume string, bookmarkFile string, fn func(ChangeEvent) error) (USNBookmark, bool, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := checkContext(ctx); err != nil {
		return USNBookmark{}, false, err
	}
	if bookmarkFile == "" {
		return USNBookmark{}, false, wrapInputError("empty bookmark file")
	}
	vol, err := normalizeVolume(volume)
	if err != nil {
		return USNBookmark{}, false, err
	}
	bookmark, err := LoadUSNBookmark(bookmarkFile)
	if err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			return USNBookmark{}, false, err
		}
		bookmark = USNBookmark{}
	}

	journal, serial, err := queryUSNJournalState(vol)
	if err != nil {
		return USNBookmark{}, false, err
	}

	fromUSN := bookmark.BookmarkUSN
	rescanned := false
	if bookmark.BookmarkUSN == 0 {
		fromUSN = 0
	}
	if bookmark.Volume != "" {
		if bookmark.Volume != vol || bookmark.VolumeSerial != serial || bookmark.UsnJournalID != uint64(journal.UsnJournalID) || bookmark.BookmarkUSN < uint64(journal.FirstUsn) {
			fromUSN = uint64(journal.FirstUsn)
			rescanned = true
		}
	}

	nextUSN, journalID, _, err := watchVolumeChanges(ctx, vol, fromUSN, defaultWatchMaxBatches, fn)
	if err != nil {
		return USNBookmark{}, rescanned, err
	}

	out := USNBookmark{
		Volume:       vol,
		VolumeSerial: serial,
		UsnJournalID: journalID,
		BookmarkUSN:  nextUSN,
		UpdatedAt:    time.Now(),
	}
	if err := SaveUSNBookmark(bookmarkFile, out); err != nil {
		return USNBookmark{}, rescanned, err
	}
	return out, rescanned, nil
}

// SaveUSNBookmark writes bookmark state to disk.
func SaveUSNBookmark(path string, bookmark USNBookmark) error {
	if path == "" {
		return wrapInputError("empty bookmark file")
	}
	dir := filepath.Dir(path)
	if dir != "" && dir != "." {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return err
		}
	}
	content, err := json.MarshalIndent(bookmark, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, content, 0600)
}

// LoadUSNBookmark reads bookmark state from disk.
func LoadUSNBookmark(path string) (USNBookmark, error) {
	if path == "" {
		return USNBookmark{}, wrapInputError("empty bookmark file")
	}
	content, err := os.ReadFile(path)
	if err != nil {
		return USNBookmark{}, err
	}
	var bookmark USNBookmark
	if err := json.Unmarshal(content, &bookmark); err != nil {
		return USNBookmark{}, err
	}
	return bookmark, nil
}

type usnMetaResolver struct {
	volume       string
	entries      map[win32api.DWORDLONG]usn.FileEntry
	pathCache    map[win32api.DWORDLONG]string
	volumeHandle syscall.Handle
}

func newUSNMetaResolver(volume string, entries map[win32api.DWORDLONG]usn.FileEntry, openVolumeHandle bool) *usnMetaResolver {
	resolver := &usnMetaResolver{
		volume:       volume,
		entries:      entries,
		pathCache:    make(map[win32api.DWORDLONG]string, len(entries)),
		volumeHandle: syscall.InvalidHandle,
	}
	if !openVolumeHandle {
		return resolver
	}
	pDriver := `\\.\` + strings.TrimSuffix(volume, `\`)
	if handle, err := usn.CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL); err == nil {
		resolver.volumeHandle = handle
	}
	return resolver
}

func (r *usnMetaResolver) Close() {
	if r == nil {
		return
	}
	if r.volumeHandle != 0 && r.volumeHandle != syscall.InvalidHandle {
		_ = syscall.Close(r.volumeHandle)
		r.volumeHandle = syscall.InvalidHandle
	}
}

func (r *usnMetaResolver) Path(id win32api.DWORDLONG) string {
	if r == nil {
		return ""
	}
	if path, ok := r.pathCache[id]; ok {
		return path
	}
	path := usn.GetFullUsnPath(r.volume, r.entries, id)
	r.pathCache[id] = path
	return path
}

func (r *usnMetaResolver) ApplyStat(id win32api.DWORDLONG, meta *FileMeta) {
	if r == nil {
		return
	}
	applyUSNStatByID(r.volumeHandle, r.entries, id, meta)
}

func applyUSNStatByID(volumeHandle syscall.Handle, entries map[win32api.DWORDLONG]usn.FileEntry, id win32api.DWORDLONG, meta *FileMeta) {
	if meta == nil {
		return
	}
	entry, ok := entries[id]
	if !ok {
		entry = usn.FileEntry{}
	}
	if volumeHandle != 0 && volumeHandle != syscall.InvalidHandle {
		openAttrs := uint32(win32api.FILE_ATTRIBUTE_NORMAL)
		if entry.Type == 1 {
			openAttrs = win32api.FILE_FLAG_BACKUP_SEMANTICS
		}
		if fileHandle, err := usn.OpenFileByIdWithfd(volumeHandle, id, syscall.O_RDONLY, openAttrs); err == nil {
			var info syscall.ByHandleFileInformation
			statErr := syscall.GetFileInformationByHandle(fileHandle, &info)
			_ = syscall.Close(fileHandle)
			if statErr == nil {
				meta.Size = uint64(info.FileSizeHigh)<<32 | uint64(info.FileSizeLow)
				meta.ModTime = time.Unix(0, info.LastWriteTime.Nanoseconds())
				if info.FileAttributes&win32api.FILE_ATTRIBUTE_DIRECTORY != 0 {
					meta.IsDir = true
				}
				return
			}
		}
	}
	if meta.Path == "" {
		return
	}
	stat, err := os.Stat(meta.Path)
	if err != nil {
		return
	}
	if size := stat.Size(); size >= 0 {
		meta.Size = uint64(size)
	}
	meta.ModTime = stat.ModTime()
	if stat.IsDir() {
		meta.IsDir = true
	}
}

func watchVolumeChanges(ctx context.Context, volume string, fromUSN uint64, maxBatches int, fn func(ChangeEvent) error) (nextUSN uint64, journalID uint64, firstUSN uint64, err error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := checkContext(ctx); err != nil {
		return fromUSN, 0, 0, err
	}
	if fn == nil {
		return fromUSN, 0, 0, wrapInputError("nil callback")
	}
	vol, err := normalizeVolume(volume)
	if err != nil {
		return fromUSN, 0, 0, err
	}
	if maxBatches <= 0 {
		maxBatches = defaultWatchMaxBatches
	}

	pathCache, err := usn.ListUsnFile(vol)
	if err != nil {
		return fromUSN, 0, 0, err
	}

	pDriver := `\\.\` + strings.TrimSuffix(vol, `\`)
	fd, err := usn.CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return fromUSN, 0, 0, err
	}
	defer syscall.Close(fd)

	var journal win32api.USN_JOURNAL_DATA
	var done uint32
	if err := usn.DeviceIoControl(fd, win32api.FSCTL_QUERY_USN_JOURNAL, []byte{}, &journal, &done); err != nil {
		return fromUSN, 0, 0, err
	}
	journalID = uint64(journal.UsnJournalID)
	firstUSN = uint64(journal.FirstUsn)

	if fromUSN == 0 {
		fromUSN = uint64(journal.NextUsn)
	}
	if fromUSN < uint64(journal.FirstUsn) {
		fromUSN = uint64(journal.FirstUsn)
	}

	readReq := win32api.READ_USN_JOURNAL_DATA{
		StartUsn:          win32api.USN(fromUSN),
		ReasonMask:        0xFFFFFFFF,
		ReturnOnlyOnClose: 0,
		Timeout:           0,
		BytesToWaitFor:    0,
		UsnJournalID:      journal.UsnJournalID,
	}

	nextUSN = fromUSN
	for batch := 0; batch < maxBatches; batch++ {
		if err := checkContext(ctx); err != nil {
			return nextUSN, journalID, firstUSN, err
		}
		buf := make([]byte, 0x10000)
		done = 0
		if err := usn.DeviceIoControl(fd, win32api.FSCTL_READ_USN_JOURNAL, &readReq, buf, &done); err != nil {
			return nextUSN, journalID, firstUSN, err
		}
		if done <= uint32(watchUSNBufferHeaderSize) {
			return nextUSN, journalID, firstUSN, nil
		}
		if int(done) > len(buf) {
			return nextUSN, journalID, firstUSN, fmt.Errorf("usn output length %d exceeds buffer %d", done, len(buf))
		}

		next := binary.LittleEndian.Uint64(buf[:watchUSNBufferHeaderSize])
		if next == nextUSN {
			return nextUSN, journalID, firstUSN, nil
		}

		if err := parseWatchUSNRecords(buf, done, func(event usnWatchRecord) error {
			if err := checkContext(ctx); err != nil {
				return err
			}
			monitor := usn.FileMonitor{
				Name:   event.FileName,
				Self:   win32api.DWORDLONG(event.FileReferenceNumber),
				Parent: win32api.DWORDLONG(event.ParentFileReferenceNumber),
				Type:   0,
				Reason: usnReasonString(event.Reason),
			}
			if event.FileAttributes&win32api.FILE_ATTRIBUTE_DIRECTORY != 0 {
				monitor.Type = 1
			}

			path := usn.GetFullUsnPathEntry(vol, pathCache, monitor)
			meta := FileMeta{
				ID:       event.FileReferenceNumber,
				ParentID: event.ParentFileReferenceNumber,
				Name:     monitor.Name,
				Path:     path,
				IsDir:    monitor.Type == 1,
				Source:   "usn",
			}
			applyUSNStatByID(fd, pathCache, monitor.Self, &meta)

			return fn(ChangeEvent{USN: event.Usn, Reason: monitor.Reason, File: meta, At: time.Now()})
		}); err != nil {
			return nextUSN, journalID, firstUSN, err
		}

		nextUSN = next
		readReq.StartUsn = win32api.USN(next)
		if nextUSN >= uint64(journal.NextUsn) {
			return nextUSN, journalID, firstUSN, nil
		}
	}

	return nextUSN, journalID, firstUSN, nil
}

func checkContext(ctx context.Context) error {
	select {
	case <-ctx.Done():
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			return wrapTimeoutError(ctx.Err().Error())
		}
		return ctx.Err()
	default:
		return nil
	}
}

func queryUSNJournalState(volume string) (win32api.USN_JOURNAL_DATA, uint32, error) {
	info, err := usn.GetDiskInfo(volume)
	if err != nil {
		return win32api.USN_JOURNAL_DATA{}, 0, err
	}
	pDriver := `\\.\` + strings.TrimSuffix(volume, `\`)
	fd, err := usn.CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return win32api.USN_JOURNAL_DATA{}, 0, err
	}
	defer syscall.Close(fd)

	var journal win32api.USN_JOURNAL_DATA
	var done uint32
	if err := usn.DeviceIoControl(fd, win32api.FSCTL_QUERY_USN_JOURNAL, []byte{}, &journal, &done); err != nil {
		return win32api.USN_JOURNAL_DATA{}, 0, err
	}
	return journal, info.SerialNumber, nil
}