add limits to websharing

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
 		}

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)

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"