wincmd/ntfs_index.go

package wincmd

import (
	"context"
	"encoding/binary"
	"fmt"
	"strings"
	"syscall"
	"time"
	"unsafe"

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

const (
	watchUSNBufferHeaderSize            = int(unsafe.Sizeof(win32api.USN(0)))
	watchUSNRecordMinSize               = int(unsafe.Offsetof(win32api.USN_RECORD{}.FileName))
	watchUSNRecordOffsetUsn             = int(unsafe.Offsetof(win32api.USN_RECORD{}.Usn))
	watchUSNRecordOffsetFileReference   = int(unsafe.Offsetof(win32api.USN_RECORD{}.FileReferenceNumber))
	watchUSNRecordOffsetParentReference = int(unsafe.Offsetof(win32api.USN_RECORD{}.ParentFileReferenceNumber))
	watchUSNRecordOffsetReason          = int(unsafe.Offsetof(win32api.USN_RECORD{}.Reason))
	watchUSNRecordOffsetFileAttributes  = int(unsafe.Offsetof(win32api.USN_RECORD{}.FileAttributes))
	watchUSNRecordOffsetFileNameLength  = int(unsafe.Offsetof(win32api.USN_RECORD{}.FileNameLength))
	watchUSNRecordOffsetFileNameOffset  = int(unsafe.Offsetof(win32api.USN_RECORD{}.FileNameOffset))
)

// FileMeta is the unified file metadata model used by wincmd task-level APIs.
type FileMeta struct {
	ID       uint64
	ParentID uint64
	Name     string
	Path     string
	IsDir    bool
	Size     uint64
	ModTime  time.Time
	Source   string
}

// IndexOptions controls how BuildVolumeIndex collects metadata.
type IndexOptions struct {
	IncludeUSN      bool
	IncludeMFT      bool
	IncludeFileStat bool
	MaxEntries      int
	Filter          func(FileMeta) bool
}

// VolumeIndex is an in-memory query index for a volume.
type VolumeIndex struct {
	Volume      string
	BuiltAt     time.Time
	BookmarkUSN uint64
	ByID        map[uint64]FileMeta
	ByPath      map[string]uint64
}

// ChangeEvent is a normalized USN change record.
type ChangeEvent struct {
	USN    uint64
	Reason string
	File   FileMeta
	At     time.Time
}

// BuildVolumeIndex builds a unified file index from USN/MFT sources.
func BuildVolumeIndex(volume string, opts IndexOptions) (*VolumeIndex, error) {
	return BuildVolumeIndexContext(context.Background(), volume, opts)
}

// ResolveFileByID resolves a file id to a canonical metadata entry.
func ResolveFileByID(volume string, id uint64) (FileMeta, error) {
	vol, err := normalizeVolume(volume)
	if err != nil {
		return FileMeta{}, err
	}

	fileMap, err := usn.ListUsnFile(vol)
	if err != nil {
		return FileMeta{}, err
	}

	did := win32api.DWORDLONG(id)
	entry, ok := fileMap[did]
	if !ok {
		return FileMeta{}, wrapNotFoundError(fmt.Sprintf("file id %d", id))
	}

	meta := FileMeta{
		ID:       id,
		ParentID: uint64(entry.Parent),
		Name:     entry.Name,
		Path:     usn.GetFullUsnPath(vol, fileMap, did),
		IsDir:    entry.Type == 1,
		Source:   "usn",
	}

	if stat, err := usn.GetUsnFileInfo(vol, fileMap, did); err == nil {
		meta.Size = uint64(stat.Size())
		meta.ModTime = stat.ModTime()
		if stat.IsDir() {
			meta.IsDir = true
		}
	}
	return meta, nil
}

// WalkFiles streams files from USN and invokes callback for each matching entry.
func WalkFiles(volume string, filter func(FileMeta) bool, fn func(FileMeta) error) error {
	return WalkFilesContext(context.Background(), volume, filter, fn)
}

// WatchVolumeChanges consumes one or more USN batches from a bookmark and emits normalized events.
// If fromUSN is 0, it tails from the journal's current NextUsn (new changes only).
func WatchVolumeChanges(volume string, fromUSN uint64, fn func(ChangeEvent) error) (nextUSN uint64, err error) {
	return WatchVolumeChangesContext(context.Background(), volume, fromUSN, fn)
}

type usnWatchRecord struct {
	Usn                       uint64
	FileReferenceNumber       uint64
	ParentFileReferenceNumber uint64
	Reason                    uint32
	FileAttributes            uint32
	FileName                  string
}

func parseWatchUSNRecords(buf []byte, done uint32, fn func(usnWatchRecord) error) error {
	for offset := watchUSNBufferHeaderSize; offset < int(done); {
		remaining := int(done) - offset
		if remaining < watchUSNRecordMinSize {
			return fmt.Errorf("usn record header truncated: remaining=%d", remaining)
		}

		recordLength := int(binary.LittleEndian.Uint32(buf[offset:]))
		if recordLength < watchUSNRecordMinSize {
			return fmt.Errorf("invalid usn record length %d", recordLength)
		}
		if recordLength > remaining {
			return fmt.Errorf("usn record length %d exceeds remaining %d", recordLength, remaining)
		}

		record := buf[offset : offset+recordLength]
		nameLen := int(binary.LittleEndian.Uint16(record[watchUSNRecordOffsetFileNameLength:]))
		nameOff := int(binary.LittleEndian.Uint16(record[watchUSNRecordOffsetFileNameOffset:]))
		if nameLen < 0 || nameLen%2 != 0 {
			return fmt.Errorf("invalid usn name length %d", nameLen)
		}
		if nameOff < watchUSNRecordMinSize || nameOff > recordLength {
			return fmt.Errorf("invalid usn name offset %d", nameOff)
		}
		if nameOff+nameLen > recordLength {
			return fmt.Errorf("usn name out of record boundary")
		}

		fileName, err := decodeUTF16Bytes(record[nameOff : nameOff+nameLen])
		if err != nil {
			return err
		}

		event := usnWatchRecord{
			Usn:                       binary.LittleEndian.Uint64(record[watchUSNRecordOffsetUsn:]),
			FileReferenceNumber:       binary.LittleEndian.Uint64(record[watchUSNRecordOffsetFileReference:]),
			ParentFileReferenceNumber: binary.LittleEndian.Uint64(record[watchUSNRecordOffsetParentReference:]),
			Reason:                    binary.LittleEndian.Uint32(record[watchUSNRecordOffsetReason:]),
			FileAttributes:            binary.LittleEndian.Uint32(record[watchUSNRecordOffsetFileAttributes:]),
			FileName:                  fileName,
		}
		if err := fn(event); err != nil {
			return err
		}

		offset += recordLength
	}
	return nil
}

func decodeUTF16Bytes(data []byte) (string, error) {
	if len(data)%2 != 0 {
		return "", fmt.Errorf("invalid utf16 byte length %d", len(data))
	}
	chars := make([]uint16, len(data)/2)
	for i := range chars {
		chars[i] = binary.LittleEndian.Uint16(data[i*2:])
	}
	return syscall.UTF16ToString(chars), nil
}

func currentUSNBookmark(volume string) (uint64, error) {
	pDriver := `\\.\` + strings.TrimSuffix(volume, `\`)
	fd, err := usn.CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return 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 0, err
	}
	return uint64(journal.NextUsn), nil
}

func normalizeVolume(volume string) (string, error) {
	v := strings.TrimSpace(strings.ReplaceAll(volume, "/", "\\"))
	if v == "" {
		return "", wrapInputError("empty volume")
	}
	if len(v) == 2 && v[1] == ':' {
		v += "\\"
	}
	if len(v) < 3 || v[1] != ':' {
		return "", wrapVolumeError(volume, nil)
	}
	if v[len(v)-1] != '\\' {
		v += "\\"
	}
	return strings.ToUpper(v[:1]) + v[1:], nil
}

func normalizePathKey(path string) string {
	if path == "" {
		return ""
	}
	return strings.ToLower(strings.ReplaceAll(strings.TrimSpace(path), "/", "\\"))
}

func indexMergeMeta(idx *VolumeIndex, incoming FileMeta, maxEntries int) {
	if idx == nil {
		return
	}
	if current, ok := idx.ByID[incoming.ID]; ok {
		merged := mergeFileMeta(current, incoming)
		idx.ByID[incoming.ID] = merged
		if key := normalizePathKey(merged.Path); key != "" {
			idx.ByPath[key] = merged.ID
		}
		return
	}
	if maxEntries > 0 && len(idx.ByID) >= maxEntries {
		return
	}
	idx.ByID[incoming.ID] = incoming
	if key := normalizePathKey(incoming.Path); key != "" {
		idx.ByPath[key] = incoming.ID
	}
}

func mergeFileMeta(current FileMeta, incoming FileMeta) FileMeta {
	merged := current
	if merged.Name == "" {
		merged.Name = incoming.Name
	}
	if merged.Path == "" {
		merged.Path = incoming.Path
	}
	if merged.ParentID == 0 {
		merged.ParentID = incoming.ParentID
	}
	if incoming.IsDir {
		merged.IsDir = true
	}
	if merged.Size == 0 && incoming.Size != 0 {
		merged.Size = incoming.Size
	}
	if merged.ModTime.IsZero() && !incoming.ModTime.IsZero() {
		merged.ModTime = incoming.ModTime
	}
	if merged.Source == "" {
		merged.Source = incoming.Source
	} else if incoming.Source != "" && merged.Source != incoming.Source && !strings.Contains(merged.Source, incoming.Source) {
		merged.Source = merged.Source + "+" + incoming.Source
	}
	return merged
}

func applyIndexFilter(idx *VolumeIndex, filter func(FileMeta) bool) {
	if idx == nil || filter == nil {
		return
	}
	for id, meta := range idx.ByID {
		if filter(meta) {
			continue
		}
		delete(idx.ByID, id)
		if key := normalizePathKey(meta.Path); key != "" {
			delete(idx.ByPath, key)
		}
	}
}

func usnReasonString(reason uint32) string {
	return usn.USNReasonString(win32api.DWORD(reason))
}