add limits for confluence

Author: cx-leonardo-fontes

plugins/confluence.go

@@ -27,6 +27,10 @@ const (
 	flagUsername  = "username"
 	flagToken     = "token"
 	flagHistory   = "history"
+
+	flagMaxAPIResponseMB = "max-api-response-megabytes"
+	flagMaxPageBodyMB    = "max-page-body-megabytes"
+	flagMaxTotalScanMB   = "max-total-scan-megabytes"
 )
 
 // Confluence Cloud REST API v2 per-request limits (server caps by endpoint/param).
@@ -46,6 +50,9 @@ const (
 	// maxPageSize is the requested number of items per page in paginated responses.
 	// Confluence v2 accepts 1–250; we use 250 to minimize requests so we're less likely to hit rate limits.
 	maxPageSize = 250
+
+	// bytesPerMegabyte is used to convert CLI megabyte values into bytes.
+	bytesPerMegabyte = 1024 * 1024
 )
 
 type ConfluencePlugin struct {
@@ -82,6 +89,11 @@ type ConfluencePlugin struct {
 
 	// invalidPageIDs holds user-provided page-ids that failed numeric/length validation.
 	invalidPageIDs map[string]struct{}
+
+	// Optional limits (0 means "no limit").
+	maxAPIResponseBytes int64
+	maxTotalScanBytes   int64
+	maxPageBodyBytes    int64
 }
 
 // NewConfluencePlugin constructs a new Confluence plugin with a default chunker.
@@ -107,6 +119,11 @@ func (p *ConfluencePlugin) DefineCommand(items chan ISourceItem, errs chan error
 	var username string
 	var token string
 
+	// CLI values in megabytes; converted to bytes in PreRunE.
+	var maxAPIResponseMB int
+	var maxPageBodyMB int
+	var maxTotalScanMB int
+
 	cmd := &cobra.Command{
 		Use:     fmt.Sprintf("%s <URL>", p.GetName()),
 		Short:   "Scan Confluence Cloud",
@@ -118,9 +135,21 @@ func (p *ConfluencePlugin) DefineCommand(items chan ISourceItem, errs chan error
 			p.SpaceIDs = trimNonEmpty(p.SpaceIDs)
 			p.PageIDs = trimNonEmpty(p.PageIDs)
 
+			// Convert per-run/response/page limits from megabytes to bytes.
+			if maxAPIResponseMB > 0 {
+				p.maxAPIResponseBytes = int64(maxAPIResponseMB) * bytesPerMegabyte
+			}
+			if maxTotalScanMB > 0 {
+				p.maxTotalScanBytes = int64(maxTotalScanMB) * bytesPerMegabyte
+			}
+			if maxPageBodyMB > 0 {
+				p.maxPageBodyBytes = int64(maxPageBodyMB) * bytesPerMegabyte
+			}
+
 			if err := p.initialize(args[0], username, token); err != nil {
 				return err
 			}
+
 			if username == "" || token == "" {
 				log.Warn().Msg("Confluence credentials not provided. The scan will run anonymously (public pages only).")
 			}
@@ -144,6 +173,14 @@ func (p *ConfluencePlugin) DefineCommand(items chan ISourceItem, errs chan error
 	flags.StringVar(&token, flagToken, "", "Confluence API/scoped token value.")
 	flags.BoolVar(&p.History, flagHistory, false, "Also scan all page revisions (all versions).")
 
+	// Optional limits (0 disables each check).
+	flags.IntVar(&maxAPIResponseMB, flagMaxAPIResponseMB, 0,
+		"limit for per-request API response size in megabytes. When exceeded, the batch is skipped and a warning is logged (0 disables the check).")
+	flags.IntVar(&maxPageBodyMB, flagMaxPageBodyMB, 0,
+		"limit for individual page body size in megabytes. Pages above this size are skipped and logged as warnings (0 disables the check).")
+	flags.IntVar(&maxTotalScanMB, flagMaxTotalScanMB, 0,
+		"limit for total downloaded data in megabytes for this Confluence scan. When exceeded, scanning stops gracefully and logs a warning (0 disables the check).")
+
 	return cmd, nil
 }
 
@@ -185,7 +222,12 @@ func isValidURL(_ *cobra.Command, args []string) error {
 // initialize stores the normalized base wiki URL and constructs the Confluence client.
 // It validates the base by discovering the cloudId; if discovery fails, init fails.
 func (p *ConfluencePlugin) initialize(base, username, token string) error {
-	client, err := NewConfluenceClient(base, username, token)
+	client, err := NewConfluenceClient(
+		base,
+		username,
+		token,
+		WithLimits(p.maxAPIResponseBytes, p.maxTotalScanBytes, p.maxPageBodyBytes),
+	)
 	if err != nil {
 		return err
 	}
@@ -198,34 +240,46 @@ func (p *ConfluencePlugin) initialize(base, username, token string) error {
 // Also tracks which selectors yielded results and logs a single consolidated warning for
 // invalid/unresolvable/inaccessible selectors, without issuing additional API requests.
 func (p *ConfluencePlugin) walkAndEmitPages(ctx context.Context) error {
+	handleTotalLimit := func(err error) error {
+		if errors.Is(err, ErrTotalScanVolumeExceeded) {
+			log.Warn().Msg("Some Confluence pages could not be processed because the scan reached the configured total-volume limit.")
+			return nil
+		}
+		return err
+	}
+
 	if len(p.SpaceIDs) > 0 || len(p.SpaceKeys) > 0 {
 		allSpaceIDs, err := p.resolveAndCollectSpaceIDs(ctx)
 		if err != nil {
-			return err
+			return handleTotalLimit(err)
 		}
 		if err := p.scanBySpaceIDs(ctx, allSpaceIDs); err != nil {
-			return err
+			return handleTotalLimit(err)
 		}
 	}
 
 	if len(p.PageIDs) > 0 {
 		if err := p.scanByPageIDs(ctx); err != nil {
-			return err
+			return handleTotalLimit(err)
 		}
 	}
 
 	if len(p.SpaceIDs) == 0 && len(p.SpaceKeys) == 0 && len(p.PageIDs) == 0 {
 		if err := p.client.WalkAllPages(ctx, maxPageSize, func(page *Page) error {
 			return p.emitUniquePage(ctx, page)
 		}); err != nil {
-			return err
+			return handleTotalLimit(err)
 		}
 	}
 
 	if msg := p.missingSelectorsWarningMessage(); msg != "" {
 		log.Warn().Msg(msg)
 	}
 
+	if p.client.APIResponseLimitHit() {
+		log.Warn().Msg("Some Confluence pages could not be processed because one or more API responses exceeded the configured size limit.")
+	}
+
 	return nil
 }
 

plugins/confluence_client.go

@@ -12,6 +12,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/rs/zerolog/log"
 )
 
 const (
@@ -33,6 +35,16 @@ var (
 
 	// ErrUnexpectedHTTPStatus is used for other non-2xx statuses not classified above.
 	ErrUnexpectedHTTPStatus = errors.New("unexpected http status")
+
+	// ErrAPIResponseTooLarge is returned internally when a single API response exceeds
+	// the configured per-request size limit. Callers treat this as a soft failure:
+	// skip the current batch and move on when possible.
+	ErrAPIResponseTooLarge = errors.New("confluence api response exceeded configured size limit")
+
+	// ErrTotalScanVolumeExceeded is returned when the total downloaded bytes across
+	// the scan exceed the configured global limit. Callers should stop scanning
+	// gracefully and not treat this as a hard failure.
+	ErrTotalScanVolumeExceeded = errors.New("confluence total scan volume exceeded configured limit")
 )
 
 // ConfluenceClient defines the operations required by the Confluence plugin.
@@ -45,6 +57,7 @@ type ConfluenceClient interface {
 	FetchPageAtVersion(ctx context.Context, pageID string, version int) (*Page, error)
 	WalkSpacesByKeys(ctx context.Context, spaceKeys []string, limit int, visit func(*Space) error) error
 	WikiBaseURL() string
+	APIResponseLimitHit() bool
 }
 
 // httpConfluenceClient is a ConfluenceClient implementation backed by net/http.
@@ -55,14 +68,40 @@ type httpConfluenceClient struct {
 	username    string
 	token       string
 	apiBase     string
+
+	// Optional limits (0 means "no limit").
+	maxAPIResponseBytes int64
+	maxTotalScanBytes   int64
+	maxPageBodyBytes    int64
+
+	totalScanBytes int64
+
+	// apiResponseLimitHit indicates that at least one API batch exceeded the
+	// configured per-request size limit during this scan. Used by the plugin
+	// to emit a single consolidated warning instead of one per batch.
+	apiResponseLimitHit bool
+}
+
+type ConfluenceClientOption func(*httpConfluenceClient)
+
+// WithLimits configures optional per-response, total-scan, and page-body size
+// limits in bytes for the Confluence client.
+func WithLimits(maxAPIResponseBytes, maxTotalScanBytes, maxPageBodyBytes int64) ConfluenceClientOption {
+	return func(c *httpConfluenceClient) {
+		c.maxAPIResponseBytes = maxAPIResponseBytes
+		c.maxTotalScanBytes = maxTotalScanBytes
+		c.maxPageBodyBytes = maxPageBodyBytes
+	}
 }
 
 // NewConfluenceClient constructs a ConfluenceClient for the given base input URL.
 // Behavior:
 //   - Normalizes base wiki URL to "https://{host}/wiki".
 //   - Discovers cloudId and always uses platform v2:
 //     "https://api.atlassian.com/ex/confluence/{cloudId}/wiki/api/v2".
-func NewConfluenceClient(inputBaseURL, username, token string) (ConfluenceClient, error) {
+//   - Applies any provided ConfluenceClientOption values (for example, response/scan
+//     size limits) before issuing API requests.
+func NewConfluenceClient(inputBaseURL, username, token string, opts ...ConfluenceClientOption) (ConfluenceClient, error) {
 	baseWikiURL, err := normalizeWikiBase(inputBaseURL)
 	if err != nil {
 		return nil, err
@@ -73,6 +112,9 @@ func NewConfluenceClient(inputBaseURL, username, token string) (ConfluenceClient
 		username:    username,
 		token:       token,
 	}
+	for _, opt := range opts {
+		opt(c)
+	}
 	apiBase, err := c.buildAPIBase(context.Background())
 	if err != nil {
 		return nil, err
@@ -85,6 +127,12 @@ func NewConfluenceClient(inputBaseURL, username, token string) (ConfluenceClient
 // Example: "https://tenant.atlassian.net/wiki".
 func (c *httpConfluenceClient) WikiBaseURL() string { return c.baseWikiURL }
 
+// APIResponseLimitHit reports whether any API response exceeded the configured
+// per-request size limit during this scan.
+func (c *httpConfluenceClient) APIResponseLimitHit() bool {
+	return c.apiResponseLimitHit
+}
+
 // normalizeWikiBase takes any Confluence-related URL (site root, /wiki, a page URL, etc.)
 // and returns "https://{host}/wiki".
 func normalizeWikiBase(inputURL string) (string, error) {
@@ -300,17 +348,87 @@ func (c *httpConfluenceClient) walkPagesPaginated(
 		if err != nil {
 			return err
 		}
-		// Prefer Link header; body may also include _links.next.
+
+		// Prefer Link header to discover the next cursor so we can safely move on
+		// even if we have to stop decoding this batch early.
 		linkNext := nextURLFromLinkHeader(headers)
-		bodyNext, decodeErr := streamPagesFromBody(rc, visit)
+
+		// Wrap the body in a counting reader to enforce:
+		//   - per-response size limit (maxAPIResponseBytes)
+		//   - total scan volume limit (maxTotalScanBytes)
+		var reader io.Reader = rc
+		if c.maxAPIResponseBytes > 0 || c.maxTotalScanBytes > 0 {
+			reader = &countingReader{
+				r:               rc,
+				apiLimitBytes:   c.maxAPIResponseBytes,
+				apiConsumed:     0,
+				totalBytes:      &c.totalScanBytes,
+				totalLimitBytes: c.maxTotalScanBytes,
+			}
+		}
+
+		// Wrap the visitor to apply per-page body-size limit when configured.
+		limitedVisit := visit
+		if c.maxPageBodyBytes > 0 {
+			underlying := visit
+			limitedVisit = func(p *Page) error {
+				if p.Body.Storage != nil {
+					bodySize := int64(len(p.Body.Storage.Value))
+					if bodySize > c.maxPageBodyBytes {
+						log.Warn().
+							Str("page_id", p.ID).
+							Int64("body_bytes", bodySize).
+							Msg("Skipping page content because the Confluence page body exceeded the configured size limit.")
+						return nil
+					}
+				}
+				return underlying(p)
+			}
+		}
+
+		bodyNext, decodeErr := streamPagesFromBody(reader, limitedVisit)
 		closeErr := rc.Close()
+
 		if decodeErr != nil {
-			return decodeErr
+			switch {
+			case errors.Is(decodeErr, ErrAPIResponseTooLarge):
+				// record that at least one batch exceeded the per-request
+				// size limit and move to the next page if we know the cursor.
+				// The plugin will emit a single consolidated warning after the scan.
+				c.apiResponseLimitHit = true
+
+				if closeErr != nil {
+					return closeErr
+				}
+
+				cur := cursorFromURL(linkNext)
+				if cur == "" {
+					// No "next" link: nothing else to do.
+					return nil
+				}
+				nextURL = withCursor(base, cur)
+				continue
+
+			case errors.Is(decodeErr, ErrTotalScanVolumeExceeded):
+				// Global total-volume limit hit: stop scanning gracefully.
+				if closeErr != nil {
+					return closeErr
+				}
+				return ErrTotalScanVolumeExceeded
+
+			default:
+				if closeErr != nil {
+					return closeErr
+				}
+				return decodeErr
+			}
 		}
+
 		if closeErr != nil {
 			return closeErr
 		}
-		// Extract only the cursor and rebuild the next URL from our base.
+
+		// use Link header when available, fall back to body _links.next.
 		cur := firstNonEmptyString(cursorFromURL(linkNext), cursorFromURL(bodyNext))
 		if cur == "" {
 			return nil
@@ -366,6 +484,15 @@ func (c *httpConfluenceClient) getJSON(ctx context.Context, reqURL string) ([]by
 	if err != nil {
 		return nil, nil, fmt.Errorf("read body: %w", err)
 	}
+
+	// Track total scan volume for non-streamed responses as well.
+	if c.maxTotalScanBytes > 0 {
+		c.totalScanBytes += int64(len(body))
+		if c.totalScanBytes > c.maxTotalScanBytes {
+			return nil, nil, ErrTotalScanVolumeExceeded
+		}
+	}
+
 	return body, resp.Header.Clone(), nil
 }
 
@@ -679,3 +806,43 @@ func firstNonEmptyString(primary, fallback string) string {
 	}
 	return fallback
 }
+
+// countingReader wraps an io.Reader and tracks:
+//   - bytes read for the current API response, enforcing maxAPIResponseBytes
+//   - global total bytes read across the scan, enforcing totalLimitBytes.
+//
+// When a limit is exceeded, subsequent reads return a sentinel error and allow
+// callers (streaming decoders) to bail out gracefully.
+type countingReader struct {
+	r               io.Reader
+	apiLimitBytes   int64
+	apiConsumed     int64
+	totalBytes      *int64
+	totalLimitBytes int64
+}
+
+func (cr *countingReader) Read(p []byte) (int, error) {
+	n, err := cr.r.Read(p)
+	if n <= 0 {
+		return n, err
+	}
+
+	if cr.apiLimitBytes > 0 {
+		cr.apiConsumed += int64(n)
+		if cr.apiConsumed > cr.apiLimitBytes {
+			// Surface a sentinel error to the decoder so callers can classify it
+			// as warn for this batch.
+			return n, ErrAPIResponseTooLarge
+		}
+	}
+
+	if cr.totalBytes != nil {
+		*cr.totalBytes += int64(n)
+		if cr.totalLimitBytes > 0 && *cr.totalBytes > cr.totalLimitBytes {
+			// Global total-volume limit exceeded.
+			return n, ErrTotalScanVolumeExceeded
+		}
+	}
+
+	return n, err
+}