Make NetBlockRandomIPv4 responsible for tracking number of random addrs

TargetGroup.cc

@@ -111,6 +111,7 @@ public:
    * 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 @@ public:
   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,14 +443,11 @@ 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();
     }
+  } while (true);
 
   return true;
 }

targets.h

@@ -100,11 +100,10 @@ public:
   /* 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 */