diff --git a/pkg/fs/file.go b/pkg/fs/file.go index f1967d963..9808abae9 100644 --- a/pkg/fs/file.go +++ b/pkg/fs/file.go @@ -22,10 +22,13 @@ package fs import ( + "encoding/binary" + "errors" "expvar" "fmt" "os" "path/filepath" + "runtime" "strings" "sync" "sync/atomic" @@ -61,6 +64,9 @@ const ( // metadataStoreQueueSize is the capacity of the async request channel. const metadataStoreQueueSize = 1500 +const negativeTTL = time.Minute // cache TTL for negative (non-PE) entries +const maxEntries = 50000 // negative entries are not added beyond this + var fsMetadataAsyncRequestDrops = expvar.NewInt("fs.metadata.async.request.drops") var fsMetadataCount = expvar.NewInt("fs.metadata.count") @@ -82,6 +88,12 @@ var onceMetadataStore sync.Once // metadataTTL is the maximum age of an untouched cache entry before the GC evicts it. var metadataTTL = 10 * time.Minute +var ( + errNotPE = errors.New("not a PE file") + errShortFile = errors.New("file too short to classify") + errRateLimited = errors.New("rate-limited for PE parsing") +) + // windowsUpdateWildcards lists system directories commonly touched during // updates, servicing, recovery, and component store operations. // Excluded from metastore because they generate a lot of legitimate file @@ -131,12 +143,25 @@ type FileInfo struct { IsDriver bool IsDotnet bool + negative bool + // accessed is updated on every cache lookup to drive TTL-based eviction. accessed atomic.Int64 } func (f *FileInfo) keepalive() { - f.accessed.Store(time.Now().UnixNano()) + now := time.Now().UnixNano() + if now-f.accessed.Load() > int64(time.Second) { + f.accessed.Store(now) + } +} + +func (f *FileInfo) expired(now time.Time) bool { + ttl := metadataTTL + if f.negative { + ttl = negativeTTL + } + return now.Sub(f.lastAccessed()) > ttl } func (f *FileInfo) lastAccessed() time.Time { @@ -144,9 +169,9 @@ func (f *FileInfo) lastAccessed() time.Time { } type Request struct { - Path string - Response chan *FileInfo - IsSeed bool // request originated from the initial rundown + Path string + IsSeed bool // request originated from the initial rundown + Sync bool // caller is invoking a sync resolution } func GetMetadataStore() *FileMetadataStore { @@ -176,13 +201,16 @@ type FileMetadataStore struct { requests chan Request - stop chan struct{} - group singleflight.Group + stop chan struct{} + closeOnce sync.Once + group singleflight.Group purger *time.Ticker limiter *rate.Limiter + pending sync.Map + windowsUpdateWildcards windowsUpdateWildcards wellKnownDLLs moduleWildcards wellKnownExecutables moduleWildcards @@ -296,7 +324,10 @@ func newFileMetadataStore() *FileMetadataStore { }, } - const numWorkers = 6 + numWorkers := 2 + if runtime.NumCPU() >= 16 { + numWorkers = 4 + } for range numWorkers { go s.runWorker() } @@ -311,7 +342,7 @@ func newFileMetadataStore() *FileMetadataStore { // in the field accessors. func (s *FileMetadataStore) DoRequest(path string) *FileInfo { p := s.normalizePath(path) - if f := s.get(p); f != nil { + if f, ok := s.lookup(p); ok { return f } @@ -320,25 +351,13 @@ func (s *FileMetadataStore) DoRequest(path string) *FileInfo { } if s.wellKnownDLLs.Accept(p) { - s.addDLL(p) - return s.get(p) + return s.addDLL(p) } if s.wellKnownExecutables.Accept(p) { - s.addExecutable(p) - return s.get(p) + return s.addExecutable(p) } - ch := make(chan *FileInfo, 1) - r := Request{Path: p, Response: ch} - select { - case s.requests <- r: - default: - // queue full: fall back to inline check rather than - // dropping a decision that has a security consequence. - return s.getOrParse(r) - } - f := <-ch - return f + return s.resolve(Request{Path: p, Sync: true}) } func (s *FileMetadataStore) DoRequestAsync(path string, isSeed bool) { @@ -352,31 +371,48 @@ func (s *FileMetadataStore) DoRequestAsync(path string, isSeed bool) { } if s.wellKnownDLLs.Accept(p) { - s.addDLL(p) + _ = s.addDLL(p) return } if s.wellKnownExecutables.Accept(p) { - s.addExecutable(p) + _ = s.addExecutable(p) + return + } + + if _, loaded := s.pending.LoadOrStore(p, struct{}{}); loaded { return } select { case s.requests <- Request{Path: p, IsSeed: isSeed}: default: + s.pending.Delete(p) // queue is full fsMetadataAsyncRequestDrops.Add(1) } } func (s *FileMetadataStore) Close() { - s.purger.Stop() - close(s.stop) + s.closeOnce.Do(func() { + s.purger.Stop() + close(s.stop) + }) } func (s *FileMetadataStore) AddFile(path string, f *FileInfo) { + s.add(s.normalizePath(path), f) +} + +func (s *FileMetadataStore) add(path string, f *FileInfo) { + f.keepalive() s.mux.Lock() defer s.mux.Unlock() - fsMetadataCount.Add(1) + if f.negative && len(s.files) >= maxEntries { + return + } + if _, ok := s.files[path]; !ok { + fsMetadataCount.Add(1) + } s.files[path] = f } @@ -384,8 +420,10 @@ func (s *FileMetadataStore) RemoveFile(path string) { p := s.normalizePath(path) s.mux.Lock() defer s.mux.Unlock() - delete(s.files, p) - fsMetadataCount.Add(-1) + if _, ok := s.files[p]; ok { + delete(s.files, p) + fsMetadataCount.Add(-1) + } } func (s *FileMetadataStore) runWorker() { @@ -420,7 +458,7 @@ func (s *FileMetadataStore) gc() { s.mux.RLock() var paths []string for path, file := range s.files { - if now.Sub(file.lastAccessed()) > metadataTTL { + if file.expired(now) { paths = append(paths, path) } } @@ -435,7 +473,7 @@ func (s *FileMetadataStore) gc() { file := s.files[path] // re-check under the write lock: the entry may have been // refreshed between the RUnlock above and this Lock - if file != nil && now.Sub(file.lastAccessed()) > metadataTTL { + if file != nil && file.expired(now) { log.Debugf("evicting file metadata for %s", path) fsMetadataCount.Add(-1) fsMetadataEvictions.Add(1) @@ -446,24 +484,40 @@ func (s *FileMetadataStore) gc() { } func (s *FileMetadataStore) processRequest(r Request) { - f := s.getOrParse(r) - if r.Response != nil { - r.Response <- f + _ = s.resolve(r) + s.pending.Delete(r.Path) +} + +func (s *FileMetadataStore) lookup(path string) (*FileInfo, bool) { + s.mux.RLock() + f := s.files[path] + s.mux.RUnlock() + if f == nil { + return nil, false + } + fsMetadataCacheHits.Add(1) + f.keepalive() + if f.negative { + return nil, true } + return f, true } func (s *FileMetadataStore) contains(path string) bool { - return s.get(path) != nil + _, contains := s.lookup(path) + return contains } -func (s *FileMetadataStore) addDLL(path string) { +func (s *FileMetadataStore) addDLL(path string) *FileInfo { f := &FileInfo{IsDLL: true} - s.AddFile(path, f) + s.add(path, f) + return f } -func (s *FileMetadataStore) addExecutable(path string) { +func (s *FileMetadataStore) addExecutable(path string) *FileInfo { f := &FileInfo{IsExecutable: true} - s.AddFile(path, f) + s.add(path, f) + return f } func (s *FileMetadataStore) get(path string) *FileInfo { @@ -477,60 +531,108 @@ func (s *FileMetadataStore) get(path string) *FileInfo { return f } -func (s *FileMetadataStore) getOrParse(r Request) *FileInfo { +func (s *FileMetadataStore) resolve(r Request) *FileInfo { path := r.Path - if f := s.get(path); f != nil { + if f, ok := s.lookup(r.Path); ok { return f } - v, err := s.group.Do(path, func() (any, error) { - pe, err := s.parsePE(path, r.IsSeed) - if err != nil { + v, err := s.group.Do(path, func() (res any, err error) { + defer func() { // a malformed PE must not kill the worker or hang callers + if rec := recover(); rec != nil { + err = fmt.Errorf("panic parsing PE %s: %v", r.Path, rec) + } + }() + + if f, ok := s.lookup(r.Path); ok { + return f, nil // a previous flight may have filled it + } + + parsed, err := s.parsePE(path, r.IsSeed || r.Sync) + switch { + case errors.Is(err, errNotPE), errors.Is(err, sys.ErrFileSkipped): + s.add(r.Path, &FileInfo{negative: true}) + return (*FileInfo)(nil), nil + case err != nil: return nil, err } - return pe, nil + f := &FileInfo{ + IsDLL: parsed.IsDLL, IsDriver: parsed.IsDriver, + IsExecutable: parsed.IsExecutable, IsDotnet: parsed.IsDotnet, + } + s.add(r.Path, f) + fsMetadataCacheMisses.Add(1) + return f, nil }) if err != nil { return nil } + return v.(*FileInfo) +} - p := v.(*pe.PE) - - f := &FileInfo{} - f.keepalive() - f.IsDLL, f.IsDriver, f.IsExecutable, f.IsDotnet = p.IsDLL, p.IsDriver, p.IsExecutable, p.IsDotnet - s.AddFile(path, f) - fsMetadataCacheMisses.Add(1) - - return f +func errDesc(err error) string { + var pathErr *os.PathError + if errors.As(err, &pathErr) { + err = pathErr.Err + } + return err.Error() } -func (s *FileMetadataStore) parsePE(path string, isSeed bool) (*pe.PE, error) { - if !isSeed && !s.limiter.Allow() { - return nil, fmt.Errorf("path %s rate-limited for PE parsing", path) +func (s *FileMetadataStore) parsePE(path string, noLimit bool) (*pe.PE, error) { + accept := func(head []byte) error { + if err := isPE(head); err != nil { + return err + } + // the limiter only guards the expensive full read + if !noLimit && !s.limiter.Allow() { + return errRateLimited + } + return nil } const size = 8 * 1024 * 1024 // 8MB - data, err := sys.ReadFile(path, size, time.Millisecond*500) - if err != nil { - fsMetadataFileParseErrors.Add(err.Error(), 1) + data, err := sys.ReadFileWithPredicate(path, size, 1024, time.Millisecond*500, accept) + switch { + case err == nil: + case errors.Is(err, errNotPE), errors.Is(err, errShortFile), + errors.Is(err, errRateLimited), errors.Is(err, sys.ErrFileSkipped): + return nil, err + default: + fsMetadataFileParseErrors.Add(errDesc(err), 1) return nil, err } - pe, err := pe.ParseBytes(data, pe.WithSections(), pe.WithSymbols(), pe.WithCLR()) + parsed, err := pe.ParseBytes(data, pe.WithSections(), pe.WithSymbols(), pe.WithCLR()) if err != nil { - fsMetadataFileParseErrors.Add(err.Error(), 1) + fsMetadataFileParseErrors.Add(errDesc(err), 1) return nil, err } - - return pe, err + return parsed, nil } func (s *FileMetadataStore) normalizePath(path string) string { return strings.ToLower(path) } +// isPE validates the DOS and NT signatures from the first 1KB. +func isPE(h []byte) error { + if len(h) < 64 { + return errShortFile // file may not be written yet + } + if h[0] != 'M' || h[1] != 'Z' { + return errNotPE + } + off := binary.LittleEndian.Uint32(h[0x3c:0x40]) + if uint64(off)+4 > uint64(len(h)) { + return nil // NT header beyond the probe size + } + if string(h[off:off+4]) != "PE\x00\x00" { + return errNotPE + } + return nil +} + // GetFileType returns the underlying file type. The opts parameter corresponds to the NtCreateFile CreateOptions argument // that specifies the options to be applied when creating or opening the file. func GetFileType(filename string, opts uint32) FileType { diff --git a/pkg/fs/file_test.go b/pkg/fs/file_test.go index 71cae64b9..1a1029617 100644 --- a/pkg/fs/file_test.go +++ b/pkg/fs/file_test.go @@ -22,6 +22,11 @@ package fs import ( + "bytes" + "encoding/binary" + "errors" + "expvar" + "fmt" "os" "path/filepath" "strings" @@ -31,6 +36,7 @@ import ( "time" "github.com/stretchr/testify/assert" + "golang.org/x/time/rate" ) func sysRoot(t *testing.T) string { @@ -243,7 +249,7 @@ func TestFileMetadataStoreGetKeepsAlive(t *testing.T) { f := s.get(path) old := f.lastAccessed() - time.Sleep(10 * time.Millisecond) + time.Sleep(1005 * time.Millisecond) f2 := s.get(path) if !f2.lastAccessed().After(old) { t.Error("expected lastAccessed to advance on repeated get()") @@ -568,6 +574,440 @@ func TestCloseStopsWorkersAndPurger(t *testing.T) { } } +func cachedEntry(s *FileMetadataStore, path string) (*FileInfo, bool) { + s.mux.RLock() + defer s.mux.RUnlock() + f, ok := s.files[s.normalizePath(path)] + return f, ok +} + +func parseErrorTotal() (n int64) { + fsMetadataFileParseErrors.Do(func(kv expvar.KeyValue) { + if v, ok := kv.Value.(*expvar.Int); ok { + n += v.Value() + } + }) + return n +} + +func writeTestFile(t *testing.T, name string, content []byte) string { + t.Helper() + p := filepath.Join(t.TempDir(), name) + if err := os.WriteFile(p, content, 0644); err != nil { + t.Fatal(err) + } + return p +} + +// realPE returns the bytes of a genuine PE shipped with Windows. +func realPE(t *testing.T) []byte { + t.Helper() + root := os.Getenv("SystemRoot") + if root == "" { + root = `C:\Windows` + } + b, err := os.ReadFile(filepath.Join(root, "System32", "cmd.exe")) + if err != nil { + t.Skipf("cannot read cmd.exe: %v", err) + } + return b +} + +// header builds an n-byte buffer with an MZ magic, e_lfanew=lfanew and, if +// sig is non-empty and fits, sig written at lfanew. +func header(n int, lfanew uint32, sig string) []byte { + b := make([]byte, n) + b[0], b[1] = 'M', 'Z' + binary.LittleEndian.PutUint32(b[0x3c:], lfanew) + if sig != "" && uint64(lfanew)+uint64(len(sig)) <= uint64(n) { + copy(b[lfanew:], sig) + } + return b +} + +func TestCheckPEHeader(t *testing.T) { + const pesig = "PE\x00\x00" + notMZ := make([]byte, 1024) + wrongMagic := header(1024, 0x80, pesig) + wrongMagic[0], wrongMagic[1] = 'Z', 'M' + + var tests = []struct { + name string + h []byte + err error + }{ + {"nil", nil, errShortFile}, + {"empty", []byte{}, errShortFile}, + {"63 bytes", make([]byte, 63), errShortFile}, + {"64 zero bytes", make([]byte, 64), errNotPE}, + {"zeros", notMZ, errNotPE}, + {"swapped magic", wrongMagic, errNotPE}, + {"valid PE signature", header(1024, 0x80, pesig), nil}, + {"valid, header exactly 64 bytes with sig beyond", header(64, 0x80, pesig), nil}, + {"bad NT signature", header(1024, 0x80, "XXXX"), errNotPE}, + {"e_lfanew points inside the DOS header", header(1024, 0, ""), errNotPE}, + {"NT header beyond probe is undecidable", header(1024, 0x2000, ""), nil}, + {"e_lfanew 0xFFFFFFFF does not overflow", header(1024, 0xFFFFFFFF, ""), nil}, + {"sig ends exactly at probe end, valid", header(1024, 1020, pesig), nil}, + {"sig ends exactly at probe end, invalid", header(1024, 1020, "XXXX"), errNotPE}, + {"sig straddles probe end", header(1024, 1021, ""), nil}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := isPE(tt.h) + if !errors.Is(err, tt.err) { + t.Fatalf("got %v, want %v", err, tt.err) + } + }) + } +} + +func TestFileInfoExpired(t *testing.T) { + now := time.Now() + age := func(d time.Duration, negative bool) *FileInfo { + f := &FileInfo{negative: negative} + f.accessed.Store(now.Add(-d).UnixNano()) + return f + } + + if age(30*time.Second, true).expired(now) { + t.Error("negative entry younger than negativeTTL must not be expired") + } + if !age(negativeTTL+time.Second, true).expired(now) { + t.Error("negative entry older than negativeTTL must be expired") + } + if age(negativeTTL+time.Second, false).expired(now) { + t.Error("positive entry must use metadataTTL, not negativeTTL") + } + if !age(metadataTTL+time.Second, false).expired(now) { + t.Error("positive entry older than metadataTTL must be expired") + } +} + +func TestLookup(t *testing.T) { + s := newTestStore() + defer s.Close() + + if f, ok := s.lookup(`c:\missing`); f != nil || ok { + t.Errorf("missing: got (%v, %v), want (nil, false)", f, ok) + } + + pos := &FileInfo{IsDLL: true} + s.add(`c:\pos.dll`, pos) + if f, ok := s.lookup(`c:\pos.dll`); f != pos || !ok { + t.Errorf("positive: got (%v, %v)", f, ok) + } + + s.add(`c:\neg.txt`, &FileInfo{negative: true}) + if f, ok := s.lookup(`c:\neg.txt`); f != nil || !ok { + t.Errorf("negative must be (nil, true), got (%v, %v)", f, ok) + } + if !s.contains(`c:\neg.txt`) { + t.Error("contains must be true for a negative entry so async requests are skipped") + } +} + +func TestAddNegativeSkippedAtCapacity(t *testing.T) { + s := newTestStore() + defer s.Close() + + s.mux.Lock() + for i := 0; i < maxEntries; i++ { + s.files[fmt.Sprintf(`c:\fill\%d`, i)] = &FileInfo{} + } + s.mux.Unlock() + + s.add(`c:\over\neg.txt`, &FileInfo{negative: true}) + if _, ok := cachedEntry(s, `c:\over\neg.txt`); ok { + t.Error("negative entry must not be added once the map is at maxEntries") + } + + s.add(`c:\over\pos.dll`, &FileInfo{IsDLL: true}) + if _, ok := cachedEntry(s, `c:\over\pos.dll`); !ok { + t.Error("positive entries must still be added at capacity") + } +} + +func TestGCEvictsNegativeEntriesEarlier(t *testing.T) { + s := newTestStore() + defer s.Close() + + old := time.Now().Add(-(negativeTTL + 30*time.Second)).UnixNano() + + neg, pos := &FileInfo{negative: true}, &FileInfo{IsDLL: true} + s.add(`c:\gc\neg.txt`, neg) + s.add(`c:\gc\pos.dll`, pos) + neg.accessed.Store(old) + pos.accessed.Store(old) + + s.gc() + + if _, ok := cachedEntry(s, `c:\gc\neg.txt`); ok { + t.Error("negative entry past negativeTTL should be evicted") + } + if _, ok := cachedEntry(s, `c:\gc\pos.dll`); !ok { + t.Error("positive entry younger than metadataTTL must survive") + } +} + +func TestDoRequestNonPECachedNegative(t *testing.T) { + s := newTestStore() + defer s.Close() + + p := writeTestFile(t, "data.log", bytes.Repeat([]byte("x"), 4096)) + errsBefore := parseErrorTotal() + + if f := s.DoRequest(p); f != nil { + t.Fatalf("non-PE file must return nil, got %+v", f) + } + e, ok := cachedEntry(s, p) + if !ok || !e.negative { + t.Fatalf("expected negative cache entry, got (%+v, %v)", e, ok) + } + if got := parseErrorTotal(); got != errsBefore { + t.Errorf("non-PE must not count as a parse error, delta=%d", got-errsBefore) + } +} + +func TestDoRequestBadNTSignatureCachedNegative(t *testing.T) { + s := newTestStore() + defer s.Close() + + p := writeTestFile(t, "fake.dll", header(512, 0x80, "XXXX")) + if f := s.DoRequest(p); f != nil { + t.Fatalf("MZ with bad NT signature must return nil, got %+v", f) + } + if e, ok := cachedEntry(s, p); !ok || !e.negative { + t.Fatalf("expected negative cache entry, got (%+v, %v)", e, ok) + } +} + +func TestNegativeEntryServedUntilInvalidated(t *testing.T) { + s := newTestStore() + defer s.Close() + + pe := realPE(t) + p := writeTestFile(t, "later.bin", bytes.Repeat([]byte("x"), 256)) + + if f := s.DoRequest(p); f != nil { + t.Fatalf("first request: expected nil, got %+v", f) + } + + if err := os.WriteFile(p, pe, 0644); err != nil { + t.Fatal(err) + } + if f := s.DoRequest(p); f != nil { + t.Fatal("negative entry should still be served from cache before invalidation") + } + + s.RemoveFile(p) + + f := s.DoRequest(p) + if f == nil { + t.Fatal("after RemoveFile the PE must be re-read and classified") + } + if !f.IsExecutable || f.negative { + t.Errorf("unexpected classification: %+v", f) + } +} + +func TestShortFilesNotCached(t *testing.T) { + pe := realPE(t) + + for _, tt := range []struct { + name string + content []byte + }{ + {"1 byte", []byte("M")}, + {"63 bytes", bytes.Repeat([]byte("x"), 63)}, + } { + t.Run(tt.name, func(t *testing.T) { + s := newTestStore() + defer s.Close() + + p := writeTestFile(t, "growing.exe", tt.content) + errsBefore := parseErrorTotal() + + if f := s.DoRequest(p); f != nil { + t.Fatalf("expected nil, got %+v", f) + } + if _, ok := cachedEntry(s, p); ok { + t.Fatal("short file must not be cached, positively or negatively") + } + if got := parseErrorTotal(); got != errsBefore { + t.Errorf("short file must not count as parse error, delta=%d", got-errsBefore) + } + + // the writer finishes: no RemoveFile needed because nothing was cached + if err := os.WriteFile(p, pe, 0644); err != nil { + t.Fatal(err) + } + if f := s.DoRequest(p); f == nil || !f.IsExecutable { + t.Fatalf("completed file should classify, got %+v", f) + } + }) + } +} + +func TestEmptyFileNotCachedNorCounted(t *testing.T) { + s := newTestStore() + defer s.Close() + + p := writeTestFile(t, "empty.exe", nil) + + if f := s.DoRequest(p); f != nil { + t.Fatalf("expected nil, got %+v", f) + } + if _, ok := cachedEntry(s, p); ok { + t.Error("empty file must not be cached") + } +} + +func TestMissingFileCountedNotCached(t *testing.T) { + s := newTestStore() + defer s.Close() + + p := filepath.Join(t.TempDir(), "gone.dll") + errsBefore := parseErrorTotal() + + if f := s.DoRequest(p); f != nil { + t.Fatalf("expected nil, got %+v", f) + } + if _, ok := cachedEntry(s, p); ok { + t.Error("open failures are transient and must not be cached") + } + if got := parseErrorTotal(); got != errsBefore+1 { + t.Errorf("open failure should count once, delta=%d", got-errsBefore) + } +} + +func TestParseErrorKeyDoesNotContainPath(t *testing.T) { + s := newTestStore() + defer s.Close() + + p := filepath.Join(t.TempDir(), "unique-marker-gone.dll") + s.DoRequest(p) + + fsMetadataFileParseErrors.Do(func(kv expvar.KeyValue) { + if bytes.Contains([]byte(kv.Key), []byte("unique-marker-gone")) { + t.Errorf("expvar key leaks the file path: %q", kv.Key) + } + }) +} + +func TestConcurrentRequestsShareOneParse(t *testing.T) { + for _, tt := range []struct { + name string + content func(t *testing.T) []byte + negative bool + }{ + {"PE", realPE, false}, + {"non-PE", func(*testing.T) []byte { return bytes.Repeat([]byte("x"), 4096) }, true}, + } { + t.Run(tt.name, func(t *testing.T) { + s := newTestStore() + defer s.Close() + + p := writeTestFile(t, "shared.bin", tt.content(t)) + missesBefore := fsMetadataCacheMisses.Value() + + const n = 32 + results := make([]*FileInfo, n) + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < n; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + results[i] = s.DoRequest(p) + }(i) + } + close(start) + wg.Wait() + + for i := 1; i < n; i++ { + if results[i] != results[0] { + t.Fatalf("caller %d got a different *FileInfo than caller 0", i) + } + } + if tt.negative { + if results[0] != nil { + t.Fatalf("expected nil, got %+v", results[0]) + } + return + } + if results[0] == nil { + t.Fatal("expected non-nil result") + } + if got := fsMetadataCacheMisses.Value() - missesBefore; got != 1 { + t.Errorf("cache misses = %d, want 1 (one parse for all callers)", got) + } + }) + } +} + +func TestRateLimiterGuardsOnlyTheFullRead(t *testing.T) { + pe := realPE(t) + exhausted := func(s *FileMetadataStore) { s.limiter = rate.NewLimiter(0, 0) } + + t.Run("async PE is rate limited and not cached", func(t *testing.T) { + s := newTestStore() + defer s.Close() + exhausted(s) + + p := writeTestFile(t, "a.exe", pe) + errsBefore := parseErrorTotal() + + if f := s.resolve(Request{Path: s.normalizePath(p)}); f != nil { + t.Fatalf("expected nil when rate limited, got %+v", f) + } + if _, ok := cachedEntry(s, p); ok { + t.Error("rate-limited result must not be cached, the next event should retry") + } + if got := parseErrorTotal(); got != errsBefore { + t.Errorf("rate limiting is not a parse error, delta=%d", got-errsBefore) + } + }) + + t.Run("async non-PE is still classified because the probe is not limited", func(t *testing.T) { + s := newTestStore() + defer s.Close() + exhausted(s) + + p := writeTestFile(t, "a.txt", bytes.Repeat([]byte("x"), 512)) + s.resolve(Request{Path: s.normalizePath(p)}) + + if e, ok := cachedEntry(s, p); !ok || !e.negative { + t.Fatalf("non-PE should be cached negative even when limiter is empty, got (%+v, %v)", e, ok) + } + }) + + t.Run("sync request bypasses the limiter", func(t *testing.T) { + s := newTestStore() + defer s.Close() + exhausted(s) + + p := writeTestFile(t, "b.exe", pe) + f := s.resolve(Request{Path: s.normalizePath(p), Sync: true}) + if f == nil || !f.IsExecutable { + t.Fatalf("sync request must not be rate limited, got %+v", f) + } + }) + + t.Run("seed request bypasses the limiter", func(t *testing.T) { + s := newTestStore() + defer s.Close() + exhausted(s) + + p := writeTestFile(t, "c.exe", pe) + f := s.resolve(Request{Path: s.normalizePath(p), IsSeed: true}) + if f == nil || !f.IsExecutable { + t.Fatalf("seed request must not be rate limited, got %+v", f) + } + }) +} + func TestGetFileType(t *testing.T) { var tests = []struct { filename string diff --git a/pkg/sys/fs.go b/pkg/sys/fs.go index bc776181b..0cde33f08 100644 --- a/pkg/sys/fs.go +++ b/pkg/sys/fs.go @@ -21,8 +21,8 @@ package sys import ( - "fmt" - "syscall" + "errors" + "runtime" "time" "golang.org/x/sys/windows" @@ -113,23 +113,37 @@ func GetMappedFile(process windows.Handle, addr uintptr) string { // elapsed, and the object's state is nonsignaled. const WaitTimeout = 0x00000102 +// ErrSkipped signals the file is offline or a cloud placeholder and was not read. +var ErrFileSkipped = errors.New("file is offline or a cloud placeholder") + +// ErrTimeout signals that the timeout elapsed while trying to read the file. +var ErrFileTimeout = errors.New("timeout reading file") + // ReadFile reads the file asynchronously with the specified number // of bytes to read and the timeout after which the I/O operation // is cancelled. func ReadFile(path string, size int, timeout time.Duration) ([]byte, error) { - pathUTF16, err := syscall.UTF16PtrFromString(path) + return ReadFileWithPredicate(path, size, 0, timeout, nil) +} + +// ReadFileWithPredicate opens the file once and reads its first probeSize bytes. If +// accept predicate is non-nil and returns an error, the read stops there and that +// error is returned. Otherwise the rest of the file is read, up to size bytes in total. +// The buffer is never larger than the file. +func ReadFileWithPredicate(path string, size, probeSize int, timeout time.Duration, accept func(head []byte) error) ([]byte, error) { + p, err := windows.UTF16PtrFromString(path) if err != nil { return nil, err } // open file asynchronously handle, err := windows.CreateFile( - pathUTF16, + p, windows.GENERIC_READ, - windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE, + windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE|windows.FILE_SHARE_DELETE, nil, windows.OPEN_EXISTING, - windows.FILE_FLAG_OVERLAPPED, + windows.FILE_FLAG_OVERLAPPED|windows.FILE_FLAG_OPEN_NO_RECALL, 0, ) if err != nil { @@ -138,51 +152,96 @@ func ReadFile(path string, size int, timeout time.Duration) ([]byte, error) { //nolint:errcheck defer windows.CloseHandle(handle) - event, err := windows.CreateEvent(nil, 1, 0, nil) + var info windows.ByHandleFileInformation + if err := windows.GetFileInformationByHandle(handle, &info); err != nil { + return nil, err + } + if info.FileAttributes&(windows.FILE_ATTRIBUTE_OFFLINE|windows.FILE_ATTRIBUTE_RECALL_ON_OPEN|windows.FILE_ATTRIBUTE_RECALL_ON_DATA_ACCESS) != 0 { + return nil, ErrFileSkipped + } + fsize := int64(info.FileSizeHigh)<<32 | int64(info.FileSizeLow) + total := size + if fsize < int64(total) { + total = int(fsize) + } + if total == 0 { + return []byte{}, nil + } + + ev, err := windows.CreateEvent(nil, 1, 0, nil) if err != nil { return nil, err } //nolint:errcheck - defer windows.CloseHandle(event) - - var overlapped windows.Overlapped - overlapped.HEvent = event + defer windows.CloseHandle(ev) - buf := make([]byte, size) // chunk size - var n uint32 + first := total + if accept != nil && probeSize > 0 && probeSize < first { + first = probeSize + } - err = windows.ReadFile(handle, buf, &n, &overlapped) - if err != nil && err != windows.ERROR_IO_PENDING { + head := make([]byte, first) + n, err := readFileAt(handle, ev, head, 0, timeout) + if err != nil { return nil, err } - // synchronous completion - if err == nil && int(n) <= size { - return buf[:n], nil - } - // wait for I/O operation to complete - wait, err := windows.WaitForSingleObject(event, uint32(timeout.Milliseconds())) - switch wait { - case windows.WAIT_FAILED: - return nil, err - case WaitTimeout: - // cancel the I/O - _ = windows.CancelIoEx(handle, &overlapped) - - // must wait until cancellation completes - _, _ = windows.WaitForSingleObject(event, 1000) - err = windows.GetOverlappedResult(handle, &overlapped, &n, false) - if err == windows.ERROR_OPERATION_ABORTED { - return nil, fmt.Errorf("timeout reading file %s", path) + head = head[:n] + if accept != nil { + if err := accept(head); err != nil { + return nil, err } - return nil, err + } + if total <= first || n < first { + return head, nil } - // get the result - err = windows.GetOverlappedResult(handle, &overlapped, &n, false) + buf := make([]byte, total) + copy(buf, head) + m, err := readFileAt(handle, ev, buf[n:], uint64(n), timeout) if err != nil { return nil, err } + return buf[:n+m], nil +} + +// readFileAt performs one overlapped read. It never returns while the kernel can +// still write into buf. +func readFileAt(h, ev windows.Handle, buf []byte, off uint64, timeout time.Duration) (int, error) { + var ov windows.Overlapped + ov.HEvent = ev + ov.Offset = uint32(off) + ov.OffsetHigh = uint32(off >> 32) + _ = windows.ResetEvent(ev) - return buf[:n], nil + var n uint32 + err := windows.ReadFile(h, buf, &n, &ov) + switch err { + case nil: + return int(n), nil + case windows.ERROR_HANDLE_EOF: + return 0, nil + case windows.ERROR_IO_PENDING: + default: + return 0, err + } + + wait, werr := windows.WaitForSingleObject(ev, uint32(timeout.Milliseconds())) + if wait != windows.WAIT_OBJECT_0 { + _ = windows.CancelIoEx(h, &ov) + // block until the kernel is done with buf + _ = windows.GetOverlappedResult(h, &ov, &n, true) + runtime.KeepAlive(buf) + if wait == WaitTimeout { + return 0, ErrFileTimeout + } + return 0, werr + } + if err := windows.GetOverlappedResult(h, &ov, &n, false); err != nil { + if err == windows.ERROR_HANDLE_EOF { + return 0, nil + } + return 0, err + } + return int(n), nil } diff --git a/pkg/sys/fs_test.go b/pkg/sys/fs_test.go index c6375c6de..883474dc3 100644 --- a/pkg/sys/fs_test.go +++ b/pkg/sys/fs_test.go @@ -21,41 +21,202 @@ package sys import ( + "errors" "os" "path/filepath" "testing" "time" "github.com/stretchr/testify/require" + "golang.org/x/sys/windows" ) +// payload returns n deterministic, non-repeating-looking bytes so that +// offset mistakes between the probe read and the remainder read show up. +func payload(n int) []byte { + b := make([]byte, n) + for i := range b { + b[i] = byte(i % 251) // prime modulus avoids aligning with power-of-two buffers + } + return b +} + +func writeTemp(t *testing.T, content []byte) string { + t.Helper() + path := filepath.Join(t.TempDir(), "test.bin") + require.NoError(t, os.WriteFile(path, content, 0644)) + return path +} + func TestReadFile(t *testing.T) { var tests = []struct { - name string - f func() (string, error) - err error - b []byte + name string + content []byte + missing bool + size int + err error + b []byte }{ - { - "read file", - func() (string, error) { - path := filepath.Join(t.TempDir(), "test.txt") + {"read file", []byte("hello world"), false, 4096, nil, []byte("hello world")}, + {"empty file", []byte{}, false, 4096, nil, []byte{}}, + {"size smaller than file", []byte("hello world"), false, 5, nil, []byte("hello")}, + {"size equal to file", []byte("hello world"), false, 11, nil, []byte("hello world")}, + {"file larger than 1 MB", payload(1<<20 + 17), false, 4 << 20, nil, payload(1<<20 + 17)}, + {"missing file", nil, true, 4096, windows.ERROR_FILE_NOT_FOUND, nil}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "missing.bin") + if !tt.missing { + path = writeTemp(t, tt.content) + } + b, err := ReadFile(path, tt.size, time.Second) + if tt.err != nil { + require.ErrorIs(t, err, tt.err) + } else { + require.NoError(t, err) + } + require.Equal(t, tt.b, b) + }) + } +} + +func TestReadFileWithPredicate(t *testing.T) { + errRejected := errors.New("rejected by accept") + + var tests = []struct { + name string + content []byte + missing bool + size int + probe int + + // accept behaviour: nil accept disables the callback entirely + useAccept bool + acceptErr error - return path, os.WriteFile(path, []byte("hello world"), 0644) - }, - nil, - []byte{104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100}, + err error + b []byte + // expected len(head) passed to accept, -1 if accept must not be called + headLen int + }{ + { + name: "no accept reads whole file", content: payload(5000), + size: 1 << 20, probe: 1024, + b: payload(5000), headLen: -1, + }, + { + name: "accept sees only probe bytes of a larger file", content: payload(5000), + size: 1 << 20, probe: 1024, useAccept: true, + b: payload(5000), headLen: 1024, + }, + { + name: "remainder read continues at the right offset", content: payload(3<<20 + 123), + size: 8 << 20, probe: 1024, useAccept: true, + b: payload(3<<20 + 123), headLen: 1024, + }, + { + name: "probe larger than file", content: payload(300), + size: 1 << 20, probe: 1024, useAccept: true, + b: payload(300), headLen: 300, + }, + { + name: "probe equal to file", content: payload(1024), + size: 1 << 20, probe: 1024, useAccept: true, + b: payload(1024), headLen: 1024, + }, + { + name: "zero probe size means single full read", content: payload(5000), + size: 1 << 20, probe: 0, useAccept: true, + b: payload(5000), headLen: 5000, + }, + { + name: "size caps the total", content: payload(5000), + size: 2000, probe: 1024, useAccept: true, + b: payload(2000), headLen: 1024, + }, + { + name: "size smaller than probe", content: payload(5000), + size: 100, probe: 1024, useAccept: true, + b: payload(100), headLen: 100, + }, + { + name: "accept error stops the read", content: payload(5000), + size: 1 << 20, probe: 1024, useAccept: true, acceptErr: errRejected, + err: errRejected, b: nil, headLen: 1024, + }, + { + name: "accept error on a file that fits in the probe", content: payload(100), + size: 1 << 20, probe: 1024, useAccept: true, acceptErr: errRejected, + err: errRejected, b: nil, headLen: 100, + }, + { + // empty files short-circuit before any read, so accept never runs. + // Callers must therefore not rely on accept to detect empty files. + name: "empty file skips accept", content: []byte{}, + size: 1 << 20, probe: 1024, useAccept: true, + b: []byte{}, headLen: -1, + }, + { + name: "missing file", missing: true, + size: 1 << 20, probe: 1024, useAccept: true, + err: windows.ERROR_FILE_NOT_FOUND, b: nil, headLen: -1, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - path, err := tt.f() - require.NoError(t, err) - defer os.Remove(path) - b, err := ReadFile(path, 4096, time.Second*1) - require.Equal(t, tt.err, err) + path := filepath.Join(t.TempDir(), "missing.bin") + if !tt.missing { + path = writeTemp(t, tt.content) + } + + var ( + calls int + head []byte + ) + var accept func([]byte) error + if tt.useAccept { + accept = func(h []byte) error { + calls++ + head = append([]byte(nil), h...) // copy: h aliases the read buffer + return tt.acceptErr + } + } + + b, err := ReadFileWithPredicate(path, tt.size, tt.probe, time.Second, accept) + if tt.err != nil { + require.ErrorIs(t, err, tt.err) + } else { + require.NoError(t, err) + } require.Equal(t, tt.b, b) + + if tt.headLen < 0 { + require.Zero(t, calls, "accept must not be called") + return + } + require.Equal(t, 1, calls, "accept must be called exactly once") + require.Len(t, head, tt.headLen) + // the head must be a prefix of the file, not garbage from a later offset + require.Equal(t, tt.content[:tt.headLen], head) }) } } + +// A file that is open for exclusive writing elsewhere must still be readable, +// because the store opens with FILE_SHARE_READ|WRITE|DELETE. This mirrors the +// file-create-event case where the writer is still active. +func TestReadFileProbeWhileWriterHoldsHandle(t *testing.T) { + path := filepath.Join(t.TempDir(), "writer.bin") + f, err := os.Create(path) + require.NoError(t, err) + defer f.Close() + _, err = f.Write(payload(2048)) + require.NoError(t, err) + + b, err := ReadFileWithPredicate(path, 1<<20, 1024, time.Second, func([]byte) error { return nil }) + require.NoError(t, err) + require.Equal(t, payload(2048), b) +}