executeCommandWithTimeout executes a command with timeout support
(ctx context.Context, session *gossh.Session, command string)
| 276 | |
| 277 | // executeCommandWithTimeout executes a command with timeout support |
| 278 | func executeCommandWithTimeout(ctx context.Context, session *gossh.Session, command string) ([]byte, error) { |
| 279 | // Set up pipes to capture output |
| 280 | stdoutPipe, err := session.StdoutPipe() |
| 281 | if err != nil { |
| 282 | return nil, xerrors.Errorf("failed to create stdout pipe: %w", err) |
| 283 | } |
| 284 | |
| 285 | stderrPipe, err := session.StderrPipe() |
| 286 | if err != nil { |
| 287 | return nil, xerrors.Errorf("failed to create stderr pipe: %w", err) |
| 288 | } |
| 289 | |
| 290 | // Start the command |
| 291 | if err := session.Start(command); err != nil { |
| 292 | return nil, xerrors.Errorf("failed to start command: %w", err) |
| 293 | } |
| 294 | |
| 295 | // Create a thread-safe buffer for combined output |
| 296 | var output bytes.Buffer |
| 297 | var mu sync.Mutex |
| 298 | safeWriter := &syncWriter{w: &output, mu: &mu} |
| 299 | |
| 300 | // Use io.MultiWriter to combine stdout and stderr |
| 301 | multiWriter := io.MultiWriter(safeWriter) |
| 302 | |
| 303 | // Channel to signal when command completes |
| 304 | done := make(chan error, 1) |
| 305 | |
| 306 | // Start goroutine to copy output and wait for completion |
| 307 | go func() { |
| 308 | // Copy stdout and stderr concurrently |
| 309 | var wg sync.WaitGroup |
| 310 | wg.Add(2) |
| 311 | |
| 312 | go func() { |
| 313 | defer wg.Done() |
| 314 | _, _ = io.Copy(multiWriter, stdoutPipe) |
| 315 | }() |
| 316 | |
| 317 | go func() { |
| 318 | defer wg.Done() |
| 319 | _, _ = io.Copy(multiWriter, stderrPipe) |
| 320 | }() |
| 321 | |
| 322 | // Wait for all output to be copied |
| 323 | wg.Wait() |
| 324 | |
| 325 | // Wait for the command to complete |
| 326 | done <- session.Wait() |
| 327 | }() |
| 328 | |
| 329 | // Wait for either completion or context cancellation |
| 330 | select { |
| 331 | case err := <-done: |
| 332 | // Command completed normally |
| 333 | return safeWriter.Bytes(), err |
| 334 | case <-ctx.Done(): |
| 335 | // Context was canceled (timeout or other cancellation) |