controller.cc 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785
  1. extern "C" {
  2. #include <event2/listener.h>
  3. #include <event2/bufferevent.h>
  4. #include <event2/buffer.h>
  5. }
  6. #include <NTL/vec_ZZ.h>
  7. #include <NTL/ZZ.h>
  8. #include <NTL/ZZ_p.h>
  9. #include <sys/socket.h>
  10. #include <netinet/in.h>
  11. #include <arpa/inet.h>
  12. #include <sys/time.h>
  13. #include <fstream>
  14. #include <vector>
  15. #include <set>
  16. #include <map>
  17. #include <stdlib.h>
  18. #include <string.h>
  19. #include "evutils.h"
  20. #include "subproblem.h"
  21. NTL_CLIENT
  22. #define DEBUG
  23. struct SubproblemProgress;
  24. typedef std::set<struct bufferevent *> BESet;
  25. typedef std::map<struct bufferevent *, SubproblemProgress *> BEMap;
  26. void besetdump(const BESet &bes, ostream &os)
  27. {
  28. BESet::const_iterator besit;
  29. os << hex << " ";
  30. for (besit = bes.begin(); besit != bes.end(); ++besit) {
  31. os << *besit << " ";
  32. }
  33. os << dec << "\n";
  34. }
  35. void besetfree(BESet &bes)
  36. {
  37. BESet::iterator besit;
  38. for (besit = bes.begin(); besit != bes.end(); ++besit) {
  39. bufferevent_free(*besit);
  40. }
  41. bes.clear();
  42. }
  43. void bemapdump(const BEMap &bem, ostream &os)
  44. {
  45. BEMap::const_iterator bemit;
  46. os << hex << " ";
  47. for (bemit = bem.begin(); bemit != bem.end(); ++bemit) {
  48. os << bemit->first << "->" << bemit->second << " ";
  49. }
  50. os << dec << "\n";
  51. }
  52. void bemapfree(BEMap &bem)
  53. {
  54. BEMap::iterator bemit;
  55. for (bemit = bem.begin(); bemit != bem.end(); ++bemit) {
  56. bufferevent_free(bemit->first);
  57. }
  58. bem.clear();
  59. }
  60. struct Statuses {
  61. BESet idle;
  62. BEMap working;
  63. // Dump the state for debug purposes
  64. void dump(ostream &os) const {
  65. os << " idle (" << idle.size() << "):\n";
  66. besetdump(idle, os);
  67. os << " working (" << working.size() << "):\n";
  68. bemapdump(working, os);
  69. }
  70. void free(void) {
  71. besetfree(idle);
  72. bemapfree(working);
  73. }
  74. };
  75. struct FactorDecomp {
  76. ZZ factor;
  77. vec_ZZ fvec;
  78. };
  79. void vsppdump(const vector<SubproblemProgress> &spv, ostream &os);
  80. static struct ControllerState {
  81. ZZ rho;
  82. FactorDecomp p, q;
  83. ZZ base, target;
  84. bool working;
  85. struct timeval started_working;
  86. vector<SubproblemProgress> subproblems_p, subproblems_q;
  87. Statuses dpnodes, workers;
  88. unsigned int num_unsolved_subproblems;
  89. unsigned int problems_remaining;
  90. struct evconnlistener *listener;
  91. ControllerState() : working(false), num_unsolved_subproblems(0),
  92. problems_remaining(0), listener(NULL) {}
  93. // Reset the state for a new problem with the same modulus
  94. void reset(void) {
  95. base = 0;
  96. target = 0;
  97. working = false;
  98. subproblems_p.clear();
  99. subproblems_q.clear();
  100. num_unsolved_subproblems = 0;
  101. started_working.tv_sec = 0;
  102. started_working.tv_usec = 0;
  103. }
  104. // Dump the state for debug purposes
  105. void dump(ostream &os) const {
  106. if (!working) {
  107. os << "Not working\n";
  108. return;
  109. }
  110. os << "P:\n";
  111. vsppdump(subproblems_p, os);
  112. os << "Q:\n";
  113. vsppdump(subproblems_q, os);
  114. os << "dpnodes:\n";
  115. dpnodes.dump(os);
  116. os << "workers:\n";
  117. workers.dump(os);
  118. }
  119. } ctrlstate;
  120. struct IPPort {
  121. unsigned char ipport[6];
  122. IPPort(unsigned char *ipp) {
  123. memmove(ipport, ipp, 6);
  124. }
  125. void dump(ostream &os) const {
  126. os << int(ipport[0]) << "." << int(ipport[1]) << "." <<
  127. int(ipport[2]) << "." << int(ipport[3]) << ":" <<
  128. ((ipport[4] << 8) + ipport[5]) << " ";
  129. }
  130. };
  131. typedef vector<IPPort> IPPortSet;
  132. void ipportsetdump(const IPPortSet &ipps, ostream &os)
  133. {
  134. IPPortSet::const_iterator ippsit;
  135. os << " ";
  136. for (ippsit = ipps.begin(); ippsit != ipps.end(); ++ippsit) {
  137. ippsit->dump(os);
  138. }
  139. os << "\n";
  140. }
  141. struct SubproblemProgress : Subproblem {
  142. // The sets of dpnodes and workers currently working on this subproblem
  143. BESet dpnodes, workers;
  144. // The dpnode IPPorts registered for this subproblem
  145. IPPortSet ipports;
  146. // The desired number of DPnodes for this subproblem
  147. unsigned short desired_dpnodes;
  148. // The maximum number of workers useful for this subproblem
  149. unsigned int max_workers;
  150. // Have we found a solution?
  151. bool solved;
  152. // The solution, if found.
  153. ZZ solution;
  154. SubproblemProgress(unsigned short id, const ZZ &b, const ZZ &t,
  155. const ZZ &m, const ZZ &o, unsigned int dpf) :
  156. Subproblem(id, b, t, m, o, dpf), solved(false) {
  157. // How many DPnodes should we use for a problem of this size?
  158. desired_dpnodes = 2;
  159. // How many workers would we like to use?
  160. ZZ sorder = SqrRoot(order >> 46);
  161. if (NumBits(sorder) > 30) {
  162. // Just use all the workers we can find
  163. max_workers = 4294967295U; // 2^32 - 1
  164. } else {
  165. max_workers = trunc_long(sorder,31) + 1;
  166. }
  167. }
  168. // Stop all dpnodes and workers
  169. void stop(void) {
  170. BESet::iterator iter;
  171. unsigned char stopcmd[1] = { 'S' };
  172. for (BESet::iterator iter = dpnodes.begin(); iter != dpnodes.end();
  173. ++iter) {
  174. bufferevent_write(*iter, stopcmd, 1);
  175. ctrlstate.dpnodes.working.erase(*iter);
  176. ctrlstate.dpnodes.idle.insert(*iter);
  177. }
  178. for (BESet::iterator iter = workers.begin(); iter != workers.end();
  179. ++iter) {
  180. bufferevent_write(*iter, stopcmd, 1);
  181. ctrlstate.workers.working.erase(*iter);
  182. ctrlstate.workers.idle.insert(*iter);
  183. }
  184. dpnodes.clear();
  185. workers.clear();
  186. ipports.clear();
  187. }
  188. // Dump for debugging purposes
  189. void dump(ostream &os) const {
  190. os << " Subproblem " << problemid << "\n";
  191. os << " dpnodes (" << dpnodes.size() << "):\n";
  192. besetdump(dpnodes, os);
  193. os << " workers (" << workers.size() << "):\n";
  194. besetdump(workers, os);
  195. os << " ipports (" << ipports.size() << "):\n";
  196. ipportsetdump(ipports, os);
  197. if (solved) {
  198. os << " solution: " << solution << "\n\n";
  199. }
  200. }
  201. void worker_write(struct bufferevent *bev) {
  202. bev_write(bev);
  203. unsigned short num_ipports = ipports.size();
  204. bufferevent_write(bev, &num_ipports, 2);
  205. for (unsigned short i = 0; i < num_ipports; ++i) {
  206. bufferevent_write(bev, ipports[i].ipport, 6);
  207. }
  208. #ifdef DEBUG
  209. cerr << "Added worker " << bev << " to subproblem "
  210. << problemid << "\n";
  211. #endif
  212. }
  213. };
  214. // Dump the state for debug purposes
  215. void vsppdump(const vector<SubproblemProgress> &spv, ostream &os)
  216. {
  217. vector<SubproblemProgress>::const_iterator spiter;
  218. int count = 0;
  219. for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
  220. ++count;
  221. os << " " << count << ":\n";
  222. spiter->dump(os);
  223. }
  224. os << "\n";
  225. }
  226. // Find a subproblem in the given vector that could use another DPnode,
  227. // and give it one of the idle ones. Only allocate it to a subproblem
  228. // with no current DPnodes if consider_empty is true.
  229. static void find_subproblem_for_dpnode(vector<SubproblemProgress> &spv,
  230. bool consider_empty)
  231. {
  232. vector<SubproblemProgress>::iterator spiter;
  233. for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
  234. if (spiter->solved) continue;
  235. if (spiter->dpnodes.size() == 0 && consider_empty == false) continue;
  236. // How many DPnodes would we like to have for this subproblem?
  237. while (spiter->dpnodes.size() < spiter->desired_dpnodes &&
  238. ctrlstate.dpnodes.idle.size() > 0) {
  239. // Get the first idle DPnode
  240. BESet::iterator beviter = ctrlstate.dpnodes.idle.begin();
  241. struct bufferevent *firstbev = *beviter;
  242. // Allocate it to the subproblem
  243. spiter->dpnodes.insert(firstbev);
  244. ctrlstate.dpnodes.working[firstbev] = &(*spiter);
  245. ctrlstate.dpnodes.idle.erase(firstbev);
  246. // Tell it to start listening for DPs
  247. spiter->bev_write(firstbev);
  248. }
  249. }
  250. }
  251. // Find a subproblem in the given vector that has all of its DPnodes and
  252. // could use another worker, and give it one of the idle ones.
  253. static void find_subproblem_for_worker(vector<SubproblemProgress> &spv)
  254. {
  255. vector<SubproblemProgress>::iterator spiter;
  256. for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
  257. if (spiter->solved) continue;
  258. while (spiter->ipports.size() == spiter->desired_dpnodes &&
  259. spiter->workers.size() < spiter->max_workers &&
  260. ctrlstate.workers.idle.size() > 0) {
  261. // Get the first idle worker
  262. BESet::iterator beviter = ctrlstate.workers.idle.begin();
  263. struct bufferevent *firstbev = *beviter;
  264. // Allocate it to the subproblem
  265. spiter->workers.insert(firstbev);
  266. ctrlstate.workers.working[firstbev] = &(*spiter);
  267. ctrlstate.workers.idle.erase(firstbev);
  268. // Tell it to start working on the subproblem
  269. spiter->worker_write(firstbev);
  270. }
  271. }
  272. }
  273. // See if there are any idle DPnodes or workers we can put to use
  274. static void schedule(void)
  275. {
  276. // Check the DPnodes
  277. // Iterate through the subproblems, looking for one that can use
  278. // another DPnode. First look for subproblems that already have
  279. // some, but not all, of their DPnodes
  280. if (ctrlstate.dpnodes.idle.size() > 0) {
  281. find_subproblem_for_dpnode(ctrlstate.subproblems_p, false);
  282. }
  283. if (ctrlstate.dpnodes.idle.size() > 0) {
  284. find_subproblem_for_dpnode(ctrlstate.subproblems_q, false);
  285. }
  286. // If there are still more dpnodes to place, start assigning them to
  287. // subproblems with no current dpnodes
  288. if (ctrlstate.dpnodes.idle.size() > 0) {
  289. find_subproblem_for_dpnode(ctrlstate.subproblems_p, true);
  290. }
  291. if (ctrlstate.dpnodes.idle.size() > 0) {
  292. find_subproblem_for_dpnode(ctrlstate.subproblems_q, true);
  293. }
  294. // Check the workers
  295. // Iterate through the subproblems, looking for one that can use
  296. // another worker.
  297. if (ctrlstate.workers.idle.size() > 0) {
  298. find_subproblem_for_worker(ctrlstate.subproblems_p);
  299. }
  300. if (ctrlstate.workers.idle.size() > 0) {
  301. find_subproblem_for_worker(ctrlstate.subproblems_q);
  302. }
  303. // cerr << "After schedule:\n"; ctrlstate.dump(cerr);
  304. }
  305. static ZZ computation_complete_p(const vector<SubproblemProgress> &v)
  306. {
  307. ZZ curmodulus, curexp;
  308. curmodulus = 2;
  309. curexp = 0;
  310. vector<SubproblemProgress>::const_iterator vit;
  311. for(vit = v.begin(); vit != v.end(); ++vit) {
  312. CRT(curexp, curmodulus, vit->solution, vit->order);
  313. }
  314. if (curexp < 0) {
  315. curexp += curmodulus;
  316. }
  317. return curexp;
  318. }
  319. // All subproblems are solved. Combine the results.
  320. static void computation_complete(void)
  321. {
  322. ZZ exp_p = computation_complete_p(ctrlstate.subproblems_p);
  323. ZZ exp_q = computation_complete_p(ctrlstate.subproblems_q);
  324. ZZ pm1 = (ctrlstate.p.factor - 1)/2;
  325. ZZ qm1 = (ctrlstate.q.factor - 1)/2;
  326. CRT(exp_p, pm1, exp_q, qm1);
  327. if (exp_p < 0) {
  328. exp_p += pm1;
  329. }
  330. ZZ& expon = exp_p;
  331. struct timeval ended_working;
  332. gettimeofday(&ended_working, NULL);
  333. unsigned long long computation_length_ms =
  334. (ended_working.tv_sec - ctrlstate.started_working.tv_sec) * 1000 +
  335. (ended_working.tv_usec - ctrlstate.started_working.tv_usec) / 1000;
  336. char length_buf[50];
  337. sprintf(length_buf, "%lld.%03lld s", computation_length_ms / 1000,
  338. computation_length_ms % 1000);
  339. cout << "expon = " << expon << "\n";
  340. ZZ base_exp = PowerMod(ctrlstate.base, expon, ctrlstate.rho);
  341. if (base_exp == ctrlstate.target) {
  342. cout << "CORRECT in " << length_buf << "\n";
  343. } else {
  344. cout << "INCORRECT in " << length_buf << ":\n";
  345. cout << "base^exp = " << base_exp << "\n";
  346. cout << "target = " << ctrlstate.target << "\n";
  347. }
  348. ctrlstate.reset();
  349. }
  350. static unsigned short curproblemid = 0;
  351. // Take base and target mod f.factor, then decompose that into small
  352. // subproblems given our knowledge of the factors of phi(f.factor)
  353. static vector<SubproblemProgress> decomp(const ZZ_p &base, const ZZ_p &target,
  354. const FactorDecomp &f)
  355. {
  356. vector<SubproblemProgress> ret;
  357. // Compute phi(factor)
  358. const int fveclen = f.fvec.length();
  359. ZZ phi = to_ZZ(2);
  360. for (int i = 0; i < fveclen; ++i) {
  361. phi *= f.fvec[i];
  362. }
  363. ZZ_p::init(f.factor);
  364. for (int i = 0; i < fveclen; ++i) {
  365. const ZZ& order = f.fvec[i];
  366. ZZ quotient = phi / order;
  367. ZZ_p subgroup_base = to_ZZ_p(rep(base));
  368. subgroup_base = power(subgroup_base, quotient);
  369. ZZ_p subgroup_target = to_ZZ_p(rep(target));
  370. subgroup_target = power(subgroup_target, quotient);
  371. if (subgroup_base == 1) {
  372. // The original base wasn't a generator of the whole group
  373. if (subgroup_target == 1) {
  374. // But the target is in the subgroup. Lucky us.
  375. continue;
  376. } else {
  377. ret.clear();
  378. return ret;
  379. }
  380. }
  381. // By default, 1 in 1000 points are distinguihed points. The
  382. // number in the next line is 2^32/1000
  383. unsigned int dpfreq = 4294967;
  384. if (order < 1000) {
  385. // Just make every point a DP
  386. dpfreq = 4294967295U;
  387. } else if (NumBits(order) < 27) {
  388. // The frequency of DPs should be 10/sqrt(order) to avoid
  389. // a DP-free cycle, so dpfreq = (10*2^32)/sqrt(order)
  390. ZZ f = (to_ZZ(10) << 32) / SqrRoot(order);
  391. dpfreq = trunc_long(f, 31);
  392. }
  393. ++curproblemid;
  394. ret.push_back(SubproblemProgress(curproblemid, rep(subgroup_base),
  395. rep(subgroup_target),
  396. f.factor, order, dpfreq));
  397. }
  398. return ret;
  399. }
  400. static int generate_problem(struct event_base *evbase)
  401. {
  402. // If there are no more problems to generate, close the listener
  403. if (ctrlstate.problems_remaining == 0) {
  404. evconnlistener_free(ctrlstate.listener);
  405. ctrlstate.dpnodes.free();
  406. ctrlstate.workers.free();
  407. return -1;
  408. }
  409. // If there's already a problem on the go, don't generate another one
  410. if (ctrlstate.working == true) {
  411. return -1;
  412. }
  413. ctrlstate.working = true;
  414. int num_subproblems_p = 0;
  415. int num_subproblems_q = 0;
  416. // Generate a DLP mod rho (in the large odd-order subgroup)
  417. ZZ_p::init(ctrlstate.rho);
  418. do {
  419. ZZ_p base = power(random_ZZ_p(), 2);
  420. ZZ_p target = power(random_ZZ_p(), 2);
  421. gettimeofday(&ctrlstate.started_working, NULL);
  422. ctrlstate.base = rep(base);
  423. ctrlstate.target = rep(target);
  424. // Decompose it mod p and mod q
  425. ctrlstate.subproblems_p = decomp(base, target, ctrlstate.p);
  426. ctrlstate.subproblems_q = decomp(base, target, ctrlstate.q);
  427. num_subproblems_p = ctrlstate.subproblems_p.size();
  428. num_subproblems_q = ctrlstate.subproblems_q.size();
  429. } while (num_subproblems_p == 0 || num_subproblems_q == 0);
  430. ctrlstate.num_unsolved_subproblems =
  431. num_subproblems_p + num_subproblems_q;
  432. schedule();
  433. --ctrlstate.problems_remaining;
  434. return 0;
  435. }
  436. typedef enum {
  437. CCSTATE_START,
  438. CCSTATE_DPWAITRESP,
  439. CCSTATE_DPLISTENING,
  440. CCSTATE_DPEXPON,
  441. CCSTATE_END
  442. } CCState;
  443. struct ControllerConnInfo {
  444. CCState state;
  445. ControllerConnInfo() : state(CCSTATE_DPWAITRESP) {}
  446. };
  447. static void controller_dpnode_reader(struct bufferevent *bev, void *ctx)
  448. {
  449. struct evbuffer *input = bufferevent_get_input(bev);
  450. ControllerConnInfo *info = (ControllerConnInfo *)ctx;
  451. unsigned char cmd[1];
  452. ZZ expon;
  453. while(1) {
  454. size_t len = evbuffer_get_length(input);
  455. switch (info->state) {
  456. case CCSTATE_START:
  457. case CCSTATE_DPWAITRESP:
  458. if (len < 1) return;
  459. bufferevent_read(bev, cmd, 1);
  460. switch (cmd[0]) {
  461. case 'L':
  462. info->state = CCSTATE_DPLISTENING;
  463. break;
  464. case 'E':
  465. info->state = CCSTATE_DPEXPON;
  466. break;
  467. default:
  468. /* Unknown DPnode command received */
  469. fprintf(stderr, "Unknown command in "
  470. "controller_dpnode_reader: "
  471. "%c\n", cmd[0]);
  472. info->state = CCSTATE_END;
  473. break;
  474. }
  475. break;
  476. case CCSTATE_DPLISTENING:
  477. // Read 6 bytes
  478. if (len < 6) return;
  479. unsigned char ipport[6];
  480. unsigned int DPip;
  481. unsigned short DPport;
  482. bufferevent_read(bev, ipport, 6);
  483. memmove(&DPip, ipport, 4);
  484. memmove(&DPport, ipport+4, 2);
  485. {
  486. struct in_addr DPaddr = { DPip };
  487. #ifdef DEBUG
  488. fprintf(stderr, "DP node at %s:%d\n", inet_ntoa(DPaddr), ntohs(DPport));
  489. #endif
  490. if (ctrlstate.dpnodes.working.count(bev) > 0) {
  491. ctrlstate.dpnodes.working[bev]->ipports.push_back(
  492. IPPort(ipport));
  493. schedule();
  494. }
  495. }
  496. info->state = CCSTATE_DPWAITRESP;
  497. break;
  498. case CCSTATE_DPEXPON:
  499. // Read the subproblemid and the answer to the subproblem
  500. if (len < 2 + 3*sizeof(unsigned int)) return;
  501. unsigned char exponbytes[2 + 3*sizeof(unsigned int)];
  502. unsigned short problemid;
  503. bufferevent_read(bev, exponbytes, 2 + 3*sizeof(unsigned int));
  504. memmove(&problemid, exponbytes, 2);
  505. ZZFromBytes(expon, exponbytes+2, 3*sizeof(unsigned int));
  506. // Find the subproblem and check the answer
  507. if (ctrlstate.dpnodes.working.count(bev) > 0) {
  508. SubproblemProgress *spp = ctrlstate.dpnodes.working[bev];
  509. if (spp->problemid == problemid &&
  510. spp->solved == false &&
  511. spp->target ==
  512. PowerMod(spp->base, expon, spp->modulus)) {
  513. // Subproblem solved!
  514. spp->solution = expon;
  515. spp->solved = true;
  516. spp->stop();
  517. ctrlstate.num_unsolved_subproblems--;
  518. if (ctrlstate.num_unsolved_subproblems == 0) {
  519. computation_complete();
  520. generate_problem(bufferevent_get_base(bev));
  521. }
  522. schedule();
  523. }
  524. }
  525. info->state = CCSTATE_DPWAITRESP;
  526. break;
  527. case CCSTATE_END:
  528. // Shut down the connection
  529. delete info;
  530. bufferevent_free(bev);
  531. return;
  532. }
  533. }
  534. }
  535. static void controller_dpnode_event_cb(struct bufferevent *bev, short events,
  536. void *ctx)
  537. {
  538. if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
  539. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  540. fprintf(stderr, "Closing dpnode connection\n");
  541. if (ctrlstate.dpnodes.working.count(bev)) {
  542. // If we lose a dpnode from an active computation, the
  543. // computation is useless.
  544. SubproblemProgress *spp = ctrlstate.dpnodes.working[bev];
  545. ctrlstate.dpnodes.working.erase(bev);
  546. spp->dpnodes.erase(bev);
  547. spp->stop();
  548. } else {
  549. ctrlstate.dpnodes.idle.erase(bev);
  550. }
  551. delete info;
  552. bufferevent_free(bev);
  553. schedule();
  554. }
  555. }
  556. static void controller_worker_event_cb(struct bufferevent *bev, short events,
  557. void *ctx)
  558. {
  559. if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
  560. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  561. fprintf(stderr, "Closing worker connection\n");
  562. if (ctrlstate.workers.working.count(bev)) {
  563. SubproblemProgress *spp = ctrlstate.workers.working[bev];
  564. ctrlstate.workers.working.erase(bev);
  565. spp->workers.erase(bev);
  566. } else {
  567. ctrlstate.workers.idle.erase(bev);
  568. }
  569. delete info;
  570. bufferevent_free(bev);
  571. schedule();
  572. }
  573. }
  574. static void controller_event_cb(struct bufferevent *bev, short events,
  575. void *ctx)
  576. {
  577. if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
  578. fprintf(stderr, "Closing connection\n");
  579. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  580. delete info;
  581. bufferevent_free(bev);
  582. }
  583. }
  584. // We're just going to read a single byte that will tell us whether the
  585. // peer is a DPnode or a Worker
  586. static void controller_master_reader(struct bufferevent *bev, void *ctx)
  587. {
  588. struct evbuffer *input = bufferevent_get_input(bev);
  589. size_t len = evbuffer_get_length(input);
  590. if (len < 1) return;
  591. char indata[1];
  592. bufferevent_read(bev, indata, 1);
  593. switch(indata[0]) {
  594. case 'D':
  595. /* Add this DPnode to the list of available ones */
  596. ctrlstate.dpnodes.idle.insert(bev);
  597. bufferevent_setcb(bev, controller_dpnode_reader, NULL,
  598. controller_dpnode_event_cb, ctx);
  599. controller_dpnode_reader(bev, ctx);
  600. schedule();
  601. return;
  602. case 'W':
  603. ctrlstate.workers.idle.insert(bev);
  604. // We don't actually read anything from workers
  605. bufferevent_enable(bev, EV_WRITE);
  606. bufferevent_setcb(bev, NULL, NULL,
  607. controller_worker_event_cb, ctx);
  608. schedule();
  609. return;
  610. default:
  611. fprintf(stderr, "Unknown command in controller_master_reader: "
  612. "%c\n", indata[0]);
  613. ControllerConnInfo *info = (ControllerConnInfo*)ctx;
  614. delete info;
  615. bufferevent_free(bev);
  616. return;
  617. }
  618. }
  619. static void controller_accept_cb(struct evconnlistener *listener,
  620. evutil_socket_t fd, struct sockaddr *address, int socklen,
  621. void *ctx)
  622. {
  623. // Create the state of the new connection
  624. ControllerConnInfo *info = new ControllerConnInfo();
  625. // Create a bufferevent for the new connection
  626. struct event_base *base = evconnlistener_get_base(listener);
  627. struct bufferevent *bev = bufferevent_socket_new(
  628. base, fd, BEV_OPT_CLOSE_ON_FREE);
  629. bufferevent_setcb(bev, controller_master_reader, NULL,
  630. controller_event_cb, info);
  631. bufferevent_enable(bev, EV_READ|EV_WRITE);
  632. }
  633. // Create a new controller socket. bindport is the port to bind to (in
  634. // host byte order), or 0 if any port will do. ip and boundport are set
  635. // to the IP and port of the socket, in network byte order.
  636. struct evconnlistener *controller_create(struct event_base *evbase,
  637. unsigned short bindport, unsigned int *ip, unsigned short *boundport)
  638. {
  639. return listener_create(evbase, bindport, controller_accept_cb, NULL,
  640. ip, boundport, false);
  641. }
  642. int main(int argc, char **argv)
  643. {
  644. // Initialize the prng with some randomness from the kernel
  645. unsigned char randbuf[1024];
  646. ifstream urand("/dev/urandom");
  647. urand.read((char *)randbuf, sizeof(randbuf));
  648. urand.close();
  649. ZZ randzz = ZZFromBytes(randbuf, sizeof(randbuf));
  650. SetSeed(randzz);
  651. // Read the modulus and the factorization of its totient from cin
  652. cin >> ctrlstate.rho >> ctrlstate.p.factor >> ctrlstate.p.fvec >>
  653. ctrlstate.q.factor >> ctrlstate.q.fvec;
  654. unsigned short bindport = 0;
  655. if (argc > 1) {
  656. bindport = strtoul(argv[1], NULL, 10);
  657. }
  658. struct event_base *evbase = event_base_new();
  659. unsigned int myip;
  660. unsigned short myport;
  661. ctrlstate.listener = controller_create(evbase, bindport, &myip, &myport);
  662. struct in_addr myaddr = { myip };
  663. printf("Bound to %s:%d\n", inet_ntoa(myaddr), ntohs(myport));
  664. ctrlstate.problems_remaining = 2;
  665. // Kick off the first problem to solve
  666. generate_problem(evbase);
  667. event_base_dispatch(evbase);
  668. return 0;
  669. }