1 /*
2 * ZMap Copyright 2013 Regents of the University of Michigan
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
5 * use this file except in compliance with the License. You may obtain a copy
6 * of the License at http://www.apache.org/licenses/LICENSE-2.0
7 */
8
9 #include <stdint.h>
10 #include <assert.h>
11
12 #include <gmp.h>
13
14 #include "../lib/includes.h"
15 #include "../lib/blacklist.h"
16 #include "shard.h"
17 #include "state.h"
18
shard_init(shard_t * shard,uint8_t shard_id,uint8_t num_shards,uint8_t sub_id,uint8_t num_subshards,const cycle_t * cycle,shard_complete_cb cb,void * arg)19 void shard_init(shard_t* shard,
20 uint8_t shard_id,
21 uint8_t num_shards,
22 uint8_t sub_id,
23 uint8_t num_subshards,
24 const cycle_t* cycle,
25 shard_complete_cb cb,
26 void *arg)
27 {
28 // Start out by figuring out the multiplication factor for this shard.
29 // With one shard, this would just be the generator, but with n shards,
30 // f = g^n.
31
32 // Then on top of that, we want to shard internally (subshards) per
33 // thread. With t threads, f = g^(nr).
34 //
35 // tot_shards = nr
36 uint32_t tot_shards = (uint32_t) num_shards * (uint32_t) num_subshards;
37 uint64_t num_elts = cycle->group->prime - 1;
38 mpz_t start, generator, prime, result, power;
39 mpz_init_set_ui(start, cycle->offset);
40 mpz_init_set_ui(generator, cycle->generator);
41 mpz_init_set_ui(power, tot_shards);
42 mpz_init_set_ui(prime, cycle->group->prime);
43 mpz_init(result);
44 mpz_powm(result, generator, power, prime);
45 shard->params.factor = (uint64_t) mpz_get_ui(result);
46 shard->params.modulus = cycle->group->prime;
47
48 // e = p - 1 = num_elts
49 // begin_idx = s + tr
50 // end_idx = [e - (e % nr) + (s + tr)] % e
51 // = [e - (e % nr) + begin_idx] % e
52 uint64_t begin_idx = shard_id + sub_id*num_shards;
53 uint64_t end_idx = (num_elts - (num_elts % tot_shards) + begin_idx) % num_elts;
54 if (end_idx >= tot_shards) {
55 end_idx += tot_shards;
56 end_idx %= num_elts;
57 }
58 mpz_powm_ui(result, generator, begin_idx + 1, prime);
59 shard->params.first = (uint64_t) mpz_get_ui(result);
60 shard->params.first *= cycle->offset;
61 shard->params.first %= shard->params.modulus;
62 mpz_powm_ui(result, generator, end_idx + 1, prime);
63 shard->params.last = (uint64_t) mpz_get_ui(result);
64 shard->params.last *= cycle->offset;
65 shard->params.last %= shard->params.modulus;
66 shard->current = shard->params.first;
67 // Handle scanning a sample
68 if (zsend.targets != zsend.max_index) {
69 shard->state.max_targets = zsend.targets / num_subshards;
70 uint32_t leftover = zsend.targets % num_subshards;
71 if (leftover > sub_id) {
72 shard->state.max_targets++;
73 }
74 } else {
75 shard->state.max_targets = zsend.targets;
76 }
77
78
79 // Set the (thread) id
80 shard->id = sub_id;
81
82 // Set the callbacks
83 shard->cb = cb;
84 shard->arg = arg;
85
86 if (shard->current - 1 >= zsend.max_index) {
87 shard_get_next_ip(shard);
88 }
89
90 // Clear everything
91 mpz_clear(start);
92 mpz_clear(generator);
93 mpz_clear(prime);
94 mpz_clear(power);
95 mpz_clear(result);
96 }
97
shard_get_cur_ip(shard_t * shard)98 uint32_t shard_get_cur_ip(shard_t *shard)
99 {
100 return (uint32_t) blacklist_lookup_index(shard->current - 1);
101 }
102
shard_get_next_elem(shard_t * shard)103 static inline uint32_t shard_get_next_elem(shard_t *shard)
104 {
105 do {
106 shard->current *= shard->params.factor;
107 shard->current %= shard->params.modulus;
108 } while (shard->current >= (1LL << 32));
109 return (uint32_t) shard->current;
110 }
111
shard_get_next_ip(shard_t * shard)112 uint32_t shard_get_next_ip(shard_t *shard)
113 {
114 while (1) {
115 uint32_t candidate = shard_get_next_elem(shard);
116 if (candidate == shard->params.last) {
117 return 0;
118 }
119 if (candidate - 1 < zsend.max_index) {
120 shard->state.whitelisted++;
121 return blacklist_lookup_index(candidate - 1);
122 }
123 shard->state.blacklisted++;
124 }
125 }
126