1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
package eval

import (
	"archive/tar"
	"compress/bzip2"
	"compress/gzip"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"strings"
)

// Realizer is an optional interface that Store implementations may provide
// to support import-from-derivation (IFD).  When readFile or import
// encounters a /nix/store path that doesn't exist on disk, the evaluator
// calls EnsureRealized to build it via the Nix daemon.
type Realizer interface {
	EnsureRealized(storePath string) error
}

// DaemonStore wraps MemStore and adds IFD support by connecting to the
// Nix daemon on demand.  All hash-computation methods are delegated to the
// embedded MemStore; EnsureRealized handles the actual build step.
type DaemonStore struct {
	*MemStore

	socketPath string
	daemon     *DaemonConn // nil until first IFD

	// registered tracks .drv store paths already sent to the daemon in this
	// session, to avoid redundant AddToStoreNar calls.
	registered map[string]bool

	// outputToDrv is a reverse index: output store path → drv store path.
	// Populated in RecordDerivation.
	outputToDrv map[string]string
}

// NewDaemonStore creates a DaemonStore that will connect to the Nix daemon at
// socketPath when the first IFD is triggered.
func NewDaemonStore(socketPath string) *DaemonStore {
	return &DaemonStore{
		MemStore:    NewMemStore(),
		socketPath:  socketPath,
		registered:  make(map[string]bool),
		outputToDrv: make(map[string]string),
	}
}

// RecordDerivation overrides MemStore to also populate the outputToDrv index.
func (ds *DaemonStore) RecordDerivation(drvPath string, drv *Derivation) {
	ds.MemStore.RecordDerivation(drvPath, drv)
	for _, out := range drv.Outputs {
		if out.Path != "" {
			ds.outputToDrv[out.Path] = drvPath
		}
	}
}

// EnsureRealized checks whether storePath exists on disk; if not, it finds
// the derivation that produces it, registers the full .drv closure with the
// daemon, and asks the daemon to build it.
func (ds *DaemonStore) EnsureRealized(storePath string) error {
	// Fast path: already on disk.
	if _, err := os.Stat(storePath); err == nil {
		return nil
	}

	// Lazily connect to the daemon.
	if err := ds.ensureConnected(); err != nil {
		return fmt.Errorf("EnsureRealized: %w", err)
	}

	// Find the .drv that produces this output.
	drvPath, ok := ds.outputToDrv[storePath]
	if !ok {
		// No derivation recorded — this is a fixed-output path (e.g. from
		// fetchTarball/fetchurl). Ask the daemon to substitute it directly.
		if err := ds.daemon.EnsurePath(storePath); err != nil {
			return fmt.Errorf("EnsureRealized: substitute %s: %w", storePath, err)
		}
	} else {
		// Walk the full .drv closure and register each one with the daemon.
		if err := ds.registerClosure(drvPath); err != nil {
			return fmt.Errorf("EnsureRealized: register closure: %w", err)
		}

		// Ask the daemon to build it.
		if err := ds.daemon.BuildPaths([]string{drvPath}); err != nil {
			return fmt.Errorf("EnsureRealized: build %s: %w", drvPath, err)
		}
	}

	// Verify it now exists.
	if _, err := os.Stat(storePath); err != nil {
		return fmt.Errorf("EnsureRealized: build succeeded but %s still missing: %w", storePath, err)
	}
	return nil
}

// registerClosure recursively registers drvPath and all its InputDrvs with
// the daemon via AddToStoreNar, in a post-order traversal (inputs before
// dependents).  Already-registered paths are skipped.
func (ds *DaemonStore) registerClosure(drvPath string) error {
	if ds.registered[drvPath] {
		return nil
	}

	drv := ds.MemStore.GetDerivation(drvPath)
	if drv == nil {
		// We don't have this derivation recorded — it may already be in the
		// daemon's store (e.g. a nixpkgs derivation built in a previous run).
		// Skip it; the daemon will handle it or fail with a clear error.
		return nil
	}

	// Recurse into InputDrvs first (post-order).
	inputDrvPaths := make([]string, 0, len(drv.InputDrvs))
	for p := range drv.InputDrvs {
		inputDrvPaths = append(inputDrvPaths, p)
	}
	// Sort for determinism (not strictly required, but helpful for debugging).
	for i := 0; i < len(inputDrvPaths); i++ {
		for j := i + 1; j < len(inputDrvPaths); j++ {
			if inputDrvPaths[i] > inputDrvPaths[j] {
				inputDrvPaths[i], inputDrvPaths[j] = inputDrvPaths[j], inputDrvPaths[i]
			}
		}
	}
	for _, dep := range inputDrvPaths {
		if err := ds.registerClosure(dep); err != nil {
			return err
		}
	}

	// Get the ATerm for this .drv.
	aterm, ok := ds.MemStore.ATermFor(drvPath)
	if !ok {
		// Same as above — not recorded, assume daemon has it.
		return nil
	}

	// Compute NAR and hash.
	narData := NARFile([]byte(aterm))
	narHash := NARFileHash([]byte(aterm))

	// Collect references: InputDrvs paths + InputSrcs paths.
	refs := drvReferences(drv)

	if err := ds.daemon.AddToStoreNar(drvPath, narHash, refs, narData); err != nil {
		return fmt.Errorf("register %s: %w", drvPath, err)
	}
	ds.registered[drvPath] = true
	return nil
}

// drvReferences returns all store paths referenced by a .drv file —
// i.e. the union of its InputDrvs keys and InputSrcs values.
// These become the "references" field in AddToStoreNar.
func drvReferences(drv *Derivation) []string {
	seen := make(map[string]bool)
	var refs []string
	add := func(p string) {
		if !seen[p] && strings.HasPrefix(p, "/nix/store/") {
			seen[p] = true
			refs = append(refs, p)
		}
	}
	for p := range drv.InputDrvs {
		add(p)
	}
	for p := range drv.InputSrcs {
		add(p)
	}
	return refs
}

// -------------------------------------------------------------------
// Fetcher implementation — pure Go download + op 7 (AddToStore)
// -------------------------------------------------------------------

const fetchMaxDecompressedBytes = 1 << 30 // 1 GiB

// FetchTarball implements Fetcher.  It downloads url, unpacks the tarball
// into an in-memory tree, and registers it with the daemon via op 7
// (AddToStore, recursive/NAR hash).  The resulting store path is verified
// against the expected storePath computed from hash.
func (ds *DaemonStore) FetchTarball(storePath, url, name string, hash []byte) error {
	if _, err := os.Stat(storePath); err == nil {
		return nil // fast path: already on disk
	}
	if err := ds.ensureConnected(); err != nil {
		return fmt.Errorf("FetchTarball: %w", err)
	}

	tree, err := downloadAndUnpackTarball(url)
	if err != nil {
		return fmt.Errorf("FetchTarball %s: %w", url, err)
	}

	got, err := ds.daemon.AddToStore(name, "fixed:r:sha256", func(w io.Writer) error {
		return WriteNAR(w, tree)
	})
	if err != nil {
		return fmt.Errorf("FetchTarball %s: AddToStore: %w", url, err)
	}
	if got != storePath {
		return fmt.Errorf("FetchTarball %s: hash mismatch: expected %s, got %s", url, storePath, got)
	}
	return nil
}

// FetchURL implements Fetcher.  It downloads url as a flat file and registers
// it with the daemon via op 7 (AddToStore, flat hash).
func (ds *DaemonStore) FetchURL(storePath, url, name string, hash []byte) error {
	if _, err := os.Stat(storePath); err == nil {
		return nil // fast path: already on disk
	}
	if err := ds.ensureConnected(); err != nil {
		return fmt.Errorf("FetchURL: %w", err)
	}

	content, err := downloadFlat(url)
	if err != nil {
		return fmt.Errorf("FetchURL %s: %w", url, err)
	}

	got, err := ds.daemon.AddToStore(name, "fixed:sha256", func(w io.Writer) error {
		// For a flat file, op 7 expects raw file bytes (not NAR-wrapped).
		_, werr := w.Write(content)
		return werr
	})
	if err != nil {
		return fmt.Errorf("FetchURL %s: AddToStore: %w", url, err)
	}
	if got != storePath {
		return fmt.Errorf("FetchURL %s: hash mismatch: expected %s, got %s", url, storePath, got)
	}
	return nil
}

// downloadAndUnpackTarball fetches url, decompresses (gzip or bzip2 detected
// by Content-Type or URL suffix), and unpacks the tar entries into an
// in-memory MemFSNode tree.  It strips the single mandatory top-level
// directory (matching Lix's fetchTarball semantics).
// Returns an error if the decompressed size exceeds fetchMaxDecompressedBytes.
func downloadAndUnpackTarball(url string) (*MemFSNode, error) {
	resp, err := httpGet(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Wrap body in a decompressor.  Detect from Content-Type then URL suffix.
	cr, err := makeDecompressor(resp, url)
	if err != nil {
		return nil, fmt.Errorf("decompressor: %w", err)
	}
	defer cr.Close()

	// Enforce 1 GiB decompressed limit.
	limited := &limitedReader{r: cr, n: fetchMaxDecompressedBytes}

	tr := tar.NewReader(limited)

	// topDir is the single top-level directory all paths must share.
	topDir := ""
	root := NewMemFSDir()

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("tar: %w", err)
		}

		// Skip PAX global headers and other metadata-only entry types.
		switch hdr.Typeflag {
		case tar.TypeXGlobalHeader, tar.TypeXHeader:
			continue
		}

		// Clean the path and enforce the single-top-level-dir invariant.
		p := path.Clean(hdr.Name)
		if p == "." || p == "" {
			continue
		}
		// Strip leading slash if present.
		if len(p) > 0 && p[0] == '/' {
			p = p[1:]
		}

		// Identify the top-level component.
		slash := strings.IndexByte(p, '/')
		var top string
		if slash < 0 {
			top = p
		} else {
			top = p[:slash]
		}

		if topDir == "" {
			topDir = top
		} else if top != topDir {
			return nil, fmt.Errorf("tarball has multiple top-level entries (%q and %q); expected exactly one", topDir, top)
		}

		// Strip the top-level directory from the path.
		var rel string
		if slash < 0 {
			// This entry IS the top-level directory itself — skip it.
			if hdr.Typeflag == tar.TypeDir {
				continue
			}
			rel = p // top-level file (unusual)
		} else {
			rel = p[slash+1:]
		}
		if rel == "" || rel == "." {
			continue
		}

		// Build the node.
		switch hdr.Typeflag {
		case tar.TypeReg, tar.TypeRegA:
			content, err := io.ReadAll(tr)
			if err != nil {
				return nil, fmt.Errorf("tar read %s: %w", hdr.Name, err)
			}
			node := &MemFSNode{
				Content:      content,
				IsExecutable: hdr.Mode&0111 != 0,
			}
			if err := root.Insert(rel, node); err != nil {
				return nil, fmt.Errorf("tar insert %s: %w", rel, err)
			}

		case tar.TypeSymlink:
			node := &MemFSNode{
				IsSymlink:     true,
				SymlinkTarget: hdr.Linkname,
			}
			if err := root.Insert(rel, node); err != nil {
				return nil, fmt.Errorf("tar insert symlink %s: %w", rel, err)
			}

		case tar.TypeDir:
			// Ensure the directory node exists (Insert creates intermediates,
			// but explicit dir entries may arrive before their children).
			node := NewMemFSDir()
			_ = root.Insert(rel, node) // ignore "already exists" — merge

		case tar.TypeLink: // hard link — copy content from already-inserted node
			// Hard links reference a path already seen; look it up.
			// For simplicity we re-read. In practice hard links are rare.
			// We skip for now (they'll be absent from output), or handle below.
			// TODO: resolve hard links properly.

		default:
			// Ignore device files, fifos, etc.
		}
	}

	if topDir == "" {
		return nil, fmt.Errorf("tarball is empty or contains no files")
	}

	return root, nil
}

// downloadFlat fetches url and returns the raw response body, enforcing the
// 1 GiB limit.
func downloadFlat(url string) ([]byte, error) {
	resp, err := httpGet(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	lr := &limitedReader{r: resp.Body, n: fetchMaxDecompressedBytes}
	data, err := io.ReadAll(lr)
	if err != nil {
		return nil, fmt.Errorf("read response: %w", err)
	}
	return data, nil
}

// httpGet performs a GET request with redirect following.
func httpGet(url string) (*http.Response, error) {
	resp, err := http.Get(url) //nolint:noctx
	if err != nil {
		return nil, fmt.Errorf("http get %s: %w", url, err)
	}
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return nil, fmt.Errorf("http get %s: status %d", url, resp.StatusCode)
	}
	return resp, nil
}

// makeDecompressor wraps resp.Body in the appropriate decompressor based on
// Content-Type header or URL suffix.  Supported: gzip, bzip2.  Falls back to
// raw (no decompression) if unrecognised.
func makeDecompressor(resp *http.Response, rawURL string) (io.ReadCloser, error) {
	ct := resp.Header.Get("Content-Type")
	u := strings.ToLower(rawURL)

	switch {
	case strings.Contains(ct, "gzip") ||
		strings.Contains(ct, "x-gzip") ||
		strings.HasSuffix(u, ".gz") ||
		strings.HasSuffix(u, ".tgz"):
		gz, err := gzip.NewReader(resp.Body)
		if err != nil {
			return nil, fmt.Errorf("gzip: %w", err)
		}
		return gz, nil

	case strings.Contains(ct, "bzip2") ||
		strings.Contains(ct, "x-bzip2") ||
		strings.HasSuffix(u, ".bz2"):
		// compress/bzip2 only has a Reader, no ReadCloser; wrap it.
		return io.NopCloser(bzip2.NewReader(resp.Body)), nil

	default:
		// Unknown — try gzip speculatively (many servers omit Content-Type).
		// If gzip.NewReader fails we fall back to raw.
		if gz, err := gzip.NewReader(resp.Body); err == nil {
			return gz, nil
		}
		return resp.Body, nil
	}
}

// limitedReader wraps an io.Reader and returns an error if more than n bytes
// are read in total.
type limitedReader struct {
	r io.Reader
	n int64
}

func (l *limitedReader) Read(p []byte) (int, error) {
	if l.n <= 0 {
		return 0, fmt.Errorf("fetchTarball: decompressed size exceeds %d byte limit", fetchMaxDecompressedBytes)
	}
	if int64(len(p)) > l.n {
		p = p[:l.n]
	}
	n, err := l.r.Read(p)
	l.n -= int64(n)
	return n, err
}

// ensureConnected lazily opens the daemon connection.
func (ds *DaemonStore) ensureConnected() error {
	if ds.daemon != nil {
		return nil
	}
	conn, err := Connect(ds.socketPath)
	if err != nil {
		return fmt.Errorf("connect to daemon: %w", err)
	}
	ds.daemon = conn
	return nil
}

// Close shuts down the daemon connection if one was established.
func (ds *DaemonStore) Close() error {
	if ds.daemon != nil {
		return ds.daemon.Close()
	}
	return nil
}