win32_resolv.c

AnyQ服务端源代码(2004/10/28)源码

  if (dstBuf[offset - 1] != 0) { 
    // if the last character was NOT '.', normal update is OK
    dstBuf[lastCountOffset] = offset - lastCountOffset -1;
    // And woo! We've already GOT a terminating NULL character from the
    // string copy
    offset++;
  } else {
    // last character was a '.'. Therefore, we should use THAT as the
    // terminating null, and not the extra one which is now present at the 
    // end of the string. Therefore, just don't bother incrementing the count!
  }

  // Add in the type and class
  *((unsigned short*) (dstBuf + offset)) = htons(qType);
  offset+=2;
  *((unsigned short*) (dstBuf + offset)) = htons(qClass);
  offset+=2;

  // Finally, return the length
  return offset;
}



/**
 * Perform a DNS lookup
 *
 * @param name Name of service to look up
 * @param domain Domain To look it up in
 * @param qClass Class of query
 * @param qType Type of query
 * @param dnsServers DNS servers to try
 * @param dstBuffer Where to put received DNS packet
 * @param dstBufferSize Size of dstBuffer
 * @return Size of received packet, or <0 on failure
 */
int dnsLookup(const char* name, const char* domain, 
              u_int16_t qClass, u_int16_t qType,
              struct in_addr* dnsServers, 
              char* dstBuffer, int dstBufferSize) {
  char tmpBuf[BUF_SIZE];
  int pktSize;
  int dnsServerCount;
  int retries;
  int count;
  int expectedId;


  // start off with expectedId 0
  expectedId = 0;
  
  // OK, format the packet & check
  if ((pktSize = formatPacket(tmpBuf, BUF_SIZE,
                              name, domain,
                              qClass, qType)) < 0) {
    return -1;
  }
  
  // retry the lookup
  for(retries=0; retries < 5; retries++) {

    // Now, try each DNS server
    dnsServerCount = 0;
    while(dnsServers[dnsServerCount].s_addr) {
      // try UDP DNS if the packet ain't too big
      if (pktSize <= NS_PACKETSZ) {
        count = dnsUdp(dnsServers[dnsServerCount], 
                       tmpBuf, pktSize,
                       dstBuffer, dstBufferSize,
                       &expectedId);
        switch(count) {
        case RESP_TRYTCP:
        case RESP_TRYNEXTSERVER:
          break;

        case RESP_UNKNOWNHOST:
          return -1;

        default:
          return count;
        }
      } else {
        count = RESP_TRYTCP;
      }

      // try TCP connection?
      if (count == RESP_TRYTCP) {
        count = dnsTcp(dnsServers[dnsServerCount], 
                       tmpBuf, pktSize,
                       dstBuffer, dstBufferSize,
                       &expectedId);
        switch(count) {
        case RESP_TRYTCP:
        case RESP_TRYNEXTSERVER:
          break;
          
        case RESP_UNKNOWNHOST:
          return -1;
          
        default:
          return count;
        }
      }

      // Move on to next server
      dnsServerCount++;
    }
  }

  // if we get here, we have not found a valid address
  return -1;
}


/**
 * Perform a DNS lookup using UDP
 *
 * @param dnsServer DNS Server to look up
 * @param sendPktBuf Buffer containing the packet to send
 * @param sendPktSize Size of the packet to send
 * @param recvPktBuf Buffer to put received packet in
 * @param recvPktSize Size of the recvPktBuf buffer
 * @param expectedId Id counter to identify response packets
 * @return Received buffer size, or <0 on failure
 */
int dnsUdp(struct in_addr dnsServer, 
           char* sendPktBuf, int sendPktSize,
           char* recvPktBuf, int recvPktSize,
           int* expectedId) {
  int sockFd;
  struct sockaddr_in server;
  struct sockaddr_in local;
  struct sockaddr_in from;
  int count;
  int tmp;
  struct pollfd sockPollFd;
  HEADER* header;


  // setup sockaddr for local machine
  memset(&local, 0, sizeof(struct sockaddr_in));
  local.sin_family = AF_INET;
  local.sin_port = 0;
  local.sin_addr.s_addr = INADDR_ANY;

  // setup sockaddr for remote server
  memset(&server, 0, sizeof(struct sockaddr_in));
  server.sin_family = AF_INET;
  server.sin_port = htons(DNS_PORT);
  server.sin_addr.s_addr = dnsServer.s_addr;

  // Create UDP socket 
  if ((sockFd = socket(AF_INET, SOCK_DGRAM, UDP_PROTO)) < 0) {
    return RESP_TRYNEXTSERVER;
  }

  // bind input socket to local machine
  if (bind(sockFd, 
           (struct sockaddr*) &local, sizeof(struct sockaddr_in)) < 0) {
    close(sockFd);
    return RESP_TRYNEXTSERVER;
  }

  // set the expectedId in the packet
  ((HEADER*) sendPktBuf)->id = htons(*expectedId);
  (*expectedId)++;

  // send the packet to the DNS server
  if (sendto(sockFd, 
             sendPktBuf, sendPktSize,
             0,
             (struct sockaddr*) &server,
             sizeof(struct sockaddr_in)) < 0) {
    close(sockFd);
    return RESP_TRYNEXTSERVER;
  }

  // setup sockaddr for received packet
  memset(&from, 0, sizeof(from));
  from.sin_family = AF_INET;
  tmp = sizeof(from);
  memset(recvPktBuf, 0, recvPktSize);


  // This is where we listen for connections
  while(1) {

    // Wait for a response
    sockPollFd.fd = sockFd;
    sockPollFd.events = POLLIN | POLLERR;
    sockPollFd.revents = 0;
    if (poll(&sockPollFd, 1, 7000) <= 0) {
      close(sockFd);
      return RESP_TRYNEXTSERVER;
    }
    
    // OK, check we have got some data
    if (sockPollFd.revents != POLLIN) {
      close(sockFd);
      return RESP_TRYNEXTSERVER;
    }    
    
    // read it
    if ((count = recvfrom(sockFd,
                          recvPktBuf,
                          recvPktSize,
                          0,
                          (struct sockaddr*) &from,
                          &tmp)) < 0) {
      close(sockFd);
      return RESP_TRYNEXTSERVER;
    }
    
    // check packet is big enough
    header = (HEADER*) recvPktBuf;
    if (count < sizeof(HEADER)) {
      // packet too small. try next server
      close(sockFd);
      return RESP_TRYNEXTSERVER;
    }
    
    // incorrect ID => old response. start listening again
    if (ntohs(header->id) != ((*expectedId) - 1)) {
      continue;
    }

    // packet truncated. try tcp
    if (header->tc) {
      close(sockFd);
      return RESP_TRYTCP;
    }

    // not a response. try next server
    if (!header->qr) {
      close(sockFd);
      return RESP_TRYNEXTSERVER;
    }

    // unknown host. exit lookup
    if (header->rcode == RCODE_NAMEERROR) {
      close(sockFd);
      return RESP_UNKNOWNHOST;
    }
  
    // error. try next server
    if (header->rcode != RCODE_OK) {
      close(sockFd);
      return RESP_TRYNEXTSERVER;
    }
  
    // There were no actual records! try next server
    if (header->ancount == 0) {
      close(sockFd);
      return RESP_TRYNEXTSERVER;
    }   

    // If we get here, we are OK! exit the loop
    break;
  }

  // OK! return the count
  close(sockFd);
  return count;
}



/**
 * Perform a DNS lookup using TCP
 *
 * @param dnsServer DNS Server to look up
 * @param sendPktBuf Buffer containing the packet to send
 * @param sendPktSize Size of the packet to send
 * @param recvPktBuf Buffer to put received packet in
 * @param recvPktSize Size of the recvPktBuf buffer
 * @param expectedId Id counter to identify response packets
 * @return Received buffer size, or <0 on failure
 */
int dnsTcp(struct in_addr dnsServer, 
           char* sendPktBuf, int sendPktSize,
           char* recvPktBuf, int recvPktSize,
           int* expectedId) {
  int sockFd;
  struct sockaddr_in server;
  struct sockaddr_in local;
  int count;
  struct pollfd sockPollFd;
  char tmpBuf[10];
  HEADER* header;



  // setup sockaddr for local machine
  memset(&local, 0, sizeof(struct sockaddr_in));
  local.sin_family = AF_INET;
  local.sin_port = 0;
  local.sin_addr.s_addr = INADDR_ANY;

  // setup sockaddr for remote server
  memset(&server, 0, sizeof(struct sockaddr_in));
  server.sin_family = AF_INET;
  server.sin_port = htons(DNS_PORT);
  server.sin_addr.s_addr = dnsServer.s_addr;
  
  // Create TCP socket 
  if ((sockFd = socket(AF_INET, SOCK_STREAM, TCP_PROTO)) < 0) {
    return RESP_TRYNEXTSERVER;
  }

  // bind input socket to local machine
  if (bind(sockFd, 
           (struct sockaddr*) &local, sizeof(struct sockaddr_in)) < 0) {
    close(sockFd);
    return RESP_TRYNEXTSERVER;
  }

  // connect to the remote server
  if (connect(sockFd,
              (struct sockaddr*) &server, sizeof(struct sockaddr_in)) < 0) {
    close(sockFd);
    return RESP_TRYNEXTSERVER;
  }
  
  // set the expectedId in the packet
  ((HEADER*) sendPktBuf)->id = htons(*expectedId);
  (*expectedId)++;

  // send the size to the DNS server
  *((unsigned short*) tmpBuf) = htons(sendPktSize);
  if (send(sockFd, 
           tmpBuf, 2,
           0) < 0) {
    close(sockFd);
    return RESP_TRYNEXTSERVER;
  } 

  // send the packet to the DNS server
  if (send(sockFd, 
           sendPktBuf, sendPktSize,
           0) < 0) {
    close(sockFd);
    return RESP_TRYNEXTSERVER;
  }

  // This is where we listen for connections
  while(1) {

    // Wait for a response
    sockPollFd.fd = sockFd;
    sockPollFd.events = POLLIN | POLLERR;
    sockPollFd.revents = 0;