|
1 | 1 | package apidump
|
2 | 2 |
|
3 | 3 | import (
|
| 4 | + "math/rand" |
| 5 | + "os" |
| 6 | + "strconv" |
4 | 7 | "time"
|
5 | 8 |
|
6 | 9 | "github.com/akitasoftware/akita-libs/akid"
|
@@ -44,15 +47,56 @@ var (
|
44 | 47 | maxWitnessSize_bytes int
|
45 | 48 | dockerExtensionMode bool
|
46 | 49 | healthCheckPort int
|
| 50 | + randomizedStart int |
47 | 51 | )
|
48 | 52 |
|
| 53 | +// This function will either startup apidump normally, or never return, with probability |
| 54 | +// determined by randomizedStart/100. The value of the command-line flag may be |
| 55 | +// overridden by an environment variable to make it easier to apply. |
| 56 | +// This function should be called as early as possible in case termination of |
| 57 | +// the agent causes deployment problems. |
| 58 | +// |
| 59 | +// Negative values are effectively treated as 0 probability, instead of being validated. |
| 60 | +func applyRandomizedStart() { |
| 61 | + prob := randomizedStart |
| 62 | + |
| 63 | + if env := os.Getenv("POSTMAN_AGENT_RANDOM_START"); env != "" { |
| 64 | + override, err := strconv.Atoi(env) |
| 65 | + if err == nil { |
| 66 | + prob = override |
| 67 | + } |
| 68 | + } |
| 69 | + |
| 70 | + if prob < 100 { |
| 71 | + printer.Stdout.Infof("Starting trace with probability %d%%.\n", prob) |
| 72 | + |
| 73 | + // Pre-1.20, Go does not seed the default Random object :( |
| 74 | + rng := rand.New(rand.NewSource(time.Now().UnixNano())) |
| 75 | + r := rng.Intn(100) // in range [0,100), |
| 76 | + // so 1% probability means < 1, not <= 1: |
| 77 | + if r < prob { |
| 78 | + return |
| 79 | + } |
| 80 | + |
| 81 | + printer.Stdout.Infof("This agent instance will not begin capturing data.\n") |
| 82 | + |
| 83 | + // Wait forever |
| 84 | + select {} |
| 85 | + |
| 86 | + // Unreachable |
| 87 | + os.Exit(0) |
| 88 | + } |
| 89 | +} |
| 90 | + |
49 | 91 | var Cmd = &cobra.Command{
|
50 | 92 | Use: "apidump",
|
51 | 93 | Short: "Capture requests/responses from network traffic.",
|
52 | 94 | Long: "Capture and store a sequence of requests/responses to a service by observing network traffic.",
|
53 | 95 | SilenceUsage: true,
|
54 | 96 | Args: cobra.ExactArgs(0),
|
55 | 97 | RunE: func(cmd *cobra.Command, _ []string) error {
|
| 98 | + applyRandomizedStart() |
| 99 | + |
56 | 100 | traceTags, err := util.ParseTagsAndWarn(tagsFlag)
|
57 | 101 | if err != nil {
|
58 | 102 | return err
|
@@ -344,4 +388,12 @@ func init() {
|
344 | 388 | "Port to listen on for Docker extension health checks. This is an internal flag used by the Akita Docker extension.",
|
345 | 389 | )
|
346 | 390 | _ = Cmd.Flags().MarkHidden("health-check-port")
|
| 391 | + |
| 392 | + Cmd.Flags().IntVar( |
| 393 | + &randomizedStart, |
| 394 | + "randomized-start", |
| 395 | + 100, |
| 396 | + "Probability that the apidump command will start intercepting traffic.", |
| 397 | + ) |
| 398 | + _ = Cmd.Flags().MarkHidden("randomized-start") |
347 | 399 | }
|
0 commit comments