main pico / pkg / apps / pgs / uploader_test.go
Eric Bower  ·  2026-08-17
  1package pgs
  2
  3import (
  4	"encoding/pem"
  5	"fmt"
  6	"io"
  7	"log/slog"
  8	"net"
  9	"os"
 10	"os/exec"
 11	"path/filepath"
 12	"strings"
 13	"testing"
 14	"time"
 15
 16	pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
 17	"github.com/picosh/pico/pkg/db"
 18	"github.com/picosh/pico/pkg/pssh"
 19	"github.com/picosh/pico/pkg/shared"
 20	"github.com/picosh/pico/pkg/storage"
 21	"github.com/pkg/sftp"
 22	"github.com/prometheus/client_golang/prometheus"
 23)
 24
 25// TestRsyncDeleteDirectoryWithKeepDir verifies that rsync --delete can
 26// successfully delete a directory that contains . _pico_keep_dir markers.
 27//
 28// Regression test for: "remove /storage/.../project/... directory not empty"
 29// The bug occurred because . _pico_keep_dir files were not cleaned up when
 30// their parent directory was explicitly deleted, causing os.Remove() to fail.
 31func TestRsyncDeleteDirectoryWithKeepDir(t *testing.T) {
 32	opts := &slog.HandlerOptions{
 33		AddSource: true,
 34		Level:     slog.LevelDebug,
 35	}
 36	logger := slog.New(
 37		slog.NewTextHandler(os.Stdout, opts),
 38	)
 39	slog.SetDefault(logger)
 40	dbpool := pgsdb.NewDBMemory(logger)
 41	dbpool.SetupTestData()
 42
 43	// Use filesystem storage so the "directory not empty" bug manifests
 44	tmpDir, err := os.MkdirTemp("", "pgs-test-storage-*")
 45	if err != nil {
 46		t.Fatalf("failed to create temp dir: %v", err)
 47	}
 48	defer func() {
 49		_ = os.RemoveAll(tmpDir)
 50	}()
 51
 52	st, err := storage.NewStorageFS(logger, tmpDir)
 53	if err != nil {
 54		t.Fatalf("failed to create storage: %v", err)
 55	}
 56
 57	pubsub := NewPubsubChan()
 58	defer func() {
 59		_ = pubsub.Close()
 60	}()
 61
 62	_ = os.Setenv("PGS_SSH_PORT", "0")
 63	_ = os.Setenv("PGS_PROM_PORT", "0")
 64
 65	cfg := NewPgsConfig(logger, dbpool, st, pubsub)
 66	done := make(chan error)
 67	readyCh := make(chan *pssh.SSHServer)
 68	prometheus.DefaultRegisterer = prometheus.NewRegistry()
 69
 70	go StartSshServerForTesting(cfg, done, readyCh)
 71
 72	server := <-readyCh
 73	if server == nil {
 74		t.Fatal("failed to create ssh server")
 75	}
 76
 77	var actualAddr string
 78	for i := 0; i < 100; i++ {
 79		server.Mu.Lock()
 80		listener := server.Listener
 81		server.Mu.Unlock()
 82		if listener != nil {
 83			actualAddr = listener.Addr().String()
 84			break
 85		}
 86		time.Sleep(10 * time.Millisecond)
 87	}
 88	if actualAddr == "" {
 89		t.Fatal("server listener not ready")
 90	}
 91
 92	user := GenerateUser()
 93	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
 94		ID:     "test-pubkey-keepdir",
 95		UserID: dbpool.Users[0].ID,
 96		Key:    shared.KeyForKeyText(user.signer.PublicKey()),
 97	})
 98
 99	conn, err := user.NewClientAddr(actualAddr)
100	if err != nil {
101		t.Fatalf("failed to connect: %v", err)
102	}
103	defer func() {
104		_ = conn.Close()
105	}()
106
107	client, err := sftp.NewClient(conn)
108	if err != nil {
109		t.Fatalf("failed to create sftp client: %v", err)
110	}
111	defer func() {
112		_ = client.Close()
113	}()
114
115	// Create temp directory with nested structure
116	name, err := os.MkdirTemp("", "rsync-dir-test-*")
117	if err != nil {
118		t.Fatalf("failed to create temp dir: %v", err)
119	}
120	defer func() {
121		_ = os.RemoveAll(name)
122	}()
123
124	// Create: project/subdir/nested/file.txt
125	nestedDir := filepath.Join(name, "subdir", "nested")
126	err = os.MkdirAll(nestedDir, 0755)
127	if err != nil {
128		t.Fatalf("failed to create nested dir: %v", err)
129	}
130	err = os.WriteFile(filepath.Join(nestedDir, "deep.txt"), []byte("deep content"), 0644)
131	if err != nil {
132		t.Fatalf("failed to write deep.txt: %v", err)
133	}
134
135	// Create: project/subdir/file.txt
136	err = os.WriteFile(filepath.Join(name, "subdir", "file.txt"), []byte("subdir content"), 0644)
137	if err != nil {
138		t.Fatalf("failed to write file.txt: %v", err)
139	}
140
141	// Create: project/index.html (stays after delete)
142	err = os.WriteFile(filepath.Join(name, "index.html"), []byte("index content"), 0644)
143	if err != nil {
144		t.Fatalf("failed to write index.html: %v", err)
145	}
146
147	block := &pem.Block{
148		Type:  "OPENSSH PRIVATE KEY",
149		Bytes: user.privateKey,
150	}
151	keyFile := filepath.Join(name, "id_ed25519")
152	err = os.WriteFile(keyFile, pem.EncodeToMemory(block), 0600)
153	if err != nil {
154		t.Fatalf("failed to write key file: %v", err)
155	}
156
157	_, port, err := net.SplitHostPort(actualAddr)
158	if err != nil {
159		t.Fatalf("failed to parse server address: %v", err)
160	}
161
162	eCmd := fmt.Sprintf(
163		"ssh -p %s -o IdentitiesOnly=yes -i %s -o StrictHostKeyChecking=no",
164		port, keyFile,
165	)
166
167	// Upload files including the subdir/ directory
168	cmd := exec.Command("rsync", "-rv", "-e", eCmd, name+"/", "localhost:/testdir")
169	result, err := cmd.CombinedOutput()
170	if err != nil {
171		t.Fatalf("rsync upload failed: %v\noutput: %s", err, result)
172	}
173
174	// Verify files exist on the server
175	_, err = client.Lstat("/testdir/subdir/file.txt")
176	if err != nil {
177		t.Fatalf("subdir/file.txt not found after upload: %v", err)
178	}
179	_, err = client.Lstat("/testdir/subdir/nested/deep.txt")
180	if err != nil {
181		t.Fatalf("subdir/nested/deep.txt not found after upload: %v", err)
182	}
183
184	// Now remove the entire subdir/ from the local directory
185	err = os.RemoveAll(filepath.Join(name, "subdir"))
186	if err != nil {
187		t.Fatalf("failed to remove local subdir: %v", err)
188	}
189
190	// Run rsync --delete - this should delete the subdir/ directory
191	// WITHOUT failing with "directory not empty"
192	delCmd := exec.Command("rsync", "-rv", "--delete", "-e", eCmd, name+"/", "localhost:/testdir")
193	result, err = delCmd.CombinedOutput()
194	if err != nil {
195		t.Fatalf("rsync --delete failed (this was the bug - 'directory not empty'): %v\noutput: %s", err, result)
196	}
197
198	// Verify the subdir and its contents are gone
199	_, err = client.Lstat("/testdir/subdir")
200	if err == nil {
201		t.Fatal("subdir/ should have been deleted but still exists")
202	}
203	// SFTP can return "no such file" or "file does not exist" depending on version
204	if !strings.Contains(err.Error(), "no such file") && !strings.Contains(err.Error(), "does not exist") {
205		t.Fatalf("expected 'not found' error, got: %v", err)
206	}
207
208	// Verify the . _pico_keep_dir markers are also cleaned up
209	_, err = client.Lstat("/testdir/subdir/._pico_keep_dir")
210	if err == nil {
211		t.Fatal("subdir/._pico_keep_dir should have been deleted but still exists")
212	}
213
214	// Verify index.html still exists (it wasn't deleted)
215	fi, err := client.Lstat("/testdir/index.html")
216	if err != nil {
217		t.Fatalf("index.html should still exist: %v", err)
218	}
219	if fi.Size() != 13 {
220		t.Errorf("index.html has wrong size: got %d, want 13", fi.Size())
221	}
222
223	close(done)
224	time.Sleep(100 * time.Millisecond)
225}
226
227// TestRsyncDeleteNestedEmptyDirectories verifies that rsync --delete handles
228// deeply nested directory structures where intermediate directories become empty
229// and need . _pico_keep_dir markers managed correctly.
230func TestRsyncDeleteNestedEmptyDirectories(t *testing.T) {
231	opts := &slog.HandlerOptions{
232		AddSource: true,
233		Level:     slog.LevelDebug,
234	}
235	logger := slog.New(
236		slog.NewTextHandler(io.Discard, opts),
237	)
238	slog.SetDefault(logger)
239	dbpool := pgsdb.NewDBMemory(logger)
240	dbpool.SetupTestData()
241
242	tmpDir, err := os.MkdirTemp("", "pgs-test-nested-*")
243	if err != nil {
244		t.Fatalf("failed to create temp dir: %v", err)
245	}
246	defer func() {
247		_ = os.RemoveAll(tmpDir)
248	}()
249
250	st, err := storage.NewStorageFS(logger, tmpDir)
251	if err != nil {
252		t.Fatalf("failed to create storage: %v", err)
253	}
254
255	pubsub := NewPubsubChan()
256	defer func() {
257		_ = pubsub.Close()
258	}()
259
260	_ = os.Setenv("PGS_SSH_PORT", "0")
261	_ = os.Setenv("PGS_PROM_PORT", "0")
262
263	cfg := NewPgsConfig(logger, dbpool, st, pubsub)
264	done := make(chan error)
265	readyCh := make(chan *pssh.SSHServer)
266	prometheus.DefaultRegisterer = prometheus.NewRegistry()
267
268	go StartSshServerForTesting(cfg, done, readyCh)
269
270	server := <-readyCh
271	if server == nil {
272		t.Fatal("failed to create ssh server")
273	}
274
275	var actualAddr string
276	for i := 0; i < 100; i++ {
277		server.Mu.Lock()
278		listener := server.Listener
279		server.Mu.Unlock()
280		if listener != nil {
281			actualAddr = listener.Addr().String()
282			break
283		}
284		time.Sleep(10 * time.Millisecond)
285	}
286	if actualAddr == "" {
287		t.Fatal("server listener not ready")
288	}
289
290	user := GenerateUser()
291	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
292		ID:     "test-pubkey-nested",
293		UserID: dbpool.Users[0].ID,
294		Key:    shared.KeyForKeyText(user.signer.PublicKey()),
295	})
296
297	conn, err := user.NewClientAddr(actualAddr)
298	if err != nil {
299		t.Fatalf("failed to connect: %v", err)
300	}
301	defer func() {
302		_ = conn.Close()
303	}()
304
305	client, err := sftp.NewClient(conn)
306	if err != nil {
307		t.Fatalf("failed to create sftp client: %v", err)
308	}
309	defer func() {
310		_ = client.Close()
311	}()
312
313	// Create temp directory with deeply nested structure
314	name, err := os.MkdirTemp("", "rsync-nested-test-*")
315	if err != nil {
316		t.Fatalf("failed to create temp dir: %v", err)
317	}
318	defer func() {
319		_ = os.RemoveAll(name)
320	}()
321
322	// Create: a/b/c/d/file.txt
323	deepDir := filepath.Join(name, "a", "b", "c", "d")
324	err = os.MkdirAll(deepDir, 0755)
325	if err != nil {
326		t.Fatalf("failed to create deep dir: %v", err)
327	}
328	err = os.WriteFile(filepath.Join(deepDir, "file.txt"), []byte("deep"), 0644)
329	if err != nil {
330		t.Fatalf("failed to write file: %v", err)
331	}
332
333	block := &pem.Block{
334		Type:  "OPENSSH PRIVATE KEY",
335		Bytes: user.privateKey,
336	}
337	keyFile := filepath.Join(name, "id_ed25519")
338	err = os.WriteFile(keyFile, pem.EncodeToMemory(block), 0600)
339	if err != nil {
340		t.Fatalf("failed to write key file: %v", err)
341	}
342
343	_, port, err := net.SplitHostPort(actualAddr)
344	if err != nil {
345		t.Fatalf("failed to parse server address: %v", err)
346	}
347
348	eCmd := fmt.Sprintf(
349		"ssh -p %s -o IdentitiesOnly=yes -i %s -o StrictHostKeyChecking=no",
350		port, keyFile,
351	)
352
353	// Upload
354	cmd := exec.Command("rsync", "-rv", "-e", eCmd, name+"/", "localhost:/deep")
355	result, err := cmd.CombinedOutput()
356	if err != nil {
357		t.Fatalf("rsync upload failed: %v\noutput: %s", err, result)
358	}
359
360	// Remove the entire a/ directory locally
361	err = os.RemoveAll(filepath.Join(name, "a"))
362	if err != nil {
363		t.Fatalf("failed to remove local a/: %v", err)
364	}
365
366	// rsync --delete should handle the deeply nested structure
367	delCmd := exec.Command("rsync", "-rv", "--delete", "-e", eCmd, name+"/", "localhost:/deep")
368	result, err = delCmd.CombinedOutput()
369	if err != nil {
370		t.Fatalf("rsync --delete failed on nested dirs: %v\noutput: %s", err, result)
371	}
372
373	// Verify the entire tree is gone
374	_, err = client.Lstat("/deep/a")
375	if err == nil {
376		t.Fatal("/deep/a should have been deleted but still exists")
377	}
378
379	close(done)
380	time.Sleep(100 * time.Millisecond)
381}
382
383// setupPgsTestServer spins up an in-memory PGS SSH server for testing and
384// returns the server address and a teardown func.
385func setupPgsTestServer(t *testing.T, dbpool *pgsdb.MemoryDB) (string, func()) {
386	t.Helper()
387
388	opts := &slog.HandlerOptions{Level: slog.LevelDebug}
389	logger := slog.New(slog.NewTextHandler(io.Discard, opts))
390	slog.SetDefault(logger)
391
392	st, err := storage.NewStorageMemory(map[string]map[string]string{})
393	if err != nil {
394		t.Fatalf("failed to create storage: %v", err)
395	}
396
397	pubsub := NewPubsubChan()
398
399	_ = os.Setenv("PGS_SSH_PORT", "0")
400	_ = os.Setenv("PGS_PROM_PORT", "0")
401
402	cfg := NewPgsConfig(logger, dbpool, st, pubsub)
403	done := make(chan error)
404	readyCh := make(chan *pssh.SSHServer)
405	prometheus.DefaultRegisterer = prometheus.NewRegistry()
406
407	go StartSshServerForTesting(cfg, done, readyCh)
408
409	server := <-readyCh
410	if server == nil {
411		t.Fatal("failed to create ssh server")
412	}
413
414	var actualAddr string
415	for i := 0; i < 100; i++ {
416		server.Mu.Lock()
417		listener := server.Listener
418		server.Mu.Unlock()
419		if listener != nil {
420			actualAddr = listener.Addr().String()
421			break
422		}
423		time.Sleep(10 * time.Millisecond)
424	}
425	if actualAddr == "" {
426		t.Fatal("server listener not ready")
427	}
428
429	teardown := func() {
430		_ = pubsub.Close()
431		close(done)
432		time.Sleep(100 * time.Millisecond)
433	}
434
435	return actualAddr, teardown
436}
437
438// TestPlusFlagAllowsUpload verifies that a user with a valid pico+ feature flag
439// can successfully upload files via SFTP.
440func TestPlusFlagAllowsUpload(t *testing.T) {
441	dbpool := pgsdb.NewDBMemory(slog.Default())
442	// SetupTestData creates a user with a valid (non-expired) plus feature flag.
443	dbpool.SetupTestData()
444
445	addr, teardown := setupPgsTestServer(t, dbpool)
446	defer teardown()
447
448	user := GenerateUser()
449	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
450		ID:     "plus-user-pubkey",
451		UserID: dbpool.Users[0].ID,
452		Key:    shared.KeyForKeyText(user.signer.PublicKey()),
453	})
454
455	conn, err := user.NewClientAddr(addr)
456	if err != nil {
457		t.Fatalf("ssh dial failed: %v", err)
458	}
459	defer func() { _ = conn.Close() }()
460
461	client, err := sftp.NewClient(conn)
462	if err != nil {
463		t.Fatalf("sftp client failed: %v", err)
464	}
465	defer func() { _ = client.Close() }()
466
467	// Upload should succeed for a plus user.
468	f, err := client.Create("myproject/index.html")
469	if err != nil {
470		t.Fatalf("plus user upload was blocked unexpectedly: %v", err)
471	}
472	if _, err := f.Write([]byte("<html>hello</html>")); err != nil {
473		t.Fatalf("write failed: %v", err)
474	}
475	_ = f.Close()
476
477	// Confirm the file is accessible via SFTP stat.
478	fi, err := client.Lstat("myproject/index.html")
479	if err != nil {
480		t.Fatalf("file not found after upload: %v", err)
481	}
482	if fi.Size() == 0 {
483		t.Errorf("uploaded file has zero size")
484	}
485
486	// Confirm the project was created in the DB (no empty-project side-effect).
487	projects, err := dbpool.FindProjectsByUser(dbpool.Users[0].ID)
488	if err != nil {
489		t.Fatalf("FindProjectsByUser: %v", err)
490	}
491	found := false
492	for _, p := range projects {
493		if p.Name == "myproject" {
494			found = true
495			break
496		}
497	}
498	if !found {
499		t.Errorf("project 'myproject' was not created in the DB after upload")
500	}
501}
502
503// TestNoPlusFlagBlocksUpload verifies that a user whose pico+ subscription has
504// expired cannot upload files. The Validate() step (introduced in the last
505// commit) must reject the connection before any project is created, so we also
506// confirm that no empty project leaks into the DB.
507func TestNoPlusFlagBlocksUpload(t *testing.T) {
508	dbpool := pgsdb.NewDBMemory(slog.Default())
509	// Create a user but give them an *expired* plus feature flag.
510	dbpool.SetupTestData()
511	expired := time.Now().Add(-time.Hour)
512	dbpool.Feature.ExpiresAt = &expired // mark as expired
513
514	addr, teardown := setupPgsTestServer(t, dbpool)
515	defer teardown()
516
517	user := GenerateUser()
518	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
519		ID:     "no-plus-user-pubkey",
520		UserID: dbpool.Users[0].ID,
521		Key:    shared.KeyForKeyText(user.signer.PublicKey()),
522	})
523
524	conn, err := user.NewClientAddr(addr)
525	if err != nil {
526		t.Fatalf("ssh dial failed: %v", err)
527	}
528	defer func() { _ = conn.Close() }()
529
530	// Validate() fires at the SFTP-subsystem session level, so the server may
531	// close the connection before the SFTP version handshake completes.
532	// Accept rejection at either sftp.NewClient or client.Create as valid.
533	client, sftpErr := sftp.NewClient(conn)
534	if sftpErr != nil {
535		// Server rejected the session during Validate() — this is the expected
536		// behavior for a non-plus user. Verify no project was leaked.
537		t.Logf("upload correctly blocked at SFTP handshake: %v", sftpErr)
538	} else {
539		defer func() { _ = client.Close() }()
540
541		// Upload must be rejected because the user's pico+ has expired.
542		_, err = client.Create("someproject/index.html")
543		if err == nil {
544			t.Fatal("expected upload to be blocked for non-plus user, but it succeeded")
545		}
546		t.Logf("upload correctly blocked at Create: %v", err)
547	}
548
549	// Most importantly: no empty project should have been created in the DB,
550	// because the refactored Validate() checks the feature flag *before*
551	// upserting the project bucket/quota.
552	projects, err := dbpool.FindProjectsByUser(dbpool.Users[0].ID)
553	if err != nil {
554		t.Fatalf("FindProjectsByUser: %v", err)
555	}
556	if len(projects) != 0 {
557		t.Errorf("expected 0 projects for non-plus user, got %d (empty project leak)", len(projects))
558	}
559}
560
561func TestFindFeatureFlag(t *testing.T) {
562	logger := slog.Default()
563	cfg := &PgsConfig{
564		MaxSize:            uint64(100 * shared.MB),
565		MaxAssetSize:       int64(10 * shared.MB),
566		MaxSpecialFileSize: int64(5 * shared.KB),
567	}
568
569	validExpires := time.Now().Add(24 * time.Hour)
570	expiredExpires := time.Now().Add(-24 * time.Hour)
571	userID := "user-1"
572
573	t.Run("plus valid returns plus", func(t *testing.T) {
574		dbpool := pgsdb.NewDBMemory(logger)
575		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
576		plusFF.ExpiresAt = &validExpires
577		dbpool.Features = []*db.FeatureFlag{plusFF}
578
579		ff, err := findFeatureFlag(dbpool, cfg, userID)
580		if err != nil {
581			t.Fatalf("unexpected error: %v", err)
582		}
583		if ff.Name != "plus" {
584			t.Errorf("expected plus, got %s", ff.Name)
585		}
586	})
587
588	t.Run("pgs valid returns pgs when no plus", func(t *testing.T) {
589		dbpool := pgsdb.NewDBMemory(logger)
590		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
591		pgsFF.ExpiresAt = &validExpires
592		dbpool.Features = []*db.FeatureFlag{pgsFF}
593
594		ff, err := findFeatureFlag(dbpool, cfg, userID)
595		if err != nil {
596			t.Fatalf("unexpected error: %v", err)
597		}
598		if ff.Name != "pgs" {
599			t.Errorf("expected pgs, got %s", ff.Name)
600		}
601	})
602
603	t.Run("plus preferred over pgs when both valid", func(t *testing.T) {
604		dbpool := pgsdb.NewDBMemory(logger)
605		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
606		plusFF.ExpiresAt = &validExpires
607		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30*shared.MB), int64(3*shared.MB), int64(1*shared.KB))
608		pgsFF.ExpiresAt = &validExpires
609		dbpool.Features = []*db.FeatureFlag{plusFF, pgsFF}
610
611		ff, err := findFeatureFlag(dbpool, cfg, userID)
612		if err != nil {
613			t.Fatalf("unexpected error: %v", err)
614		}
615		if ff.Name != "plus" {
616			t.Errorf("expected plus to be picked first, got %s", ff.Name)
617		}
618	})
619
620	t.Run("pgs picked when plus is expired", func(t *testing.T) {
621		dbpool := pgsdb.NewDBMemory(logger)
622		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
623		plusFF.ExpiresAt = &expiredExpires
624		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30*shared.MB), int64(3*shared.MB), int64(1*shared.KB))
625		pgsFF.ExpiresAt = &validExpires
626		dbpool.Features = []*db.FeatureFlag{plusFF, pgsFF}
627
628		ff, err := findFeatureFlag(dbpool, cfg, userID)
629		if err != nil {
630			t.Fatalf("unexpected error: %v", err)
631		}
632		if ff.Name != "pgs" {
633			t.Errorf("expected pgs when plus is expired, got %s", ff.Name)
634		}
635	})
636
637	t.Run("error when both plus and pgs expired", func(t *testing.T) {
638		dbpool := pgsdb.NewDBMemory(logger)
639		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
640		plusFF.ExpiresAt = &expiredExpires
641		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30*shared.MB), int64(3*shared.MB), int64(1*shared.KB))
642		pgsFF.ExpiresAt = &expiredExpires
643		dbpool.Features = []*db.FeatureFlag{plusFF, pgsFF}
644
645		_, err := findFeatureFlag(dbpool, cfg, userID)
646		if err == nil {
647			t.Fatal("expected error, got nil")
648		}
649	})
650
651	t.Run("error when neither plus nor pgs exists", func(t *testing.T) {
652		dbpool := pgsdb.NewDBMemory(logger)
653		_, err := findFeatureFlag(dbpool, cfg, userID)
654		if err == nil {
655			t.Fatal("expected error, got nil")
656		}
657	})
658}
659
660func TestPgsFlagAllowsUpload(t *testing.T) {
661	dbpool := pgsdb.NewDBMemory(slog.Default())
662	dbpool.SetupTestData()
663	// Replace default plus flag with a valid pgs flag
664	dbpool.Feature = nil
665	valid := time.Now().Add(24 * time.Hour)
666	pgsFF := db.NewFeatureFlag(
667		dbpool.Users[0].ID,
668		"pgs",
669		uint64(25*shared.MB),
670		int64(10*shared.MB),
671		int64(5*shared.KB),
672	)
673	pgsFF.ExpiresAt = &valid
674	dbpool.Features = []*db.FeatureFlag{pgsFF}
675
676	addr, teardown := setupPgsTestServer(t, dbpool)
677	defer teardown()
678
679	user := GenerateUser()
680	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
681		ID:     "pgs-user-pubkey",
682		UserID: dbpool.Users[0].ID,
683		Key:    shared.KeyForKeyText(user.signer.PublicKey()),
684	})
685
686	conn, err := user.NewClientAddr(addr)
687	if err != nil {
688		t.Fatalf("ssh dial failed: %v", err)
689	}
690	defer func() { _ = conn.Close() }()
691
692	client, err := sftp.NewClient(conn)
693	if err != nil {
694		t.Fatalf("sftp client failed: %v", err)
695	}
696	defer func() { _ = client.Close() }()
697
698	// Upload should succeed for a pgs user.
699	f, err := client.Create("myproject/index.html")
700	if err != nil {
701		t.Fatalf("pgs user upload was blocked unexpectedly: %v", err)
702	}
703	if _, err := f.Write([]byte("<html>hello from pgs</html>")); err != nil {
704		t.Fatalf("write failed: %v", err)
705	}
706	_ = f.Close()
707
708	// Confirm the file is accessible via SFTP stat.
709	fi, err := client.Lstat("myproject/index.html")
710	if err != nil {
711		t.Fatalf("file not found after upload: %v", err)
712	}
713	if fi.Size() == 0 {
714		t.Errorf("uploaded file has zero size")
715	}
716}