Make NetBlockRandomIPv4 responsible for tracking number of random addrs

TargetGroup.cc

@@ -111,6 +111,7 @@ class NetBlock {
  * the return value to the pointer that this method was called through.
  * On error, return NULL. */
  virtual NetBlock *resolve() { return this; }
+ virtual void reject_last_host() {}
  virtual bool next(struct sockaddr_storage *ss, size_t *sslen) = 0;
  virtual void apply_netmask(int bits) = 0;
  virtual std::string str() const = 0;
@@ -120,12 +121,15 @@ class NetBlockRandomIPv4 : public NetBlock {
 public:
  NetBlockRandomIPv4();
 
+ void reject_last_host() { count++; }
+ void set_num_random(int num) { count = num; }
  bool next(struct sockaddr_storage *ss, size_t *sslen);
  void apply_netmask(int bits) {}
  std::string str() const {return "Random IPv4 addresses";}
 
 private:
  struct sockaddr_in base;
+ int count;
 };
 
 class NetBlockIPv4Ranges : public NetBlock {
@@ -337,17 +341,21 @@ bool NetBlock::is_resolved_address(const struct sockaddr_storage *ss) const {
  return false;
 }
 
-NetBlockRandomIPv4::NetBlockRandomIPv4() {
+NetBlockRandomIPv4::NetBlockRandomIPv4() : count(0) {
  memset(&base, 0, sizeof(base));
  base.sin_family = AF_INET;
 }
 
 bool NetBlockRandomIPv4::next(struct sockaddr_storage *ss, size_t *sslen) {
+ if (count <= 0) {
+ return false;
+ }
  do {
  base.sin_addr.s_addr = get_random_unique_u32();
  } while (ip_is_reserved(&base.sin_addr));
  memcpy(ss, &base, sizeof(base));
  *sslen = sizeof(base);
+ count--;
  return true;
 }
 
@@ -803,9 +811,10 @@ int TargetGroup::parse_expr(const char *target_expr, int af) {
  return 1;
 }
 
-void TargetGroup::generate_random_ips() {
+void TargetGroup::generate_random_ips(int num_random) {
  assert(this->netblock == NULL);
  this->netblock = new NetBlockRandomIPv4();
+ this->netblock->set_num_random(num_random);
 }
 
 /* Grab the next host from this expression (if any) and updates its internal

TargetGroup.h

@@ -101,7 +101,11 @@ class TargetGroup {
  const std::list<struct sockaddr_storage> &get_unscanned_addrs(void) const;
  /* is the current expression a named host */
  int get_namedhost() const;
- void generate_random_ips();
+ void generate_random_ips(int num_random);
+ void reject_last_host();
+
+ private:
+ std::list<NetBlock *>netblocks;
 };
 
 #endif /* TARGETGROUP_H */

targets.cc

@@ -296,9 +296,8 @@ HostGroupState::HostGroupState(int lookahead, int rnd, int num_random, int argc,
  current_batch_sz = 0;
  next_batch_no = 0;
  randomize = rnd;
- this->num_random = num_random;
  if (num_random > 0) {
- current_group.generate_random_ips();
+ current_group.generate_random_ips(num_random);
  }
 }
 
@@ -444,15 +443,12 @@ bool HostGroupState::get_next_host(struct sockaddr_storage *ss, size_t *sslen, s
  }
  }
  /* Check exclude list. */
- if (!hostInExclude((struct sockaddr *) ss, *sslen, exclude_group))
+ if (!hostInExclude((struct sockaddr *) ss, *sslen, exclude_group)) {
+ current_group.reject_last_host();
  break;
+ }
  } while (true);
 
- // If this is the last of the random IPs, pop in a new expression
- if (num_random > 0 && --num_random == 0) {
- process_next_expression();
- }
-
  return true;
 }
 

targets.h

@@ -100,11 +100,10 @@ class HostGroupState {
  /* Returns true iff the defer buffer is not yet full. */
  bool defer(Target *t);
  void undefer();
+ const char *next_expression();
  bool get_next_host(struct sockaddr_storage *ss, size_t *sslen, struct addrset *exclude_group);
  private:
- int num_random;
  bool process_next_expression();
- const char *next_expression();
 };
 
 /* ports is used to pass information about what ports to use for host discovery */