network.cpp

从Beevers的个人网站上下载 的源码

/*
  ktracker (c) 2006 Kris Beevers

  This file is part of ktracker.

  ktracker is free software; you can redistribute it and/or modify it
  under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  ktracker is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with ktracker; if not, write to the Free Software Foundation,
  Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA

  $Id: network.cpp,v 1.1.1.1 2006/10/10 20:41:29 beevek Exp $
*/

#include <stdio.h>
#include <errno.h>
#include <netinet/in.h>
#include <netdb.h>
#include <signal.h>
#include <fcntl.h>
#include <string.h>
#include <sys/select.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <inttypes.h>
#include <vector>
#include <algorithm>
#include "network.hpp"

// a bunch of stuff is stolen and modified from robot / netrobotd.cpp

int tcp_sock = -1;
std::vector<int> clients;
packet_handler_t packet_handler = 0;

inline void setup_sockaddr(struct sockaddr_in *sa)
{
  memset(sa, 0, sizeof(*sa));
  sa->sin_family = AF_INET;
  sa->sin_addr.s_addr = htonl(INADDR_ANY);
  sa->sin_port = htons(KTRACKER_PORT);
}

int listen_init(void)
{
  tcp_sock = socket(AF_INET, SOCK_STREAM, 0);
  if(tcp_sock < 0)
    return -1;
  int i = 1;
  if(setsockopt(tcp_sock, SOL_SOCKET, SO_REUSEADDR,
		(char *)&i, sizeof(&i)) < 0) {
    close(tcp_sock);
    return -1;
  }
  struct sockaddr_in sa;
  setup_sockaddr(&sa);
  if(bind(tcp_sock, (struct sockaddr *)&sa, sizeof(sa)) < 0 ||
     listen(tcp_sock, 5) < 0) {
    close(tcp_sock);
    return -1;
  }
  return 0;
}

const char * name_from_sock(int sock)
{
  struct sockaddr_in peer;
  int peerlen = sizeof(peer);
  getpeername(sock, (struct sockaddr *)&peer, (socklen_t *)&peerlen);
  return inet_ntoa(peer.sin_addr);
}

void disconnect_tcp_client(int sock)
{
  fprintf(stderr, "%s disconnected\n", name_from_sock(sock));
  std::vector<int>::iterator i = find(clients.begin(), clients.end(), sock);
  if(i != clients.end())
    clients.erase(i);
}

int accept_client()
{
  struct sockaddr_in ca;
  uint32_t calen = sizeof(ca);
  int client = accept(tcp_sock, (struct sockaddr *)&ca, &calen);
  if(client < 0)
    return -1;
  fcntl(client, F_SETOWN, getpid());
  fcntl(client, F_SETSIG, SIGIO); // SIGIO will be triggered when data is available
  fcntl(client, F_SETFL, O_ASYNC);
  clients.push_back(client);
  fprintf(stderr, "%s connected\n", name_from_sock(client));
  return 0;
}

void handle_network_request(int sig)
{
  int max = tcp_sock;
  for(uint32_t i = 0; i < clients.size(); ++i)
    if(clients[i] > max) max = clients[i];

  struct timeval tv;
  tv.tv_sec = tv.tv_usec = 0;
  int n;
  fd_set rfds;
  do { // loop until there's no waiting data
    FD_ZERO(&rfds);
    FD_SET(tcp_sock, &rfds);
    for(uint32_t i = 0; i < clients.size(); ++i)
      FD_SET(clients[i], &rfds);
    n = select(max + 1, &rfds, 0, 0, &tv);
    if(n < 0) {
      if(errno == EINTR)
	continue;
      perror("select()");
      exit(1);
    }

    if(FD_ISSET(tcp_sock, &rfds) && accept_client() < 0) {
      perror("accept_client()");
      exit(1);
    }

    for(uint32_t i = 0; i < clients.size(); ++i)
      if(FD_ISSET(clients[i], &rfds) && !packet_handler(clients[i]))
	disconnect_tcp_client(clients[i]);

  } while(n != 0);
}

int init_network(packet_handler_t handler)
{
  if(listen_init() < 0) {
    perror("init_network()");
    return -1;
  }
  packet_handler = handler;
  signal(SIGIO, handle_network_request);
  fcntl(tcp_sock, F_SETOWN, getpid());
  fcntl(tcp_sock, F_SETSIG, SIGIO); // SIGIO will be triggered when data is available
  return fcntl(tcp_sock, F_SETFL, O_ASYNC);
}