diff --git a/pkg/cmdrunner/cmdrunner.go b/pkg/cmdrunner/cmdrunner.go index 1b97fd2c..1b1a72b0 100644 --- a/pkg/cmdrunner/cmdrunner.go +++ b/pkg/cmdrunner/cmdrunner.go @@ -1704,7 +1704,7 @@ func ScreenWebShareCommand(ctx context.Context, pk *scpacket.FeCommandPacketType } viewKey := base64.RawURLEncoding.EncodeToString(viewKeyBytes) webShareOpts := sstore.ScreenWebShareOpts{ShareName: shareName, ViewKey: viewKey} - err = sstore.CanScreenWebShare(screen) + err = sstore.CanScreenWebShare(ctx, screen) if err != nil { return nil, err } diff --git a/pkg/sstore/dbops.go b/pkg/sstore/dbops.go index 3d6f9833..c4ca7cc9 100644 --- a/pkg/sstore/dbops.go +++ b/pkg/sstore/dbops.go @@ -2437,6 +2437,54 @@ func PurgeHistoryByIds(ctx context.Context, historyIds []string) ([]*HistoryItem }) } +func CountScreenWebShares(ctx context.Context) (int, error) { + return WithTxRtn(ctx, func(tx *TxWrap) (int, error) { + query := `SELECT count(*) FROM screen WHERE sharemode = ?` + count := tx.GetInt(query, ShareModeWeb) + return count, nil + }) +} + +func CountScreenLines(ctx context.Context, screenId string) (int, error) { + return WithTxRtn(ctx, func(tx *TxWrap) (int, error) { + query := `SELECT count(*) FROM line WHERE screenid = ? AND NOT archived` + lineCount := tx.GetInt(query, screenId) + return lineCount, nil + }) +} + +func CanScreenWebShare(ctx context.Context, screen *ScreenType) error { + if screen == nil { + return fmt.Errorf("cannot share screen, not found") + } + if screen.ShareMode == ShareModeWeb { + return fmt.Errorf("screen is already shared to web") + } + if screen.ShareMode != ShareModeLocal { + return fmt.Errorf("screen cannot be shared, invalid current share mode %q (must be local)", screen.ShareMode) + } + if screen.Archived { + return fmt.Errorf("screen cannot be shared, must un-archive before sharing") + } + webShareCount, err := CountScreenWebShares(ctx) + if err != nil { + return fmt.Errorf("screen cannot be share: error getting webshare count: %v", err) + } + if webShareCount >= MaxWebShareScreenCount { + go UpdateCurrentActivity(context.Background(), ActivityUpdate{WebShareLimit: 1}) + return fmt.Errorf("screen cannot be shared, limited to a maximum of %d shared screen(s)", MaxWebShareScreenCount) + } + lineCount, err := CountScreenLines(ctx, screen.ScreenId) + if err != nil { + return fmt.Errorf("screen cannot be share: error getting screen line count: %v", err) + } + if lineCount > MaxWebShareLineCount { + go UpdateCurrentActivity(context.Background(), ActivityUpdate{WebShareLimit: 1}) + return fmt.Errorf("screen cannot be shared, limited to a maximum of %d lines", MaxWebShareLineCount) + } + return nil +} + func ScreenWebShareStart(ctx context.Context, screenId string, shareOpts ScreenWebShareOpts) error { return WithTx(ctx, func(tx *TxWrap) error { query := `SELECT screenid FROM screen WHERE screenid = ?` @@ -2450,12 +2498,6 @@ func ScreenWebShareStart(ctx context.Context, screenId string, shareOpts ScreenW if shareMode != ShareModeLocal { return fmt.Errorf("screen cannot be shared, invalid current share mode %q (must be local)", shareMode) } - query = `SELECT count(*) FROM line WHERE screenid = ? AND NOT archived` - lineCount := tx.GetInt(query, screenId) - if lineCount > MaxWebShareLineCount { - return fmt.Errorf("screen cannot be shared, limited to a maximum of %d lines", MaxWebShareLineCount) - go UpdateCurrentActivity(context.Background(), ActivityUpdate{WebShareLimit: 1}) - } query = `UPDATE screen SET sharemode = ?, webshareopts = ? WHERE screenid = ?` tx.Exec(query, ShareModeWeb, quickJson(shareOpts), screenId) insertScreenNewUpdate(tx, screenId) diff --git a/pkg/sstore/sstore.go b/pkg/sstore/sstore.go index 70df6358..3bbad9e3 100644 --- a/pkg/sstore/sstore.go +++ b/pkg/sstore/sstore.go @@ -34,6 +34,7 @@ const LineNoHeight = -1 const DBFileName = "prompt.db" const DBFileNameBackup = "backup.prompt.db" const MaxWebShareLineCount = 50 +const MaxWebShareScreenCount = 5 const DefaultSessionName = "default" const LocalRemoteAlias = "local"