pirclient.cc 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  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. if (extract(queryid, pirresponse, plainresponse)) {
  93. response_len = plainresponse.length();
  94. } else {
  95. response_len = 0;
  96. }
  97. *(uint32_t*)(header+PIRCLIENT_HDR_SIZE-4) = htonl(response_len);
  98. res = write_all(header, PIRCLIENT_HDR_SIZE);
  99. if (res <= 0) return;
  100. if (response_len > 0) {
  101. res = write_all(plainresponse.c_str(), response_len);
  102. if (res <= 0) return;
  103. }
  104. }
  105. break;
  106. }
  107. // Clean up for the next request.
  108. free(body);
  109. body = NULL;
  110. bodylen = 0;
  111. }
  112. }