chore(pgs): uploader tests confirming pico+ only

Commit `538c3e9`

Eric Bower · 2026-08-13 10:05:13 -0400 EDT

parent e5d5430

chore(pgs): uploader tests confirming pico+ only

1 files changed, +178, -0

M pkg/apps/pgs/uploader_test.go

+178, -0

......@@ -379,3 +379,181 @@ func TestRsyncDeleteNestedEmptyDirectories(t *testing.T) {
	close(done)
	time.Sleep(100 * time.Millisecond)
}

// setupPgsTestServer spins up an in-memory PGS SSH server for testing and
// returns the server address and a teardown func.
func setupPgsTestServer(t *testing.T, dbpool *pgsdb.MemoryDB) (string, func()) {
	t.Helper()

	opts := &slog.HandlerOptions{Level: slog.LevelDebug}
	logger := slog.New(slog.NewTextHandler(io.Discard, opts))
	slog.SetDefault(logger)

	st, err := storage.NewStorageMemory(map[string]map[string]string{})
	if err != nil {
		t.Fatalf("failed to create storage: %v", err)
	}

	pubsub := NewPubsubChan()

	_ = os.Setenv("PGS_SSH_PORT", "0")
	_ = os.Setenv("PGS_PROM_PORT", "0")

	cfg := NewPgsConfig(logger, dbpool, st, pubsub)
	done := make(chan error)
	readyCh := make(chan *pssh.SSHServer)
	prometheus.DefaultRegisterer = prometheus.NewRegistry()

	go StartSshServerForTesting(cfg, done, readyCh)

	server := <-readyCh
	if server == nil {
		t.Fatal("failed to create ssh server")
	}

	var actualAddr string
	for i := 0; i < 100; i++ {
		server.Mu.Lock()
		listener := server.Listener
		server.Mu.Unlock()
		if listener != nil {
			actualAddr = listener.Addr().String()
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if actualAddr == "" {
		t.Fatal("server listener not ready")
	}

	teardown := func() {
		_ = pubsub.Close()
		close(done)
		time.Sleep(100 * time.Millisecond)
	}

	return actualAddr, teardown
}

// TestPlusFlagAllowsUpload verifies that a user with a valid pico+ feature flag
// can successfully upload files via SFTP.
func TestPlusFlagAllowsUpload(t *testing.T) {
	dbpool := pgsdb.NewDBMemory(slog.Default())
	// SetupTestData creates a user with a valid (non-expired) plus feature flag.
	dbpool.SetupTestData()

	addr, teardown := setupPgsTestServer(t, dbpool)
	defer teardown()

	user := GenerateUser()
	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
		ID:     "plus-user-pubkey",
		UserID: dbpool.Users[0].ID,
		Key:    shared.KeyForKeyText(user.signer.PublicKey()),
	})

	conn, err := user.NewClientAddr(addr)
	if err != nil {
		t.Fatalf("ssh dial failed: %v", err)
	}
	defer func() { _ = conn.Close() }()

	client, err := sftp.NewClient(conn)
	if err != nil {
		t.Fatalf("sftp client failed: %v", err)
	}
	defer func() { _ = client.Close() }()

	// Upload should succeed for a plus user.
	f, err := client.Create("myproject/index.html")
	if err != nil {
		t.Fatalf("plus user upload was blocked unexpectedly: %v", err)
	}
	if _, err := f.Write([]byte("<html>hello</html>")); err != nil {
		t.Fatalf("write failed: %v", err)
	}
	_ = f.Close()

	// Confirm the file is accessible via SFTP stat.
	fi, err := client.Lstat("myproject/index.html")
	if err != nil {
		t.Fatalf("file not found after upload: %v", err)
	}
	if fi.Size() == 0 {
		t.Errorf("uploaded file has zero size")
	}

	// Confirm the project was created in the DB (no empty-project side-effect).
	projects, err := dbpool.FindProjectsByUser(dbpool.Users[0].ID)
	if err != nil {
		t.Fatalf("FindProjectsByUser: %v", err)
	}
	found := false
	for _, p := range projects {
		if p.Name == "myproject" {
			found = true
			break
		}
	}
	if !found {
		t.Errorf("project 'myproject' was not created in the DB after upload")
	}
}

// TestNoPlusFlagBlocksUpload verifies that a user whose pico+ subscription has
// expired cannot upload files. The Validate() step (introduced in the last
// commit) must reject the connection before any project is created, so we also
// confirm that no empty project leaks into the DB.
func TestNoPlusFlagBlocksUpload(t *testing.T) {
	dbpool := pgsdb.NewDBMemory(slog.Default())
	// Create a user but give them an *expired* plus feature flag.
	dbpool.SetupTestData()
	expired := time.Now().Add(-time.Hour)
	dbpool.Feature.ExpiresAt = &expired // mark as expired

	addr, teardown := setupPgsTestServer(t, dbpool)
	defer teardown()

	user := GenerateUser()
	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
		ID:     "no-plus-user-pubkey",
		UserID: dbpool.Users[0].ID,
		Key:    shared.KeyForKeyText(user.signer.PublicKey()),
	})

	conn, err := user.NewClientAddr(addr)
	if err != nil {
		t.Fatalf("ssh dial failed: %v", err)
	}
	defer func() { _ = conn.Close() }()

	// Validate() fires at the SFTP-subsystem session level, so the server may
	// close the connection before the SFTP version handshake completes.
	// Accept rejection at either sftp.NewClient or client.Create as valid.
	client, sftpErr := sftp.NewClient(conn)
	if sftpErr != nil {
		// Server rejected the session during Validate() — this is the expected
		// behavior for a non-plus user. Verify no project was leaked.
		t.Logf("upload correctly blocked at SFTP handshake: %v", sftpErr)
	} else {
		defer func() { _ = client.Close() }()

		// Upload must be rejected because the user's pico+ has expired.
		_, err = client.Create("someproject/index.html")
		if err == nil {
			t.Fatal("expected upload to be blocked for non-plus user, but it succeeded")
		}
		t.Logf("upload correctly blocked at Create: %v", err)
	}

	// Most importantly: no empty project should have been created in the DB,
	// because the refactored Validate() checks the feature flag *before*
	// upserting the project bucket/quota.
	projects, err := dbpool.FindProjectsByUser(dbpool.Users[0].ID)
	if err != nil {
		t.Fatalf("FindProjectsByUser: %v", err)
	}
	if len(projects) != 0 {
		t.Errorf("expected 0 projects for non-plus user, got %d (empty project leak)", len(projects))
	}
}

...	...		@@ -379,3 +379,181 @@ func TestRsyncDeleteNestedEmptyDirectories(t *testing.T) {
379	379		close(done)
380	380		time.Sleep(100 * time.Millisecond)
381	381		}
	382	+
	383	+	// setupPgsTestServer spins up an in-memory PGS SSH server for testing and
	384	+	// returns the server address and a teardown func.
	385	+	func 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.
	440	+	func 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.
	507	+	func 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	+	}

Commit 538c3e9

Commit `538c3e9`