pirclient.cc 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  1. // Implementation of the main loop of the pirserver, responsible for the
  2. // communication with the tor process. All the actual private lookup
  3. // work is done by an appropriate subclass of PIRServer.
  4. #include <stdlib.h>
  5. #include <unistd.h>
  6. #include <arpa/inet.h>
  7. #include <string.h>
  8. #include "pirclient.h"
  9. #define PIRCLIENT_HDR_SIZE 13
  10. #define PIRCLIENT_REQUEST_CREATE 0x41
  11. #define PIRCLIENT_REQUEST_EXTRACT 0x42
  12. #define PIRCLIENT_RESPONSE_CREATE 0xBF
  13. #define PIRCLIENT_RESPONSE_EXTRACT 0xBE
  14. static int
  15. read_all(char *buf, size_t len)
  16. {
  17. int tot_read = 0;
  18. while(len > 0) {
  19. int res = read(0, buf, len);
  20. if (res <= 0) return res;
  21. buf += res;
  22. len -= res;
  23. tot_read += res;
  24. }
  25. return tot_read;
  26. }
  27. static int
  28. write_all(const char *buf, size_t len)
  29. {
  30. int tot_written = 0;
  31. while(len > 0) {
  32. int res = write(1, buf, len);
  33. if (res <= 0) return res;
  34. buf += res;
  35. len -= res;
  36. tot_written += res;
  37. }
  38. return tot_written;
  39. }
  40. void
  41. PIRClient::mainloop()
  42. {
  43. char header[PIRCLIENT_HDR_SIZE];
  44. size_t bodylen = 0;
  45. char *body = NULL;
  46. string query, response;
  47. size_t response_len;
  48. while(1) {
  49. // Read the request from stdin
  50. int res = read_all(header, PIRCLIENT_HDR_SIZE);
  51. if (res <= 0) return; // stdin has reached EOF (or error); we
  52. // will terminate
  53. bodylen = ntohl(*(uint32_t*)(header+PIRCLIENT_HDR_SIZE-4));
  54. if (bodylen > 0) {
  55. body = (char *)malloc(bodylen);
  56. res = read_all(body, bodylen);
  57. if (res <= 0) return;
  58. }
  59. // We have a complete request. Dispatch it.
  60. switch(header[8]) {
  61. case PIRCLIENT_REQUEST_CREATE:
  62. if (bodylen >= 32) {
  63. string plainquery(body, 32);
  64. string params(body+32, bodylen-32);
  65. void *queryid = NULL;
  66. string pirquery;
  67. create(plainquery, params, queryid, pirquery);
  68. size_t pirquery_len = pirquery.length();
  69. header[8] = PIRCLIENT_RESPONSE_CREATE;
  70. *(uint32_t*)(header+PIRCLIENT_HDR_SIZE-4) =
  71. htonl(8+pirquery_len);
  72. res = write_all(header, PIRCLIENT_HDR_SIZE);
  73. if (res <= 0) return;
  74. res = write_all((const char *)&queryid, sizeof(queryid));
  75. if (res <= 0) return;
  76. if (sizeof(queryid) < 8) {
  77. res = write_all("\0\0\0\0\0\0\0\0", 8-sizeof(queryid));
  78. if (res <= 0) return;
  79. }
  80. if (pirquery_len > 0) {
  81. res = write_all(pirquery.c_str(), pirquery_len);
  82. if (res <= 0) return;
  83. }
  84. }
  85. break;
  86. case PIRCLIENT_REQUEST_EXTRACT:
  87. if (bodylen >= 8) {
  88. void *queryid;
  89. memmove(&queryid, body, sizeof(queryid));
  90. string pirresponse(body+8, bodylen-8);
  91. string plainresponse;
  92. header[8] = PIRCLIENT_RESPONSE_EXTRACT;
  93. if (extract(queryid, pirresponse, plainresponse)) {
  94. response_len = plainresponse.length();
  95. } else {
  96. response_len = 0;
  97. }
  98. *(uint32_t*)(header+PIRCLIENT_HDR_SIZE-4) = htonl(response_len);
  99. res = write_all(header, PIRCLIENT_HDR_SIZE);
  100. if (res <= 0) return;
  101. if (response_len > 0) {
  102. res = write_all(plainresponse.c_str(), response_len);
  103. if (res <= 0) return;
  104. }
  105. }
  106. break;
  107. }
  108. // Clean up for the next request.
  109. free(body);
  110. body = NULL;
  111. bodylen = 0;
  112. }
  113. }