Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
456 lines (433 sloc) 11.8 KB
#include "tcpsocket.h"
TCPSocket::TCPSocket(unsigned int max_recv_buf_len,
unsigned int max_post_buf_len)
: port_(0),
recv_buf_(NULL),
post_buf_(NULL),
socket_(INVALID_SOCKET),
connect_status_(CLOSED),
socket_type_(COMM_SOCKET),
post_begin_(0), post_end_(0),
recv_begin_(0), recv_end_(0),
max_recv_buf_len_(max_recv_buf_len),
max_post_buf_len_(max_post_buf_len)
{
memset(ip_, 0, sizeof(ip_));
memset(&addr_, 0, sizeof(addr_));
recv_buf_ = new unsigned char[max_recv_buf_len_];
post_buf_ = new unsigned char[max_post_buf_len_];
}
TCPSocket::~TCPSocket()
{
if (OPENED == connect_status_
&& socket_ != INVALID_SOCKET) {
Close();
}
if (post_buf_) {
delete[] post_buf_;
post_buf_ = NULL;
}
if (recv_buf_) {
delete[] recv_buf_;
recv_buf_ = NULL;
}
}
TCPSocket& TCPSocket::operator=(const TCPSocket& t)
{
if (this != &t) {
t.get_ip(ip_);
addr_ = t.get_addr();
port_ = t.get_port();
socket_ = t.get_socket();
socket_type_ = t.get_socket_type();
connect_status_ = t.get_connect_status();
recv_begin_ = t.get_recv_begin();
recv_end_ = t.get_recv_end();
post_begin_ = t.get_post_begin();
post_end_ = t.get_post_end();
max_recv_buf_len_ = t.get_max_recv_buf_len();
max_post_buf_len_ = t.get_max_post_buf_len();
if (recv_buf_)
memcpy(recv_buf_, t.get_recv_buf(),
max_recv_buf_len_);
if (post_buf_)
memcpy(post_buf_, t.get_post_buf(),
max_post_buf_len_);
}
return *this;
}
int TCPSocket::Init(unsigned int max_recv_buf_len,
unsigned int max_post_buf_len)
{
port_ = 0;
socket_ = INVALID_SOCKET;
connect_status_ = CLOSED;
socket_type_ = COMM_SOCKET;
recv_begin_ = recv_end_ = 0;
post_begin_ = post_end_ = 0;
max_recv_buf_len_ = max_recv_buf_len;
max_post_buf_len_ = max_post_buf_len;
memset(&ip_, 0, sizeof(ip_));
if (recv_buf_) {
delete[] recv_buf_;
recv_buf_ = NULL;
}
if (post_buf_) {
delete[] post_buf_;
post_buf_ = NULL;
}
recv_buf_ = new unsigned char[max_recv_buf_len_];
post_buf_ = new unsigned char[max_post_buf_len_];
return 0;
}
int TCPSocket::CreateServer(unsigned short port, char* ip)
{
if (connect_status_ != CLOSED
&& socket_ != INVALID_SOCKET) {
Close();
}
socket_type_ = LISTEN_SOCKET;
socket_ = socket(AF_INET, SOCK_STREAM, 0);
if (INVALID_SOCKET == socket_) {
connect_status_ = CLOSED;
return -1;
}
memset(&addr_, 0, sizeof(addr_));
addr_.sin_family = AF_INET;
addr_.sin_port = htons(port);
if (ip) {
addr_.sin_addr.s_addr = inet_addr(ip);
} else {
addr_.sin_addr.s_addr = htonl(INADDR_ANY);
}
if (bind(socket_,
(const sockaddr*)&addr_,
sizeof(addr_)) < 0) {
Close();
return -2;
}
if (listen(socket_, 10) < 0) {
Close();
return -3;
}
SetReuseAddr();
recv_begin_ = recv_end_ = 0;
post_begin_ = post_end_ = 0;
connect_status_ = OPENED;
return 0;
}
int TCPSocket::CreateClient(char* ip)
{
if (connect_status_ != CLOSED
&& socket_ != INVALID_SOCKET) {
Close();
}
socket_type_ = COMM_SOCKET;
socket_ = socket(AF_INET, SOCK_STREAM, 0);
if (INVALID_SOCKET == socket_) {
connect_status_ = CLOSED;
return -1;
}
if (ip) {
sockaddr_in addr;
memset(&addr, 0, sizeof(addr));
addr.sin_family = AF_INET;
addr.sin_addr.s_addr = inet_addr(ip);
if (bind(socket_, (const sockaddr*)&addr, sizeof(addr)) < 0) {
Close();
return -2;
}
}
recv_begin_ = recv_end_ = 0;
post_begin_ = post_end_ = 0;
connect_status_ = OPENED;
return 0;
}
int TCPSocket::SetNonBlock()
{
int flags = fcntl(socket_, F_GETFL, 0);
if (flags < 0) return flags;
flags |= O_NONBLOCK;
flags |= O_NDELAY;
return fcntl(socket_, F_SETFL, flags);
}
int TCPSocket::SetReuseAddr()
{
int reuse_port_flag = 1;
if (setsockopt(socket_, SOL_SOCKET, SO_REUSEADDR,
&reuse_port_flag, sizeof(reuse_port_flag)) < 0) {
return -1;
}
return 0;
}
int TCPSocket::Close()
{
if (socket_ != INVALID_SOCKET) {
close(socket_);
socket_ = INVALID_SOCKET;
}
connect_status_ = CLOSED;
post_begin_ = post_end_ = 0;
if (post_buf_) {
delete[] post_buf_;
post_buf_ = NULL;
}
return 0;
}
int TCPSocket::Accept(TCPSocket& comm_socket)
{
int tmp_socket;
sockaddr_in addr;
memset(&addr, 0, sizeof(addr));
socklen_t len = sizeof(addr);
if ((tmp_socket = accept(socket_,
(sockaddr*)&addr,
&len)) < 0) {
return -1;
}
comm_socket.Init();
comm_socket.set_addr(addr);
comm_socket.set_socket(tmp_socket);
comm_socket.set_socket_type(COMM_SOCKET);
comm_socket.set_connect_status(CONNECTED);
comm_socket.SetNonBlock();
return 0;
}
int TCPSocket::ConnectTo(char* ip, unsigned short port)
{
if (!ip) return -1;
if (connect_status_ != OPENED
|| INVALID_SOCKET == socket_) return -2;
sockaddr_in addr;
memset(&addr, 0, sizeof(addr));
addr.sin_family = AF_INET;
addr.sin_port = htons(port);
addr.sin_addr.s_addr = inet_addr(ip);
if (connect(socket_, (const sockaddr*)&addr,
sizeof(addr)) < 0) {
Close();
return -3;
}
SetNonBlock();
memcpy(ip_, ip, IP_LEN);
port_ = port;
post_begin_ = post_end_ = 0;
recv_begin_ = recv_end_ = 0;
connect_status_ = CONNECTED;
return 0;
}
int TCPSocket::RecvBytes()
{
if (INVALID_SOCKET == socket_
|| connect_status_ != CONNECTED) {
return ERR_RECV_NOSOCK;
}
if (recv_begin_ == recv_end_) {
recv_begin_ = recv_end_ = 0;
}
int recvd_bytes_cnt = 0;
int retval = 0;
do {
if (recv_end_ >= max_recv_buf_len_) {
retval = ERR_RECV_NOAPPBUFF;
break;
}
recvd_bytes_cnt = recv(socket_,
&recv_buf_[recv_end_],
max_recv_buf_len_ - recv_end_, 0);
if (recvd_bytes_cnt > 0) {
recv_end_ += recvd_bytes_cnt;
} else if (0 == recvd_bytes_cnt) {
retval = ERR_RECV_REMOTE;
break;
} else if (EAGAIN != errno) {
retval = ERR_RECV_FAILED;
break;
}
} while (recvd_bytes_cnt > 0);
return retval;
}
int TCPSocket::GetOneCode(unsigned int& code_len, unsigned char* code)
{
if (!code) return -1;
int retval = 0;
int max_code_buf_len = code_len;
int recvd_bytes_cnt
= recv_end_ - recv_begin_;
if (recvd_bytes_cnt <= MSG_PREFIX_LEN) {
return -2;
}
int need_bytes_cnt = ntohl(*((int*)&recv_buf_[recv_begin_]));
if (need_bytes_cnt <= 0) {
recv_begin_ = recv_end_ = 0;
return -3;
}
if (code_len < need_bytes_cnt + MSG_PREFIX_LEN)
return -4;
if (int(recvd_bytes_cnt - MSG_PREFIX_LEN) < need_bytes_cnt) {
// check if the rest of recv_buf_ is enough.
if (recv_begin_ + need_bytes_cnt + MSG_PREFIX_LEN
>= max_recv_buf_len_) {
memmove(recv_buf_,
&recv_buf_[recv_begin_], recvd_bytes_cnt);
recv_begin_ = 0;
recv_end_ = recvd_bytes_cnt;
if (recv_begin_ + need_bytes_cnt + MSG_PREFIX_LEN
> max_recv_buf_len_) {
// recv_buf_ not enough.
return -5;
}
}
return -6;
}
code_len = need_bytes_cnt + MSG_PREFIX_LEN;
memcpy(code, &recv_buf_[recv_begin_], code_len);
recv_begin_ += code_len;
if (recv_begin_ == recv_end_)
recv_begin_ = recv_end_ = 0;
return 0;
}
int TCPSocket::SendReservedBytes()
{
return SendBytes();
}
int TCPSocket::GetReservedBytes()
{
return (int)(post_end_ - post_begin_);
}
int TCPSocket::SendOneCode(unsigned int code_len,
unsigned char* code, bool is_reserve)
{
if (INVALID_SOCKET == socket_
|| connect_status_ != CONNECTED) {
return ERR_SEND_NOSOCK;
}
int retval = SendBytes();
if (retval == ERR_SEND_NOTCPBUFF) {
if (is_reserve) {
if (ReserveBytes(code_len, code) < 0)
return ERR_SEND_FAILED;
}
return ERR_SEND_NOTCPBUFF;
}
int sent_bytes = 0;
unsigned char* tmp_code = code;
unsigned int left_bytes = code_len;
while (1) {
sent_bytes = send(socket_,
tmp_code, left_bytes, 0);
if (sent_bytes > 0) {
tmp_code += sent_bytes;
left_bytes -= sent_bytes;
if (!left_bytes) {
return sent_bytes;
}
} else if (sent_bytes < 0) {
if (EAGAIN == errno) {
if (!is_reserve) {
if (left_bytes == code_len)
return ERR_SEND_NOTCPBUFF;
}
if (ReserveBytes(left_bytes, tmp_code) < 0) {
return ERR_SEND_FAILED;
} else
return ERR_SEND_NOTCPBUFF;
}
else
return ERR_SEND_FAILED;
} else
return ERR_SEND_FAILED;
}
return sent_bytes;
}
int TCPSocket::GetTCPSendBufLen()
{
int retval = -1;
int opt_val = 0;
socklen_t opt_len = sizeof(opt_val);
if (!getsockopt(socket_,
SOL_SOCKET, SO_SNDBUF,
&opt_val, &opt_len))
retval = opt_val;
return retval;
}
int TCPSocket::GetTCPRecvBufLen()
{
int retval = -1;
int opt_val = 0;
socklen_t opt_len = sizeof(opt_val);
if (!getsockopt(socket_,
SOL_SOCKET, SO_RCVBUF,
&opt_val, &opt_len))
retval = opt_val;
return retval;
}
int TCPSocket::SetTCPSendBufLen(int len)
{
int retval = -1;
if (len < 0) return retval;
int opt_val = len;
socklen_t opt_len = sizeof(opt_val);
if (!setsockopt(socket_,
SOL_SOCKET, SO_SNDBUF,
(const void*)&opt_val,
opt_len))
retval = opt_val;
return retval;
}
int TCPSocket::SetTCPRecvBufLen(int len)
{
int retval = -1;
if (len < 0) return retval;
int opt_val = len;
socklen_t opt_len = sizeof(opt_val);
if (!setsockopt(socket_,
SOL_SOCKET, SO_RCVBUF,
(const void*)&opt_val,
opt_len))
retval = opt_val;
return retval;
}
int TCPSocket::SendBytes()
{
int sent_bytes = 0;
int left_bytes = post_end_ - post_begin_;
if (!left_bytes) return 0;
unsigned char* tmp_post_buf = &post_buf_[post_begin_];
while (1) {
sent_bytes = send(socket_,
tmp_post_buf, left_bytes, 0);
if (sent_bytes > 0) {
tmp_post_buf += sent_bytes;
left_bytes -= sent_bytes;
post_begin_ += sent_bytes;
if (!left_bytes) {
post_begin_ = post_end_ = 0;
return sent_bytes;
}
} else if (sent_bytes < 0) {
if (EAGAIN == errno)
return ERR_SEND_NOTCPBUFF;
else
return ERR_SEND_FAILED;
} else
return ERR_SEND_FAILED;
}
return sent_bytes;
}
int TCPSocket::ReserveBytes(unsigned int bytes_cnt, unsigned char* bytes)
{
if (!bytes || !bytes_cnt) return -1;
int existing_bytes_cnt = post_end_ - post_begin_;
if (existing_bytes_cnt + bytes_cnt < max_post_buf_len_) {
if (post_begin_ > 0) {
memmove(post_buf_, &post_buf_[post_begin_],
existing_bytes_cnt);
post_begin_ = 0;
}
post_end_ = existing_bytes_cnt;
memcpy(&post_buf_[post_end_], bytes, bytes_cnt);
post_end_ += bytes_cnt;
} else return -2;
return 0;
}