Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 32 additions & 18 deletions cmdutil/line_writer.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,29 +24,38 @@ func newLineWriter(handler OutputLineHandler) *lineWriter {

func (lw *lineWriter) Write(p []byte) (n int, err error) {
lw.mu.Lock()
defer lw.mu.Unlock()

// Write to the actual output first
if lw.output != nil {
n, err = lw.output.Write(p)
if err != nil {
return n, err
}
} else {
n = len(p)
}
// Capture references under the lock so they can't change mid-call.
output := lw.output
handler := lw.handler

// Add to buffer and process complete lines
// Buffer incoming data and extract complete lines while holding the lock.
lw.buf = append(lw.buf, p...)
var lines []string
for {
idx := bytes.IndexByte(lw.buf, '\n')
if idx < 0 {
break
}
line := string(lw.buf[:idx])
lines = append(lines, string(lw.buf[:idx]))
lw.buf = lw.buf[idx+1:]
if lw.handler != nil {
lw.handler(line)
}
lw.mu.Unlock()

// Perform blocking I/O outside the lock to avoid holding it across
// potentially slow writes or handler callbacks.
if output != nil {
n, err = output.Write(p)
if err != nil {
return n, err

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If output.Write(p) fails here, the function returns early but lines were already extracted from lw.buf at line 40. Those lines are gone from the buffer and handlers never fire.

Original code wrote to output before buffering, so a write failure meant nothing was consumed. This version buffers first, then writes.

Fix: move the output write before the buffer manipulation (still outside the lock). Sequence should be:

  1. Lock, capture refs, release lock
  2. output.Write(p) - fail early, nothing consumed yet
  3. Lock, buffer + extract lines, release lock
  4. Call handlers

}
} else {
n = len(p)
}

for _, line := range lines {
if handler != nil {
handler(line)
}
}

Expand All @@ -56,10 +65,15 @@ func (lw *lineWriter) Write(p []byte) (n int, err error) {
// Flush processes any remaining buffered data as a final line.
func (lw *lineWriter) Flush() {
lw.mu.Lock()
defer lw.mu.Unlock()

if len(lw.buf) > 0 && lw.handler != nil {
lw.handler(string(lw.buf))
handler := lw.handler
var remaining string
if len(lw.buf) > 0 {
remaining = string(lw.buf)
lw.buf = nil
}
lw.mu.Unlock()

if remaining != "" && handler != nil {
handler(remaining)
}
}
22 changes: 12 additions & 10 deletions fileutil/fileutil.go
Original file line number Diff line number Diff line change
Expand Up @@ -330,19 +330,21 @@ func LoadCacheJSON(path string, target any, opts CacheOptions) (valid bool, err
return false, fmt.Errorf("failed to read cache file: %w", err)
}

if err := json.Unmarshal(data, target); err != nil {
return false, fmt.Errorf("failed to parse cache JSON: %w", err)
}

// Try to extract metadata for validation
// We need to re-parse to get the metadata
// Validate metadata first (lightweight partial unmarshal) to avoid
// deserializing the full target when the cache is expired.
var metaWrapper struct {
Cache CacheMetadata `json:"_cache"`
}
if err := json.Unmarshal(data, &metaWrapper); err == nil {
if !metaWrapper.Cache.IsCacheValid(opts) {
return false, nil // Cache is invalid (expired or version mismatch)
}
if err := json.Unmarshal(data, &metaWrapper); err != nil {
return false, fmt.Errorf("failed to parse cache JSON: %w", err)
}
if !metaWrapper.Cache.IsCacheValid(opts) {
return false, nil // Cache is invalid (expired or version mismatch)
}

// Cache is valid - deserialize the full target.
if err := json.Unmarshal(data, target); err != nil {
return false, fmt.Errorf("failed to parse cache JSON: %w", err)
}

return true, nil
Expand Down
70 changes: 44 additions & 26 deletions httpclient/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,18 @@ type Client struct {
}

// NewClient creates a new HTTP client
// redirectContextKeyType is used to pass per-request redirect configuration
// through request context, allowing a single http.Client to handle different
// redirect policies without per-call allocation.
type redirectContextKeyType struct{}

var redirectContextKey = redirectContextKeyType{}

type redirectConfig struct {
followRedirects bool
maxRedirects int
}

func NewClient(tokenProvider TokenProvider, insecure bool, timeout time.Duration) *Client {
transport := &http.Transport{
TLSClientConfig: &tls.Config{
Expand All @@ -116,6 +128,16 @@ func NewClient(tokenProvider TokenProvider, insecure bool, timeout time.Duration
httpClient: &http.Client{
Transport: transport,
Timeout: timeout,
CheckRedirect: func(req *http.Request, via []*http.Request) error {
cfg, _ := req.Context().Value(redirectContextKey).(redirectConfig)
if !cfg.followRedirects {
return http.ErrUseLastResponse
}
if cfg.maxRedirects > 0 && len(via) >= cfg.maxRedirects {
return fmt.Errorf("stopped after %d redirects", cfg.maxRedirects)
}
return nil
},
},
tokenProvider: tokenProvider,
}
Expand All @@ -125,34 +147,25 @@ func NewClient(tokenProvider TokenProvider, insecure bool, timeout time.Duration
func (c *Client) Execute(ctx context.Context, opts RequestOptions) (*Response, error) {
startTime := time.Now()

// Configure redirect handling
// Store per-request redirect config in context instead of allocating a new http.Client.
maxRedirects := opts.MaxRedirects
if maxRedirects == 0 {
maxRedirects = DefaultMaxRedirects
}

originalClient := c.httpClient
client := &http.Client{
Transport: originalClient.Transport,
Timeout: originalClient.Timeout,
CheckRedirect: func(req *http.Request, via []*http.Request) error {
if !opts.FollowRedirects {
return http.ErrUseLastResponse
}
// via contains all previous requests including the original, so len(via) is the redirect count
if len(via) >= maxRedirects {
return fmt.Errorf("stopped after %d redirects", maxRedirects)
}
return nil
},
}
client := c.httpClient

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

handlePagination(ctx, client, opts, response) at ~line 291 receives ctx (the original function parameter), not req.Context() (which carries the redirectConfig). Since client is now the shared client with context-based CheckRedirect, pagination requests won't have the redirect config and will default to blocking all redirects.

Probably not an issue in practice (Azure pagination URLs don't redirect), but it's a semantic regression from the original code where the per-call client closure captured opts.FollowRedirects.

Fix: handlePagination(req.Context(), client, opts, response)


// Create request
req, err := http.NewRequestWithContext(ctx, opts.Method, opts.URL, opts.Body)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}

// Inject per-request redirect policy via context.
req = req.WithContext(context.WithValue(req.Context(), redirectContextKey, redirectConfig{
followRedirects: opts.FollowRedirects,
maxRedirects: maxRedirects,
}))

// Add custom headers
for key, value := range opts.Headers {
req.Header.Set(key, value)
Expand Down Expand Up @@ -460,13 +473,9 @@ func parseLinkHeader(linkHeader string) (string, bool) {
return "", false
}

// extractNextLinkFromBody extracts nextLink from JSON response body (Azure API format)
func extractNextLinkFromBody(body []byte) (string, bool) {
var data map[string]any
if err := json.Unmarshal(body, &data); err != nil {
return "", false
}

// extractNextLinkFromParsed extracts nextLink from an already-parsed JSON map,
// avoiding redundant deserialization when the caller has already unmarshaled the body.
func extractNextLinkFromParsed(data map[string]any) (string, bool) {
if nextLink, ok := data["nextLink"].(string); ok && nextLink != "" {
return nextLink, true
}
Expand Down Expand Up @@ -497,6 +506,15 @@ var ErrPaginationSizeLimitExceeded = fmt.Errorf("pagination aggregate size limit
// exceeds the maximum page count.
var ErrPaginationPageLimitExceeded = fmt.Errorf("pagination page count limit exceeded")

// extractNextLinkFromBody extracts nextLink from JSON response body (Azure API format)
func extractNextLinkFromBody(body []byte) (string, bool) {
var data map[string]any
if err := json.Unmarshal(body, &data); err != nil {
return "", false
}
return extractNextLinkFromParsed(data)
}

// handlePagination handles pagination by following next links.
// It enforces same-origin checks to prevent SSRF via server-controlled nextLink URLs.
// It also enforces aggregate size and page count limits to prevent unbounded memory growth.
Expand Down Expand Up @@ -550,7 +568,7 @@ func handlePagination(ctx context.Context, client *http.Client, opts RequestOpti
allResults = append(allResults, firstData)
}

if next, ok := extractNextLinkFromBody(currentBody); ok {
if next, ok := extractNextLinkFromParsed(firstData); ok {
nextURL = next
hasMore = true
} else if linkHeader := firstResponse.Headers.Get("Link"); linkHeader != "" {
Expand Down Expand Up @@ -637,7 +655,7 @@ func handlePagination(ctx context.Context, client *http.Client, opts RequestOpti
}

nextURL = ""
if next, ok := extractNextLinkFromBody(body); ok {
if next, ok := extractNextLinkFromParsed(pageData); ok {
nextURL = next
} else if linkHeader := resp.Header.Get("Link"); linkHeader != "" {
if next, ok := parseLinkHeader(linkHeader); ok {
Expand Down
21 changes: 17 additions & 4 deletions progress/progress.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ type MultiProgress struct {
mu sync.RWMutex
stopChan chan struct{}
stopped bool
started bool
lastLineCount int
termWidth int
}
Expand Down Expand Up @@ -147,15 +148,27 @@ func (mp *MultiProgress) GetBar(id string) *ProgressSpinner {
}

// Start starts the multi-progress display (renders all bars periodically).
// It is safe to call multiple times; redundant calls are no-ops unless Stop()
// was called in between (in which case a new render loop is started).
func (mp *MultiProgress) Start() {
// Hide cursor during progress display
fmt.Print("\033[?25l")

// Set initial line count based on number of bars
mp.mu.Lock()
if mp.started && !mp.stopped {
// Already running - avoid leaking another goroutine.
mp.mu.Unlock()
return
}
if mp.stopped {
// Restarting after a stop - create a fresh channel.
mp.stopChan = make(chan struct{})
mp.stopped = false
}
mp.started = true
mp.lastLineCount = len(mp.bars)
mp.mu.Unlock()

// Hide cursor during progress display
fmt.Print("\033[?25l")

go func() {
ticker := time.NewTicker(refreshInterval)
defer ticker.Stop()
Expand Down
Loading