chore(pgs): check for pgs feature flag

Commit `7b5a5aa`

Eric Bower · 2026-08-17 09:35:36 -0400 EDT

parent c58eb5b

chore(pgs): check for pgs feature flag

11 files changed, +290, -23

A cmd/scripts/feature-flag/main.go

M pkg/apps/pgs/cli.go

M pkg/apps/pgs/config.go

M pkg/apps/pgs/db/memory.go

A pkg/apps/pgs/feature_flag.go

M pkg/apps/pgs/uploader.go

M pkg/apps/pgs/uploader_test.go

M pkg/db/db.go

M pkg/db/postgres/storage.go

M pkg/db/postgres/storage_test.go

M pkg/db/stub/stub.go

A cmd/scripts/feature-flag/main.go

+37, -0

......@@ -0,0 +1,37 @@
package main

import (
	"log/slog"
	"os"

	"github.com/picosh/pico/pkg/db/postgres"
)

func main() {
	logger := slog.Default()
	DbURL := os.Getenv("DATABASE_URL")
	dbpool := postgres.NewDB(DbURL, logger)

	args := os.Args
	if len(args) < 3 {
		logger.Error("usage: go run ./cmd/scripts/feature-flag <username> <feature>")
		os.Exit(1)
	}

	username := args[1]
	feature := args[2]

	logger.Info(
		"Adding feature flag to user",
		"username", username,
		"feature", feature,
	)

	err := dbpool.AddFeatureUser(username, feature)
	if err != nil {
		logger.Error("Failed to add feature flag to user", "err", err)
		os.Exit(1)
	} else {
		logger.Info("Successfully added feature flag to user")
	}
}

M pkg/apps/pgs/cli.go

+3, -5

......@@ -230,13 +230,11 @@ This means only you can access the site through a web tunnel or by downloading t
}

func (c *Cmd) stats(cfgMaxSize uint64) error {
	ff, err := c.Dbpool.FindFeature(c.User.ID, "plus")
	ff, err := findFeatureFlag(c.Dbpool, c.Cfg, c.User.ID)
	if err != nil {
		ff = db.NewFeatureFlag(c.User.ID, "plus", cfgMaxSize, 0, 0)
		ff = db.NewFeatureFlag(c.User.ID, "pgs", cfgMaxSize, 0, 0)
	}
	// this is jank
	ff.Data.StorageMax = ff.FindStorageMax(cfgMaxSize)
	storageMax := ff.Data.StorageMax
	storageMax := ff.FindStorageMax(cfgMaxSize)

	bucketName := shared.GetAssetBucketName(c.User.ID)
	bucket, err := c.Store.GetBucket(bucketName)

M pkg/apps/pgs/config.go

+1, -1

......@@ -62,7 +62,7 @@ func (c *PgsConfig) StaticPath(fname string) string {
	return filepath.Join("pkg", "apps", "pgs", fname)
}

var maxSize = uint64(25 * shared.MB)
var maxSize = uint64(50 * shared.MB)
var maxAssetSize = int64(10 * shared.MB)

// Needs to be small for caching files like _headers and _redirects.

M pkg/apps/pgs/db/memory.go

+10, -1

......@@ -16,6 +16,7 @@ type MemoryDB struct {
	Projects    []*db.Project
	Pubkeys     []*db.PublicKey
	Feature     *db.FeatureFlag
	Features    []*db.FeatureFlag
	FormEntries []*db.FormEntry
}

......@@ -82,7 +83,15 @@ func (me *MemoryDB) FindUserByName(name string) (*db.User, error) {
}

func (me *MemoryDB) FindFeature(userID, name string) (*db.FeatureFlag, error) {
	return me.Feature, nil
	for _, ff := range me.Features {
		if ff.UserID == userID && ff.Name == name {
			return ff, nil
		}
	}
	if me.Feature != nil && (me.Feature.Name == name || me.Feature.Name == "") {
		return me.Feature, nil
	}
	return nil, fmt.Errorf("feature flag not found")
}

func (me *MemoryDB) Close() error {

A pkg/apps/pgs/feature_flag.go

+43, -0

......@@ -0,0 +1,43 @@
package pgs

import (
	"fmt"
	"strings"

	pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
	"github.com/picosh/pico/pkg/db"
)

func setFeatureLimits(ff *db.FeatureFlag, cfg *PgsConfig) {
	ff.Data.StorageMax = ff.FindStorageMax(cfg.MaxSize)
	ff.Data.FileMax = ff.FindFileMax(cfg.MaxAssetSize)
	ff.Data.SpecialFileMax = ff.FindSpecialFileMax(cfg.MaxSpecialFileSize)
}

func findFeatureFlag(dbpool pgsdb.PgsDB, cfg *PgsConfig, userID string) (*db.FeatureFlag, error) {
	ff, err := dbpool.FindFeature(userID, "plus")
	if err == nil {
		if ff.IsValid() {
			setFeatureLimits(ff, cfg)
			return ff, nil
		}
		err = fmt.Errorf("your pico+ has expired")
	}

	ffPgs, pgsErr := dbpool.FindFeature(userID, "pgs")
	if pgsErr == nil {
		if ffPgs.IsValid() {
			setFeatureLimits(ffPgs, cfg)
			return ffPgs, nil
		}
		pgsErr = fmt.Errorf("your pgs access has expired")
	}

	if err != nil && strings.Contains(err.Error(), "expired") {
		return nil, err
	}
	if pgsErr != nil {
		return nil, pgsErr
	}
	return nil, err
}

M pkg/apps/pgs/uploader.go

+1, -16

......@@ -15,7 +15,6 @@ import (
	"sync"
	"time"

	pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
	"github.com/picosh/pico/pkg/db"
	"github.com/picosh/pico/pkg/pssh"
	sendutils "github.com/picosh/pico/pkg/send/utils"

......@@ -219,7 +218,7 @@ func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error {
		return err
	}

	ff, err := findPlusFF(h.Cfg.DB, h.Cfg, user.ID)
	ff, err := findFeatureFlag(h.Cfg.DB, h.Cfg, user.ID)
	if err != nil {
		return err
	}

......@@ -273,20 +272,6 @@ func (h *UploadAssetHandler) findDenylist(bucket storage.Bucket, project *db.Pro
	return str, nil
}

func findPlusFF(dbpool pgsdb.PgsDB, cfg *PgsConfig, userID string) (*db.FeatureFlag, error) {
	ff, err := dbpool.FindFeature(userID, "plus")
	if err != nil {
		return nil, err
	}
	if !ff.IsValid() {
		return nil, fmt.Errorf("your pico+ has expired")
	}
	ff.Data.StorageMax = ff.FindStorageMax(cfg.MaxSize)
	ff.Data.FileMax = ff.FindFileMax(cfg.MaxAssetSize)
	ff.Data.SpecialFileMax = ff.FindSpecialFileMax(cfg.MaxSpecialFileSize)
	return ff, nil
}

func mtimeToTime(entry *sendutils.FileEntry) time.Time {
	var mtime time.Time
	if entry.Mtime > 0 {

M pkg/apps/pgs/uploader_test.go

+157, -0

......@@ -557,3 +557,160 @@ func TestNoPlusFlagBlocksUpload(t *testing.T) {
		t.Errorf("expected 0 projects for non-plus user, got %d (empty project leak)", len(projects))
	}
}

func TestFindFeatureFlag(t *testing.T) {
	logger := slog.Default()
	cfg := &PgsConfig{
		MaxSize:            uint64(100 * shared.MB),
		MaxAssetSize:       int64(10 * shared.MB),
		MaxSpecialFileSize: int64(5 * shared.KB),
	}

	validExpires := time.Now().Add(24 * time.Hour)
	expiredExpires := time.Now().Add(-24 * time.Hour)
	userID := "user-1"

	t.Run("plus valid returns plus", func(t *testing.T) {
		dbpool := pgsdb.NewDBMemory(logger)
		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
		plusFF.ExpiresAt = &validExpires
		dbpool.Features = []*db.FeatureFlag{plusFF}

		ff, err := findFeatureFlag(dbpool, cfg, userID)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if ff.Name != "plus" {
			t.Errorf("expected plus, got %s", ff.Name)
		}
	})

	t.Run("pgs valid returns pgs when no plus", func(t *testing.T) {
		dbpool := pgsdb.NewDBMemory(logger)
		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
		pgsFF.ExpiresAt = &validExpires
		dbpool.Features = []*db.FeatureFlag{pgsFF}

		ff, err := findFeatureFlag(dbpool, cfg, userID)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if ff.Name != "pgs" {
			t.Errorf("expected pgs, got %s", ff.Name)
		}
	})

	t.Run("plus preferred over pgs when both valid", func(t *testing.T) {
		dbpool := pgsdb.NewDBMemory(logger)
		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
		plusFF.ExpiresAt = &validExpires
		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30*shared.MB), int64(3*shared.MB), int64(1*shared.KB))
		pgsFF.ExpiresAt = &validExpires
		dbpool.Features = []*db.FeatureFlag{plusFF, pgsFF}

		ff, err := findFeatureFlag(dbpool, cfg, userID)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if ff.Name != "plus" {
			t.Errorf("expected plus to be picked first, got %s", ff.Name)
		}
	})

	t.Run("pgs picked when plus is expired", func(t *testing.T) {
		dbpool := pgsdb.NewDBMemory(logger)
		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
		plusFF.ExpiresAt = &expiredExpires
		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30*shared.MB), int64(3*shared.MB), int64(1*shared.KB))
		pgsFF.ExpiresAt = &validExpires
		dbpool.Features = []*db.FeatureFlag{plusFF, pgsFF}

		ff, err := findFeatureFlag(dbpool, cfg, userID)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if ff.Name != "pgs" {
			t.Errorf("expected pgs when plus is expired, got %s", ff.Name)
		}
	})

	t.Run("error when both plus and pgs expired", func(t *testing.T) {
		dbpool := pgsdb.NewDBMemory(logger)
		plusFF := db.NewFeatureFlag(userID, "plus", uint64(50*shared.MB), int64(5*shared.MB), int64(2*shared.KB))
		plusFF.ExpiresAt = &expiredExpires
		pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30*shared.MB), int64(3*shared.MB), int64(1*shared.KB))
		pgsFF.ExpiresAt = &expiredExpires
		dbpool.Features = []*db.FeatureFlag{plusFF, pgsFF}

		_, err := findFeatureFlag(dbpool, cfg, userID)
		if err == nil {
			t.Fatal("expected error, got nil")
		}
	})

	t.Run("error when neither plus nor pgs exists", func(t *testing.T) {
		dbpool := pgsdb.NewDBMemory(logger)
		_, err := findFeatureFlag(dbpool, cfg, userID)
		if err == nil {
			t.Fatal("expected error, got nil")
		}
	})
}

func TestPgsFlagAllowsUpload(t *testing.T) {
	dbpool := pgsdb.NewDBMemory(slog.Default())
	dbpool.SetupTestData()
	// Replace default plus flag with a valid pgs flag
	dbpool.Feature = nil
	valid := time.Now().Add(24 * time.Hour)
	pgsFF := db.NewFeatureFlag(
		dbpool.Users[0].ID,
		"pgs",
		uint64(25*shared.MB),
		int64(10*shared.MB),
		int64(5*shared.KB),
	)
	pgsFF.ExpiresAt = &valid
	dbpool.Features = []*db.FeatureFlag{pgsFF}

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

	user := GenerateUser()
	dbpool.Pubkeys = append(dbpool.Pubkeys, &db.PublicKey{
		ID:     "pgs-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 pgs user.
	f, err := client.Create("myproject/index.html")
	if err != nil {
		t.Fatalf("pgs user upload was blocked unexpectedly: %v", err)
	}
	if _, err := f.Write([]byte("<html>hello from pgs</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")
	}
}

M pkg/db/db.go

+1, -0

......@@ -609,6 +609,7 @@ type DB interface {
	VisitUrlNotFound(opts *SummaryOpts) ([]*VisitUrl, error)

	AddPicoPlusUser(username, email, paymentType, txId string) error
	AddFeatureUser(username, name string) error
	FindFeature(userID string, feature string) (*FeatureFlag, error)
	FindFeaturesByUser(userID string) ([]*FeatureFlag, error)
	HasFeatureByUser(userID string, feature string) bool

M pkg/db/postgres/storage.go

+11, -0

......@@ -1902,6 +1902,17 @@ func (me *PsqlDB) AddPicoPlusUser(username, email, paymentType, txId string) err
	return tx.Commit()
}

func (me *PsqlDB) AddFeatureUser(username, name string) error {
	user, err := me.FindUserByName(username)
	if err != nil {
		return err
	}

	expiresAt := me.createFeatureExpiresAt(user.ID, name)
	_, err = me.InsertFeature(user.ID, name, expiresAt)
	return err
}

func (me *PsqlDB) UpsertProject(userID, projectName, projectDir string) (*db.Project, error) {
	project, err := me.FindProjectByName(userID, projectName)
	if err == nil {

M pkg/db/postgres/storage_test.go

+22, -0

......@@ -1081,6 +1081,28 @@ func TestAddPicoPlusUser(t *testing.T) {
	}
}

func TestAddFeatureUser(t *testing.T) {
	cleanupTestData(t)

	user, _ := testDB.RegisterUser("featureuserowner", "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAI featureuserowner", "comment", "")

	err := testDB.AddFeatureUser("featureuserowner", "pgs")
	if err != nil {
		t.Fatalf("AddFeatureUser failed: %v", err)
	}

	ff, err := testDB.FindFeature(user.ID, "pgs")
	if err != nil {
		t.Fatalf("FindFeature failed: %v", err)
	}
	if ff.Name != "pgs" {
		t.Errorf("expected feature name pgs, got %s", ff.Name)
	}
	if !ff.IsValid() {
		t.Errorf("expected feature to be valid")
	}
}

// ============ Feed Items Tests ============

func TestInsertFeedItems(t *testing.T) {

M pkg/db/stub/stub.go

+4, -0

......@@ -204,6 +204,10 @@ func (me *StubDB) AddPicoPlusUser(username, email, paymentType, txId string) err
	return errNotImpl
}

func (me *StubDB) AddFeatureUser(username, name string) error {
	return errNotImpl
}

func (me *StubDB) FindUserStats(userID string) (*db.UserStats, error) {
	return nil, errNotImpl
}

...	...		@@ -230,13 +230,11 @@ This means only you can access the site through a web tunnel or by downloading t
230	230		}
231	231
232	232		func (c *Cmd) stats(cfgMaxSize uint64) error {
233		-	ff, err := c.Dbpool.FindFeature(c.User.ID, "plus")
	233	+	ff, err := findFeatureFlag(c.Dbpool, c.Cfg, c.User.ID)
234	234		if err != nil {
235		-	ff = db.NewFeatureFlag(c.User.ID, "plus", cfgMaxSize, 0, 0)
	235	+	ff = db.NewFeatureFlag(c.User.ID, "pgs", cfgMaxSize, 0, 0)
236	236		}
237		-	// this is jank
238		-	ff.Data.StorageMax = ff.FindStorageMax(cfgMaxSize)
239		-	storageMax := ff.Data.StorageMax
	237	+	storageMax := ff.FindStorageMax(cfgMaxSize)
240	238
241	239		bucketName := shared.GetAssetBucketName(c.User.ID)
242	240		bucket, err := c.Store.GetBucket(bucketName)

...	...		@@ -82,7 +83,15 @@ func (me MemoryDB) FindUserByName(name string) (db.User, error) {
82	83		}
83	84
84	85		func (me MemoryDB) FindFeature(userID, name string) (db.FeatureFlag, error) {
85		-	return me.Feature, nil
	86	+	for _, ff := range me.Features {
	87	+	if ff.UserID == userID && ff.Name == name {
	88	+	return ff, nil
	89	+	}
	90	+	}
	91	+	if me.Feature != nil && (me.Feature.Name == name \|\| me.Feature.Name == "") {
	92	+	return me.Feature, nil
	93	+	}
	94	+	return nil, fmt.Errorf("feature flag not found")
86	95		}
87	96
88	97		func (me *MemoryDB) Close() error {

...	...		@@ -273,20 +272,6 @@ func (h UploadAssetHandler) findDenylist(bucket storage.Bucket, project db.Pro
273	272		return str, nil
274	273		}
275	274
276		-	func findPlusFF(dbpool pgsdb.PgsDB, cfg PgsConfig, userID string) (db.FeatureFlag, error) {
277		-	ff, err := dbpool.FindFeature(userID, "plus")
278		-	if err != nil {
279		-	return nil, err
280		-	}
281		-	if !ff.IsValid() {
282		-	return nil, fmt.Errorf("your pico+ has expired")
283		-	}
284		-	ff.Data.StorageMax = ff.FindStorageMax(cfg.MaxSize)
285		-	ff.Data.FileMax = ff.FindFileMax(cfg.MaxAssetSize)
286		-	ff.Data.SpecialFileMax = ff.FindSpecialFileMax(cfg.MaxSpecialFileSize)
287		-	return ff, nil
288		-	}
289		-
290	275		func mtimeToTime(entry *sendutils.FileEntry) time.Time {
291	276		var mtime time.Time
292	277		if entry.Mtime > 0 {

...	...		@@ -557,3 +557,160 @@ func TestNoPlusFlagBlocksUpload(t *testing.T) {
557	557		t.Errorf("expected 0 projects for non-plus user, got %d (empty project leak)", len(projects))
558	558		}
559	559		}
	560	+
	561	+	func 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(50shared.MB), int64(5shared.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(50shared.MB), int64(5shared.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(50shared.MB), int64(5shared.MB), int64(2*shared.KB))
	606	+	plusFF.ExpiresAt = &validExpires
	607	+	pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30shared.MB), int64(3shared.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(50shared.MB), int64(5shared.MB), int64(2*shared.KB))
	623	+	plusFF.ExpiresAt = &expiredExpires
	624	+	pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30shared.MB), int64(3shared.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(50shared.MB), int64(5shared.MB), int64(2*shared.KB))
	640	+	plusFF.ExpiresAt = &expiredExpires
	641	+	pgsFF := db.NewFeatureFlag(userID, "pgs", uint64(30shared.MB), int64(3shared.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	+
	660	+	func 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	+	}

...	...		@@ -1902,6 +1902,17 @@ func (me *PsqlDB) AddPicoPlusUser(username, email, paymentType, txId string) err
1902	1902		return tx.Commit()
1903	1903		}
1904	1904
	1905	+	func (me *PsqlDB) AddFeatureUser(username, name string) error {
	1906	+	user, err := me.FindUserByName(username)
	1907	+	if err != nil {
	1908	+	return err
	1909	+	}
	1910	+
	1911	+	expiresAt := me.createFeatureExpiresAt(user.ID, name)
	1912	+	_, err = me.InsertFeature(user.ID, name, expiresAt)
	1913	+	return err
	1914	+	}
	1915	+
1905	1916		func (me PsqlDB) UpsertProject(userID, projectName, projectDir string) (db.Project, error) {
1906	1917		project, err := me.FindProjectByName(userID, projectName)
1907	1918		if err == nil {

...	...		@@ -0,0 +1,37 @@
	1	+	package main
	2	+
	3	+	import (
	4	+	"log/slog"
	5	+	"os"
	6	+
	7	+	"github.com/picosh/pico/pkg/db/postgres"
	8	+	)
	9	+
	10	+	func main() {
	11	+	logger := slog.Default()
	12	+	DbURL := os.Getenv("DATABASE_URL")
	13	+	dbpool := postgres.NewDB(DbURL, logger)
	14	+
	15	+	args := os.Args
	16	+	if len(args) < 3 {
	17	+	logger.Error("usage: go run ./cmd/scripts/feature-flag <username> <feature>")
	18	+	os.Exit(1)
	19	+	}
	20	+
	21	+	username := args[1]
	22	+	feature := args[2]
	23	+
	24	+	logger.Info(
	25	+	"Adding feature flag to user",
	26	+	"username", username,
	27	+	"feature", feature,
	28	+	)
	29	+
	30	+	err := dbpool.AddFeatureUser(username, feature)
	31	+	if err != nil {
	32	+	logger.Error("Failed to add feature flag to user", "err", err)
	33	+	os.Exit(1)
	34	+	} else {
	35	+	logger.Info("Successfully added feature flag to user")
	36	+	}
	37	+	}

...	...		@@ -62,7 +62,7 @@ func (c *PgsConfig) StaticPath(fname string) string {
62	62		return filepath.Join("pkg", "apps", "pgs", fname)
63	63		}
64	64
65		-	var maxSize = uint64(25 * shared.MB)
	65	+	var maxSize = uint64(50 * shared.MB)
66	66		var maxAssetSize = int64(10 * shared.MB)
67	67
68	68		// Needs to be small for caching files like _headers and _redirects.

...	...		@@ -16,6 +16,7 @@ type MemoryDB struct {
16	16		Projects []*db.Project
17	17		Pubkeys []*db.PublicKey
18	18		Feature *db.FeatureFlag
	19	+	Features []*db.FeatureFlag
19	20		FormEntries []*db.FormEntry
20	21		}
21	22

...	...		@@ -0,0 +1,43 @@
	1	+	package pgs
	2	+
	3	+	import (
	4	+	"fmt"
	5	+	"strings"
	6	+
	7	+	pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
	8	+	"github.com/picosh/pico/pkg/db"
	9	+	)
	10	+
	11	+	func setFeatureLimits(ff db.FeatureFlag, cfg PgsConfig) {
	12	+	ff.Data.StorageMax = ff.FindStorageMax(cfg.MaxSize)
	13	+	ff.Data.FileMax = ff.FindFileMax(cfg.MaxAssetSize)
	14	+	ff.Data.SpecialFileMax = ff.FindSpecialFileMax(cfg.MaxSpecialFileSize)
	15	+	}
	16	+
	17	+	func findFeatureFlag(dbpool pgsdb.PgsDB, cfg PgsConfig, userID string) (db.FeatureFlag, error) {
	18	+	ff, err := dbpool.FindFeature(userID, "plus")
	19	+	if err == nil {
	20	+	if ff.IsValid() {
	21	+	setFeatureLimits(ff, cfg)
	22	+	return ff, nil
	23	+	}
	24	+	err = fmt.Errorf("your pico+ has expired")
	25	+	}
	26	+
	27	+	ffPgs, pgsErr := dbpool.FindFeature(userID, "pgs")
	28	+	if pgsErr == nil {
	29	+	if ffPgs.IsValid() {
	30	+	setFeatureLimits(ffPgs, cfg)
	31	+	return ffPgs, nil
	32	+	}
	33	+	pgsErr = fmt.Errorf("your pgs access has expired")
	34	+	}
	35	+
	36	+	if err != nil && strings.Contains(err.Error(), "expired") {
	37	+	return nil, err
	38	+	}
	39	+	if pgsErr != nil {
	40	+	return nil, pgsErr
	41	+	}
	42	+	return nil, err
	43	+	}

...	...		@@ -15,7 +15,6 @@ import (
15	15		"sync"
16	16		"time"
17	17
18		-	pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
19	18		"github.com/picosh/pico/pkg/db"
20	19		"github.com/picosh/pico/pkg/pssh"
21	20		sendutils "github.com/picosh/pico/pkg/send/utils"

...	...		@@ -219,7 +218,7 @@ func (h UploadAssetHandler) Validate(s pssh.SSHServerConnSession) error {
219	218		return err
220	219		}
221	220
222		-	ff, err := findPlusFF(h.Cfg.DB, h.Cfg, user.ID)
	221	+	ff, err := findFeatureFlag(h.Cfg.DB, h.Cfg, user.ID)
223	222		if err != nil {
224	223		return err
225	224		}

...	...		@@ -609,6 +609,7 @@ type DB interface {
609	609		VisitUrlNotFound(opts SummaryOpts) ([]VisitUrl, error)
610	610
611	611		AddPicoPlusUser(username, email, paymentType, txId string) error
	612	+	AddFeatureUser(username, name string) error
612	613		FindFeature(userID string, feature string) (*FeatureFlag, error)
613	614		FindFeaturesByUser(userID string) ([]*FeatureFlag, error)
614	615		HasFeatureByUser(userID string, feature string) bool

...	...		@@ -1081,6 +1081,28 @@ func TestAddPicoPlusUser(t *testing.T) {
1081	1081		}
1082	1082		}
1083	1083
	1084	+	func TestAddFeatureUser(t *testing.T) {
	1085	+	cleanupTestData(t)
	1086	+
	1087	+	user, _ := testDB.RegisterUser("featureuserowner", "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAI featureuserowner", "comment", "")
	1088	+
	1089	+	err := testDB.AddFeatureUser("featureuserowner", "pgs")
	1090	+	if err != nil {
	1091	+	t.Fatalf("AddFeatureUser failed: %v", err)
	1092	+	}
	1093	+
	1094	+	ff, err := testDB.FindFeature(user.ID, "pgs")
	1095	+	if err != nil {
	1096	+	t.Fatalf("FindFeature failed: %v", err)
	1097	+	}
	1098	+	if ff.Name != "pgs" {
	1099	+	t.Errorf("expected feature name pgs, got %s", ff.Name)
	1100	+	}
	1101	+	if !ff.IsValid() {
	1102	+	t.Errorf("expected feature to be valid")
	1103	+	}
	1104	+	}
	1105	+
1084	1106		// ============ Feed Items Tests ============
1085	1107
1086	1108		func TestInsertFeedItems(t *testing.T) {

...	...		@@ -204,6 +204,10 @@ func (me *StubDB) AddPicoPlusUser(username, email, paymentType, txId string) err
204	204		return errNotImpl
205	205		}
206	206
	207	+	func (me *StubDB) AddFeatureUser(username, name string) error {
	208	+	return errNotImpl
	209	+	}
	210	+
207	211		func (me StubDB) FindUserStats(userID string) (db.UserStats, error) {
208	212		return nil, errNotImpl
209	213		}

Commit 7b5a5aa

Commit `7b5a5aa`