1package hints
2
3import (
4 "bufio"
5 "encoding/json"
6 "fmt"
7 "image"
8 "image/color"
9 "image/png"
10 "math"
11 "os"
12 "path/filepath"
13 "sort"
14 "strings"
15 "syscall"
16 "time"
17
18 "github.com/bnema/neferwl/internal/ports"
19)
20
21const (
22 // maxDumps is the number of dump groups kept in the directory.
23 maxDumps = 50
24 outlineSize = 2
25)
26
27// Options tunes Run.
28type Options struct {
29 // Command is the nefer-hints binary: a path, or a name looked up in PATH.
30 Command string
31 // ModelDir is the directory of the detector's model files.
32 ModelDir string
33 // Threshold overrides the minimum detection score, 0.01 to 0.99.
34 // Zero uses model.json or, if absent, the runner default.
35 Threshold float64
36 // DumpDir, when not empty, receives the scanned frames (PNG, annotated
37 // PNG and JSON). They contain the screen in clear.
38 DumpDir string
39
40 // detector replaces the subprocess in tests.
41 detector detector
42}
43
44// dumpJob is one scan to write. pix is the capture buffer of the scan, owned
45// by the job: writeDump converts it in place.
46type dumpJob struct {
47 when time.Time
48 req ports.HintScanRequest
49 pix []byte
50 m *mapping // owns pix; released after the dump
51 stride int
52 targets []ports.HintTarget
53 dets []detection // detector boxes, physical pixels
54 captureD time.Duration
55 detectD time.Duration
56}
57
58type dumpTarget struct {
59 Rect [4]int `json:"rect"` // x1, y1, x2, y2, logical output-local
60 Score float32 `json:"score"`
61}
62
63type dumpMeta struct {
64 Output string `json:"output"`
65 Width int `json:"width"`
66 Height int `json:"height"`
67 Stride int `json:"stride"`
68 Scale float64 `json:"scale"`
69 CaptureMS float64 `json:"capture_ms"`
70 DetectMS float64 `json:"detect_ms"`
71 TotalMS float64 `json:"total_ms"`
72 Targets []dumpTarget `json:"targets"`
73}
74
75// boxColor is the outline of the detector boxes.
76var boxColor = color.RGBA{0x00, 0xc8, 0x53, 0xff}
77
78func ms(d time.Duration) float64 { return float64(d.Microseconds()) / 1000 }
79
80// writeDump writes the three files of j into dir, then prunes old groups.
81func writeDump(dir string, j *dumpJob) error {
82 if err := secureDir(dir); err != nil {
83 return err
84 }
85 prefix := filepath.Join(dir, j.when.Format("20060102-150405.000")+"-"+safeName(j.req.Output))
86
87 // The frame is converted in place to RGBA and encoded straight from the
88 // capture buffer: no copy. The job owns the buffer, which is released
89 // with it.
90 for y := 0; y < j.req.Height; y++ {
91 row := j.pix[y*j.stride : y*j.stride+j.req.Width*4]
92 for x := 0; x < len(row); x += 4 {
93 row[x], row[x+2], row[x+3] = row[x+2], row[x], 0xff
94 }
95 }
96 img := &image.RGBA{Pix: j.pix, Stride: j.stride, Rect: image.Rect(0, 0, j.req.Width, j.req.Height)}
97 if err := writePNG(prefix+".png", img); err != nil {
98 return err
99 }
100
101 for _, d := range j.dets {
102 outline(img, image.Rect(int(d.x1), int(d.y1), int(math.Ceil(float64(d.x2))), int(math.Ceil(float64(d.y2)))), boxColor)
103 }
104 if err := writePNG(prefix+".annotated.png", img); err != nil {
105 return err
106 }
107
108 meta := dumpMeta{
109 Output: j.req.Output, Width: j.req.Width, Height: j.req.Height, Stride: j.stride,
110 Scale: j.req.Scale,
111 CaptureMS: ms(j.captureD), DetectMS: ms(j.detectD), TotalMS: ms(j.captureD + j.detectD),
112 Targets: make([]dumpTarget, 0, len(j.targets)),
113 }
114 for _, t := range j.targets {
115 meta.Targets = append(meta.Targets, dumpTarget{Rect: [4]int{t.Rect.X, t.Rect.Y, t.Rect.X + t.Rect.W, t.Rect.Y + t.Rect.H}, Score: t.Score})
116 }
117 data, err := json.Marshal(meta)
118 if err != nil {
119 return err
120 }
121 if err := writeFile(prefix+".json", func(f *os.File) error { _, err := f.Write(append(data, '\n')); return err }); err != nil {
122 return err
123 }
124 return prune(dir, maxDumps)
125}
126
127func safeName(s string) string {
128 return strings.Map(func(r rune) rune {
129 if r == '/' || r == 0 {
130 return '_'
131 }
132 return r
133 }, s)
134}
135
136func writePNG(path string, img image.Image) error {
137 enc := png.Encoder{CompressionLevel: png.BestSpeed}
138 return writeFile(path, func(f *os.File) error {
139 bw := bufio.NewWriterSize(f, 1<<16)
140 if err := enc.Encode(bw, img); err != nil {
141 return err
142 }
143 return bw.Flush()
144 })
145}
146
147// secureDir makes sure dir is a real directory owned by the current user and
148// private to it, whatever it was before.
149func secureDir(dir string) error {
150 if err := os.MkdirAll(dir, 0o700); err != nil {
151 return err
152 }
153 fi, err := os.Lstat(dir)
154 if err != nil {
155 return err
156 }
157 if !fi.IsDir() {
158 return fmt.Errorf("dump dir %s is not a directory", dir)
159 }
160 if st, ok := fi.Sys().(*syscall.Stat_t); ok && int(st.Uid) != os.Getuid() {
161 return fmt.Errorf("dump dir %s is not owned by the current user", dir)
162 }
163 return os.Chmod(dir, 0o700)
164}
165
166// writeFile fills a fresh private temporary file in the directory of path and
167// renames it over path: an existing file or symlink there is replaced, never
168// followed.
169func writeFile(path string, fill func(*os.File) error) error {
170 f, err := os.CreateTemp(filepath.Dir(path), ".dump-*")
171 if err != nil {
172 return err
173 }
174 tmp := f.Name()
175 err = f.Chmod(0o600)
176 if err == nil {
177 err = fill(f)
178 }
179 if cerr := f.Close(); err == nil {
180 err = cerr
181 }
182 if err == nil {
183 err = os.Rename(tmp, path)
184 }
185 if err != nil {
186 _ = os.Remove(tmp)
187 }
188 return err
189}
190
191// outline draws a 2 px frame inside r, clipped to img.
192func outline(img *image.RGBA, r image.Rectangle, c color.RGBA) {
193 r = r.Intersect(img.Rect)
194 if r.Empty() {
195 return
196 }
197 for y := r.Min.Y; y < r.Max.Y; y++ {
198 for x := r.Min.X; x < r.Max.X; x++ {
199 if x-r.Min.X < outlineSize || r.Max.X-1-x < outlineSize || y-r.Min.Y < outlineSize || r.Max.Y-1-y < outlineSize {
200 img.SetRGBA(x, y, c)
201 }
202 }
203 }
204}
205
206// groupOf returns the "<timestamp>-<output>" prefix of a dump file name.
207func groupOf(name string) string {
208 for _, suf := range []string{".annotated.png", ".png", ".json"} {
209 if strings.HasSuffix(name, suf) {
210 return strings.TrimSuffix(name, suf)
211 }
212 }
213 return ""
214}
215
216// prune keeps the keep most recent groups of dir; timestamps sort by name.
217func prune(dir string, keep int) error {
218 entries, err := os.ReadDir(dir)
219 if err != nil {
220 return err
221 }
222 groups := map[string][]string{}
223 for _, e := range entries {
224 if g := groupOf(e.Name()); g != "" && !e.IsDir() {
225 groups[g] = append(groups[g], e.Name())
226 }
227 }
228 if len(groups) <= keep {
229 return nil
230 }
231 names := make([]string, 0, len(groups))
232 for g := range groups {
233 names = append(names, g)
234 }
235 sort.Strings(names)
236 for _, g := range names[:len(names)-keep] {
237 for _, f := range groups[g] {
238 if err := os.Remove(filepath.Join(dir, f)); err != nil {
239 return fmt.Errorf("prune: %w", err)
240 }
241 }
242 }
243 return nil
244}