| 1 | package dbreads |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "fmt" |
| 8 | "io" |
| 9 | "os" |
| 10 | "path/filepath" |
| 11 | "strings" |
| 12 | "time" |
| 13 | |
| 14 | "github.com/benbjohnson/litestream" |
| 15 | lsfile "github.com/benbjohnson/litestream/file" |
| 16 | "github.com/superfly/ltx" |
| 17 | |
| 18 | "github.com/carlosframework/platform/internal/dbdoor" |
| 19 | "github.com/carlosframework/platform/internal/dbquery" |
| 20 | ) |
| 21 | |
| 22 | // WorkerOptions are ServeWorker's settings. Production (RunWorker) sets |
| 23 | // Scratch and leaves the rest zero. |
| 24 | type WorkerOptions struct { |
| 25 | // Scratch is where the read's private directory goes: /var/tmp, the |
| 26 | // unit's own PrivateTmp, which systemd removes at stop. |
| 27 | Scratch string |
| 28 | // QueryTimeout is QueryTimeout when zero; Grace is WorkerExitGrace. |
| 29 | QueryTimeout time.Duration |
| 30 | Grace time.Duration |
| 31 | // Exit ends the process (os.Exit when nil): the 5 s backstop for a step |
| 32 | // that ignores its context (spec §8.3, codex 8). |
| 33 | Exit func(code int) |
| 34 | // Self is the hello's view of the sandbox (SelfView when nil). |
| 35 | Self func() HelloReply |
| 36 | // Restore is dbdoor.RestoreLocal when nil. Tests replace it. |
| 37 | Restore func(ctx context.Context, c litestream.ReplicaClient, out string, target uint64) error |
| 38 | // MaxRestored is the largest restored-size cap a header may claim: |
| 39 | // DefaultMaxRestored when zero, and in the unit the process's own |
| 40 | // RLIMIT_FSIZE (LimitFSIZE), which a box raising the cap raises too. The |
| 41 | // header's files may total at most PlanBytesFactor times its cap. |
| 42 | MaxRestored int64 |
| 43 | } |
| 44 | |
| 45 | type frame struct { |
| 46 | typ byte |
| 47 | p []byte |
| 48 | } |
| 49 | |
| 50 | var errAgentGone = errors.New("dbreads: the agent closed the connection") |
| 51 | |
| 52 | // ServeWorker is the worker's whole job on one connection (spec §8.3): a |
| 53 | // hello, or a plan, its files, a restore, ready, then run or stop. One |
| 54 | // goroutine reads the connection from the start; EOF or an error cancels |
| 55 | // the one context everything runs under, and if the work has not returned |
| 56 | // Grace later, the process exits. |
| 57 | func ServeWorker(conn io.ReadWriter, opt WorkerOptions) error { |
| 58 | if opt.QueryTimeout == 0 { |
| 59 | opt.QueryTimeout = QueryTimeout |
| 60 | } |
| 61 | if opt.Grace == 0 { |
| 62 | opt.Grace = WorkerExitGrace |
| 63 | } |
| 64 | if opt.Exit == nil { |
| 65 | opt.Exit = os.Exit |
| 66 | } |
| 67 | if opt.Self == nil { |
| 68 | opt.Self = SelfView |
| 69 | } |
| 70 | if opt.Restore == nil { |
| 71 | opt.Restore = dbdoor.RestoreLocal |
| 72 | } |
| 73 | if opt.MaxRestored <= 0 { |
| 74 | opt.MaxRestored = DefaultMaxRestored |
| 75 | } |
| 76 | ctx, cancel := context.WithCancel(context.Background()) |
| 77 | defer cancel() |
| 78 | frames := make(chan frame) |
| 79 | done := make(chan struct{}) |
| 80 | defer close(done) |
| 81 | go func() { |
| 82 | defer close(frames) |
| 83 | for { |
| 84 | typ, p, err := readFrame(conn) |
| 85 | if err != nil { |
| 86 | cancel() |
| 87 | select { |
| 88 | case <-done: |
| 89 | case <-time.After(opt.Grace): |
| 90 | opt.Exit(3) |
| 91 | } |
| 92 | return |
| 93 | } |
| 94 | select { |
| 95 | case frames <- frame{typ, p}: |
| 96 | case <-done: |
| 97 | return |
| 98 | } |
| 99 | } |
| 100 | }() |
| 101 | next := func() (frame, error) { |
| 102 | select { |
| 103 | case f, ok := <-frames: |
| 104 | if !ok { |
| 105 | return frame{}, errAgentGone |
| 106 | } |
| 107 | return f, nil |
| 108 | case <-ctx.Done(): |
| 109 | return frame{}, errAgentGone |
| 110 | } |
| 111 | } |
| 112 | end := func(code string, err error) error { |
| 113 | msg := endMsg{Code: code} |
| 114 | if err != nil { |
| 115 | msg.Error = err.Error() |
| 116 | } |
| 117 | return writeJSONFrame(conn, frameEnd, msg) |
| 118 | } |
| 119 | |
| 120 | f, err := next() |
| 121 | if err != nil { |
| 122 | return err |
| 123 | } |
| 124 | switch f.typ { |
| 125 | case frameHello: |
| 126 | var h hello |
| 127 | if json.Unmarshal(f.p, &h) != nil { |
| 128 | return errProtocol |
| 129 | } |
| 130 | return writeJSONFrame(conn, frameHelloReply, opt.Self()) |
| 131 | case frameHeader: |
| 132 | default: |
| 133 | return errProtocol |
| 134 | } |
| 135 | var h planHeader |
| 136 | if err := json.Unmarshal(f.p, &h); err != nil || h.V != ProtocolVersion || (h.Op != OpQuery && h.Op != OpPull) { |
| 137 | return errProtocol |
| 138 | } |
| 139 | // Every byte this worker writes is bounded before the first is |
| 140 | // written: each file by its header size (receiveFile), and all of them |
| 141 | // together by PlanBytesFactor times a cap no larger than the worker's |
| 142 | // own (Astra round 1, finding 1). LimitFSIZE alone bounds one file. |
| 143 | if err := planWithinBounds(h, opt.MaxRestored); err != nil { |
| 144 | return end(CodeCheckFailed, err) |
| 145 | } |
| 146 | dir, err := os.MkdirTemp(opt.Scratch, "dbread-") |
| 147 | if err != nil { |
| 148 | return err |
| 149 | } |
| 150 | defer os.RemoveAll(dir) |
| 151 | root := filepath.Join(dir, "replica") |
| 152 | for _, lf := range h.Files { |
| 153 | path := litestream.LTXFilePath(root, lf.Level, ltx.TXID(lf.MinTXID), ltx.TXID(lf.MaxTXID)) |
| 154 | if err := receiveFile(next, path, lf.Size); err != nil { |
| 155 | if errors.Is(err, errAgentGone) { |
| 156 | return err |
| 157 | } |
| 158 | return end(CodeCheckFailed, err) |
| 159 | } |
| 160 | // litestream's restore picks files by their names and compacts them |
| 161 | // by their headers, so a file whose header is not its plan entry |
| 162 | // would be reported at a transaction it does not hold (Astra round 1 |
| 163 | // on plan 3, finding 1). A pull's files go back unrestored, so each is |
| 164 | // decoded whole as well. |
| 165 | if err := verifyLTX(path, lf, h.Op == OpPull); err != nil { |
| 166 | return end(CodeCheckFailed, err) |
| 167 | } |
| 168 | } |
| 169 | if len(h.Files) == 0 || h.Files[len(h.Files)-1].MaxTXID != h.Target { |
| 170 | return end(CodeCheckFailed, errors.New("dbreads: the plan does not end at its target")) |
| 171 | } |
| 172 | if h.Op == OpPull { |
| 173 | return holdAndSend(ctx, conn, next, end, root, h) |
| 174 | } |
| 175 | db := filepath.Join(dir, "copy.db") |
| 176 | if err := opt.Restore(ctx, lsfile.NewReplicaClient(root), db, h.Target); err != nil { |
| 177 | if ctx.Err() != nil { |
| 178 | return errAgentGone |
| 179 | } |
| 180 | return end(CodeCheckFailed, err) |
| 181 | } |
| 182 | if st, err := os.Stat(db); err != nil || st.Size() > h.MaxRestored { |
| 183 | return end(CodeTooLarge, err) |
| 184 | } |
| 185 | if err := writeJSONFrame(conn, frameReady, readyMsg{TXID: h.Target}); err != nil { |
| 186 | return err |
| 187 | } |
| 188 | f, err = next() |
| 189 | if err != nil { |
| 190 | return err |
| 191 | } |
| 192 | switch f.typ { |
| 193 | case frameStop: |
| 194 | return nil |
| 195 | case frameRun: |
| 196 | default: |
| 197 | return errProtocol |
| 198 | } |
| 199 | var run runMsg |
| 200 | if err := json.Unmarshal(f.p, &run); err != nil { |
| 201 | return errProtocol |
| 202 | } |
| 203 | qctx, qcancel := context.WithTimeout(ctx, opt.QueryTimeout) |
| 204 | defer qcancel() |
| 205 | out := &frameWriter{w: conn} |
| 206 | qerr := dbquery.RunContext(qctx, out, db, run.SQL, dbquery.Options{Immutable: true}) |
| 207 | if ferr := out.flush(); ferr != nil { |
| 208 | return ferr |
| 209 | } |
| 210 | switch { |
| 211 | case qerr == nil: |
| 212 | return end("", nil) |
| 213 | case errors.Is(qctx.Err(), context.DeadlineExceeded): |
| 214 | return end(CodeTimeout, nil) |
| 215 | case ctx.Err() != nil: |
| 216 | return errAgentGone |
| 217 | } |
| 218 | return end(CodeSQLError, qerr) |
| 219 | } |
| 220 | |
| 221 | // holdAndSend is a pull's whole worker job after the files are received and |
| 222 | // checked (spec §6: the box never restores for a pull): the files are held |
| 223 | // in this unit's private scratch, and after the agent's check B they go back |
| 224 | // in the header's order, whole, as output. |
| 225 | func holdAndSend(ctx context.Context, conn io.Writer, next func() (frame, error), end func(string, error) error, root string, h planHeader) error { |
| 226 | // ServeWorker has checked every file whole against its plan entry, and |
| 227 | // the plan's end against its target. |
| 228 | paths := make([]string, len(h.Files)) |
| 229 | for i, lf := range h.Files { |
| 230 | paths[i] = litestream.LTXFilePath(root, lf.Level, ltx.TXID(lf.MinTXID), ltx.TXID(lf.MaxTXID)) |
| 231 | } |
| 232 | if err := writeJSONFrame(conn, frameReady, readyMsg{TXID: h.Target}); err != nil { |
| 233 | return err |
| 234 | } |
| 235 | f, err := next() |
| 236 | if err != nil { |
| 237 | return err |
| 238 | } |
| 239 | switch f.typ { |
| 240 | case frameStop: |
| 241 | return nil |
| 242 | case frameRun: |
| 243 | default: |
| 244 | return errProtocol |
| 245 | } |
| 246 | out := &frameWriter{w: conn} |
| 247 | for _, p := range paths { |
| 248 | if ctx.Err() != nil { |
| 249 | return errAgentGone |
| 250 | } |
| 251 | in, err := os.Open(p) |
| 252 | if err != nil { |
| 253 | return end(CodeCheckFailed, err) |
| 254 | } |
| 255 | _, err = io.Copy(out, in) |
| 256 | in.Close() |
| 257 | if err != nil { |
| 258 | return err |
| 259 | } |
| 260 | } |
| 261 | if err := out.flush(); err != nil { |
| 262 | return err |
| 263 | } |
| 264 | return end("", nil) |
| 265 | } |
| 266 | |
| 267 | // verifyLTX checks one received file against its plan entry: its LTX |
| 268 | // header must name exactly the entry's transactions. With whole, every page |
| 269 | // and the checksums are decoded too (the decode is the worker's, as every |
| 270 | // decode of tenant bytes is). |
| 271 | func verifyLTX(path string, lf dbdoor.LTXFile, whole bool) error { |
| 272 | f, err := os.Open(path) |
| 273 | if err != nil { |
| 274 | return err |
| 275 | } |
| 276 | defer f.Close() |
| 277 | dec := ltx.NewDecoder(f) |
| 278 | if whole { |
| 279 | err = dec.Verify() |
| 280 | } else { |
| 281 | err = dec.DecodeHeader() |
| 282 | } |
| 283 | if err != nil { |
| 284 | return err |
| 285 | } |
| 286 | if h := dec.Header(); uint64(h.MinTXID) != lf.MinTXID || uint64(h.MaxTXID) != lf.MaxTXID { |
| 287 | return fmt.Errorf("dbreads: %s holds transactions %s", filepath.Base(path), dbdoor.Span(uint64(h.MinTXID), uint64(h.MaxTXID))) |
| 288 | } |
| 289 | return nil |
| 290 | } |
| 291 | |
| 292 | // receiveFile writes exactly size bytes of file frames to path, refusing a |
| 293 | // file larger than its header said (spec §8.3). |
| 294 | func receiveFile(next func() (frame, error), path string, size int64) error { |
| 295 | if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { |
| 296 | return err |
| 297 | } |
| 298 | out, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600) |
| 299 | if err != nil { |
| 300 | return err |
| 301 | } |
| 302 | defer out.Close() |
| 303 | var got int64 |
| 304 | for got < size { |
| 305 | f, err := next() |
| 306 | if err != nil { |
| 307 | return err |
| 308 | } |
| 309 | if f.typ != frameFile { |
| 310 | return errProtocol |
| 311 | } |
| 312 | if got+int64(len(f.p)) > size { |
| 313 | return fmt.Errorf("dbreads: %s is larger than its header said (%d bytes)", filepath.Base(path), size) |
| 314 | } |
| 315 | if _, err := out.Write(f.p); err != nil { |
| 316 | return err |
| 317 | } |
| 318 | got += int64(len(f.p)) |
| 319 | } |
| 320 | return out.Close() |
| 321 | } |
| 322 | |
| 323 | // frameWriter frames the query's output in outputFrameBytes pieces. |
| 324 | type frameWriter struct { |
| 325 | w io.Writer |
| 326 | buf []byte |
| 327 | } |
| 328 | |
| 329 | func (f *frameWriter) Write(p []byte) (int, error) { |
| 330 | f.buf = append(f.buf, p...) |
| 331 | for len(f.buf) >= outputFrameBytes { |
| 332 | if err := writeFrame(f.w, frameOutput, f.buf[:outputFrameBytes]); err != nil { |
| 333 | return 0, err |
| 334 | } |
| 335 | f.buf = f.buf[outputFrameBytes:] |
| 336 | } |
| 337 | return len(p), nil |
| 338 | } |
| 339 | |
| 340 | func (f *frameWriter) flush() error { |
| 341 | if len(f.buf) == 0 { |
| 342 | return nil |
| 343 | } |
| 344 | err := writeFrame(f.w, frameOutput, f.buf) |
| 345 | f.buf = nil |
| 346 | return err |
| 347 | } |
| 348 | |
| 349 | // SelfView is what the worker can see of the box: in the unit, the loopback |
| 350 | // interface alone, an empty /data and no readable /etc/carlos. |
| 351 | // |
| 352 | // Interfaces come from /proc/self/net/dev, this process's own network |
| 353 | // namespace. net.Interfaces cannot be used: on Linux it opens an AF_NETLINK |
| 354 | // socket, which the unit's RestrictAddressFamilies=AF_UNIX refuses, and an |
| 355 | // empty list would read as "no network" (plan review finding 1). Every |
| 356 | // observation that fails is reported in Errors, and a probe fails on any. |
| 357 | func SelfView() HelloReply { |
| 358 | h := HelloReply{V: ProtocolVersion} |
| 359 | if raw, err := os.ReadFile("/proc/self/net/dev"); err != nil { |
| 360 | h.Errors = append(h.Errors, "interfaces: "+err.Error()) |
| 361 | } else { |
| 362 | h.Interfaces = parseNetDev(string(raw)) |
| 363 | } |
| 364 | switch ents, err := os.ReadDir("/data"); { |
| 365 | case err == nil: |
| 366 | h.DataEntries = len(ents) |
| 367 | case !errors.Is(err, os.ErrNotExist): |
| 368 | h.Errors = append(h.Errors, "/data: "+err.Error()) |
| 369 | } |
| 370 | if f, err := os.Open("/etc/carlos/host.env"); err == nil { |
| 371 | f.Close() |
| 372 | h.EtcReadable = true |
| 373 | } |
| 374 | return h |
| 375 | } |
| 376 | |
| 377 | // parseNetDev reads /proc/net/dev's interface names: two header lines, then |
| 378 | // one "<name>: <counters>" line per interface. |
| 379 | func parseNetDev(s string) []string { |
| 380 | var out []string |
| 381 | for i, line := range strings.Split(s, "\n") { |
| 382 | if i < 2 { |
| 383 | continue |
| 384 | } |
| 385 | name, _, ok := strings.Cut(line, ":") |
| 386 | if name = strings.TrimSpace(name); ok && name != "" { |
| 387 | out = append(out, name) |
| 388 | } |
| 389 | } |
| 390 | return out |
| 391 | } |
| 392 | |
| 393 | // planWithinBounds refuses a header whose restored-size cap is not in |
| 394 | // (0, limit], or whose files have a negative size or total more than |
| 395 | // PlanBytesFactor times that cap. |
| 396 | func planWithinBounds(h planHeader, limit int64) error { |
| 397 | if h.MaxRestored <= 0 || h.MaxRestored > limit { |
| 398 | return fmt.Errorf("dbreads: the plan's restored-size cap %d is outside (0, %d]", h.MaxRestored, limit) |
| 399 | } |
| 400 | var total int64 |
| 401 | for _, f := range h.Files { |
| 402 | if f.Size < 0 || f.Size > PlanBytesFactor*h.MaxRestored-total { |
| 403 | return fmt.Errorf("dbreads: the plan's files exceed %d bytes", PlanBytesFactor*h.MaxRestored) |
| 404 | } |
| 405 | total += f.Size |
| 406 | } |
| 407 | return nil |
| 408 | } |
| 409 | |
amadan