/libdangerous-snmp-base/iosystem.cpp

#include "iosystem.hpp"

#include "dangerous/snmp/logger.hpp"

#include <iostream>
#include <stdexcept>

#include <sys/epoll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>

namespace dangerous { namespace snmp {

IoSystem::IoSystem() throw( Exception ) {
	isStopped = true;

	int result = socketpair( AF_UNIX, SOCK_DGRAM | SOCK_NONBLOCK, 0 /* protocol */, localSockets );
	if( result < 0 ) {
		throw Exception( "Could not create the local sockets." );
	}
}

IoSystem::~IoSystem() {
	if( logger.system( Logger::IO ) ) {
		logger.out() << "IoSystem::~IoSystem: Shutting down localSockets[0]." << std::endl;
	}
	shutdown( localSockets[ 0 ], SHUT_RDWR );
	close( localSockets[ 0 ] );
	
	if( logger.system( Logger::IO ) ) {
		logger.out() << "IoSystem::~IoSystem: Shutting down localSockets[1]." << std::endl;
	}
	shutdown( localSockets[ 1 ], SHUT_RDWR );
	close( localSockets[ 1 ] );
	
	if( logger.system( Logger::IO ) ) {
		logger.out() << "IoSystem::~IoSystem: Destructor complete." << std::endl;
	}
}

void IoSystem::registerInput( int handle, std::shared_ptr<ByteStream> byteStream ) {
	lock.lock();

	if( inMap.find( handle ) == inMap.end() ) {
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::registerInput: Registering input handle " << handle << "." << std::endl;
		}
		byteStream->setLive( true );
		inMap[ handle ] = byteStream;
	} else {
		// TODO: Double-register?
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::registerInput: Handle " << handle << " is double-registered." << std::endl;
		}
	}

	lock.unlock();
	
	wakeup();
}

void IoSystem::registerOutput( int handle, std::shared_ptr<ByteStream> byteStream ) {
	lock.lock();

	if( outMap.find( handle ) == outMap.end() ) {
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::registerInput: Registering output handle " << handle << "." << std::endl;
		}
		outMap[ handle ] = byteStream;
	} else {
		// TODO: Double-register?
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::registerOutput: Handle " << handle << " is double-registered." << std::endl;
		}
	}

	lock.unlock();
	
	wakeup();
}
	
void IoSystem::unregisterInput( int handle ) {
	lock.lock();

	if( inMap.find( handle ) != inMap.end() ) {
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::registerInput: Unregistering input handle " << handle << "." << std::endl;
		}
		inMap[ handle ]->setLive( false );
		inMap.erase( handle );
	} else {
		// TODO: never registered?
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::unregisterInput: Handle " << handle << " was never registered." << std::endl;
		}
	}

	lock.unlock();
	
	wakeup();
}

void IoSystem::unregisterOutput( int handle ) {
	lock.lock();

	if( outMap.find( handle ) != outMap.end() ) {
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::registerInput: Unregistering output handle " << handle << "." << std::endl;
		}
		outMap.erase( handle );
	} else {
		// TODO: never registered?
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::unregisterOutput: Handle " << handle << " was never registered." << std::endl;
		}
	}

	lock.unlock();
	
	wakeup();
}

void IoSystem::wakeup() {
	if( logger.system( Logger::IO ) ) {
		logger.out() << "IoSystem::wakeup: Sending wakeup message to handle " << localSockets[ 1 ] << "." << std::endl;
	}
	char buffer[ 1 ] = { 0 };
	send( localSockets[ 1 ], buffer, 1, 0 );
}

void IoSystem::main() {
	//! This is the file descriptor that we'll be using for "epoll".
	//! "epoll" is essentially a better "select" that doesn't run into
	//! problems with file descriptors whose values happen to be larger
	//! than one thousand.
	int epollHandle = -1; //!< Default to an invalid file descriptor.

	// Create the "epoll" handle.
	int errorNumber = 0;
	epollHandle = epoll_create1( EPOLL_CLOEXEC );
	errorNumber = errno;
	if( logger.system( Logger::IO ) ) {
		logger.out() << "IoSystem::main: epoll_create1 returned " << epollHandle << "." << std::endl;
	}
	if( epollHandle < 0 ) {
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::main: epoll_create1 error number: " << errorNumber << "." << std::endl;
		}
		throw Exception( "Could not create epoll handle." );
	}
	
	// Add the local control socket to the list of things that we want
	// to listen for.
	//
	// Note that "epoll" will KEEP this association, so we only need to
	// do it this one time.
	{
		struct epoll_event event;
		event.events = EPOLLIN;
		event.data.fd = localSockets[ 0 ];
		
		int errorNumber = 0;
		int status = epoll_ctl( epollHandle, EPOLL_CTL_ADD, event.data.fd, &event );
		errorNumber = errno;
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::main: epoll_ctl (add) returned " << status << "." << std::endl;
		}
		if( status < 0 ) {
			if( logger.system( Logger::IO ) ) {
				logger.out() << "IoSystem::main: epoll_ctl (add) error number: " << errorNumber << "." << std::endl;
			}
			throw Exception( "Could not add local socket to epoll handle." );
		}
	}

	// "epoll" will deal with a certain number of events at a time.
	// Since we're doing this in a loop, it really doesn't matter what
	// this value is.  All that matters is that we define it and use it.

	//! This is the number of events to get at a time when using "epoll" later.
	const unsigned int MAX_EVENTS = 10;

	//! This is the list of events that "epoll" will return later.
	struct epoll_event events[ MAX_EVENTS ];

	//! This is the mapping of in-map and out-map events that we CURRENTLY have
	//! set up for use with "epoll".
	std::map<int,int> existingEventMapping;

	while( ! isStopped ) {
		//! This is the mapping of in-map and out-map events that we WANT to have
		//! set up for use with "epoll".
		std::map<int,int> proposedEventMapping;

		// Update the proposed mapping.

		lock.lock();
		for( const auto& entry : inMap ) {
			// Before we add this socket to our list of sockets to listen for
			// more data from, make sure that it's still "open".  We will do
			// this by doing a quick "peek" for a single byte.  If we get back
			// exactly ZERO bytes (and NOT an error), then we know that this
			// socket has been closed; it's no good to us now.
			char byte = 0;
			int bytesAvailable = recv( entry.first, &byte, sizeof(byte), MSG_PEEK | MSG_DONTWAIT );
			if( bytesAvailable == 0 ) {
				continue;
			}

			proposedEventMapping[ entry.first ] |= EPOLLIN | EPOLLONESHOT;
		}

		for( const auto& entry : outMap ) {
			if( ! entry.second->hasNextPacket() ) {
				continue;
			}

			proposedEventMapping[ entry.first ] |= EPOLLOUT | EPOLLONESHOT;
		}
		lock.unlock();

		// Make any necessary changes.

		for( const auto& entry : proposedEventMapping ) {
			// If this mapping doesn't exist yet, then make it.
			if( existingEventMapping.find( entry.first ) == existingEventMapping.end() ) {
				struct epoll_event event;
				event.events = entry.second;
				event.data.fd = entry.first;
				
				int errorNumber = 0;
				int status = epoll_ctl( epollHandle, EPOLL_CTL_ADD, event.data.fd, &event );
				errorNumber = errno;
				if( logger.system( Logger::IO ) ) {
					logger.out() << "IoSystem::main: epoll_ctl (add) returned " << status << "." << std::endl;
				}

				// Note: we don't care about the error right now.  It either worked or it didn't.
				if( status < 0 ) {
					if( logger.system( Logger::IO ) ) {
						logger.out() << "IoSystem::main: epoll_ctl (add) error number: " << errorNumber << "." << std::endl;
					}
				}
			} else {
				// Otherwise, we need to update the mapping.

				struct epoll_event event;
				event.events = entry.second;
				event.data.fd = entry.first;
				
				int errorNumber = 0;
				int status = epoll_ctl( epollHandle, EPOLL_CTL_MOD, event.data.fd, &event );
				errorNumber = errno;
				if( logger.system( Logger::IO ) ) {
					logger.out() << "IoSystem::main: epoll_ctl (mod) returned " << status << "." << std::endl;
				}

				// Note: we don't care about the error right now.  It either worked or it didn't.
				if( status < 0 ) {
					if( logger.system( Logger::IO ) ) {
						logger.out() << "IoSystem::main: epoll_ctl (mod) error number: " << errorNumber << "." << std::endl;
					}
				}
			}
		}
		for( const auto& entry : existingEventMapping ) {
			// If we don't care about this handle anymore, then remove it from the "epoll" handle.
			if( proposedEventMapping.find( entry.first ) == proposedEventMapping.end() ) {
				struct epoll_event event;
				event.events = entry.second;
				event.data.fd = entry.first;
				
				int errorNumber = 0;
				int status = epoll_ctl( epollHandle, EPOLL_CTL_DEL, event.data.fd, &event );
				errorNumber = errno;
				if( logger.system( Logger::IO ) ) {
					logger.out() << "IoSystem::main: epoll_ctl (del) returned " << status << "." << std::endl;
				}

				// Note: we don't care about the error right now.  It either worked or it didn't.
				// For example, the file descriptor may have been closed by now.  "epoll" will
				// automatically remove it by itself in that case.
				if( status < 0 ) {
					if( logger.system( Logger::IO ) ) {
						logger.out() << "IoSystem::main: epoll_ctl (del) error number: " << errorNumber << "." << std::endl;
					}
				}
			}
		}

		// At this point, we have updated the existing mapping to match the state that we want,
		// so we can just swap the lists.
		existingEventMapping.swap( proposedEventMapping );


		//! This is how long to have "epoll_wait" wait before returning.  This CAN be "-1" to wait
		//! forever, but we're going to wake up at some point just to make sure that things are
		//! working.
		std::chrono::milliseconds timeout = std::chrono::seconds( 5 );

		//! This is the number of file handles that had events; that is, this is the number of
		//! entries in the "events" array that are valid.
		int handleCount = epoll_wait( epollHandle, events, MAX_EVENTS, timeout.count() );
		errorNumber = errno;
		if( logger.system( Logger::IO ) ) {
			logger.out() << "IoSystem::main: epoll_wait returned " << handleCount << "." << std::endl;
		}
		if( handleCount < 0 ) {
			// TODO: Error.
			if( logger.system( Logger::IO ) ) {
				logger.out() << "IoSystem::main: epoll_wait error number: " << errorNumber << "." << std::endl;
			}
			continue;
		}

		// If we have nothing to do, then just move on.
		if( handleCount == 0 ) {
			continue;
		}

		lock.lock();
		for( int i = 0; i < handleCount; i++ ) {
			struct epoll_event* event = &events[ i ];

			// Input event {{{
			if( ( event->events & EPOLLIN ) == EPOLLIN ) {
				// If "localSockets[0]" is available, we really don't care.  We were only using it
				// to wake ourselves up after making a change to "inMap" or "outMap".
				if( event->data.fd == localSockets[0] ) {
					if( logger.system( Logger::IO ) ) {
						logger.out() << "IoSystem::main: wakeup signal given." << std::endl;
					}

					const int BUFFER_SIZE = 8192;
					char buffer[ BUFFER_SIZE ];
					while( true ) {
						ssize_t bytesRead = recv( localSockets[0], buffer, (int)BUFFER_SIZE, 0 );
						if( bytesRead < 0 ) {
							int errorNumber = errno;
							if( errorNumber == EAGAIN || errorNumber == EWOULDBLOCK ) {
								break;
							} else {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: recv: Unexpected error: " << errorNumber << std::endl;
								}
								break;
							}
						}

						if( bytesRead == 0 ) {
							// TODO: End of file?
						} else {
							// Do nothing; we just need to eat up the bytes.
						}
					}
				} else {
					const auto& entryIterator = inMap.find( event->data.fd );
					if( entryIterator == inMap.end() ) {
						// This file descriptor is no longer in use.
					} else {
						const auto& entry = *entryIterator;
						if( logger.system( Logger::IO ) ) {
							logger.out() << "IoSystem::main: handle #" << entry.first << " is available for reading." << std::endl;
						}

						const int BUFFER_SIZE = 8192;
						char buffer[ BUFFER_SIZE ];
						while( true ) {
							ssize_t bytesRead = recv( entry.first, buffer, (int)BUFFER_SIZE, 0 );
							if( bytesRead < 0 ) {
								int errorNumber = errno;

								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Could not read any bytes from handle #" << entry.first << "." << std::endl;
								}
								if( errorNumber == EAGAIN || errorNumber == EWOULDBLOCK ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Socket " << entry.first << " is out of data for now." << std::endl;
									}
								} else if( errorNumber == EBADF ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Socket " << entry.first << " is an invalid file descriptor." << std::endl;
									}
								} else if( errorNumber == ECONNREFUSED ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Socket " << entry.first << " refused connection." << std::endl;
									}
								} else if( errorNumber == EFAULT ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Receive buffer is out of bounds." << std::endl;
									}
								} else if( errorNumber == EINTR ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Receive call was interrupted." << std::endl;
									}
								} else if( errorNumber == EINVAL ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Invalid argument passed to 'recv'." << std::endl;
									}
								} else if( errorNumber == ENOMEM ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Out of memory." << std::endl;
									}
								} else if( errorNumber == ENOTCONN ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Socket " << entry.first << " is not connected." << std::endl;
									}
								} else if( errorNumber == ENOTSOCK ) {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Error: Socket " << entry.first << " is not a socket." << std::endl;
									}
								} else {
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: recv: Unexpected error: " << errorNumber << std::endl;
									}
								}
								
								break;
							}
							
							if( logger.system( Logger::IO ) ) {
								logger.out() << "IoSystem::main: Read " << (int)bytesRead << " from handle #" << entry.first << "." << std::endl;
							}

							if( bytesRead == 0 ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: End of file: handle #" << entry.first << "." << std::endl;
								}
								break;
							} else {
								// Add the newly-read bytes to the buffer.
								unsigned int bytesWritten = 0;
								entry.second->writeBytes( buffer, bytesRead, bytesWritten );

								// Also, if this socket is a "datagram" socket, then we can issue the "endPacket"
								// call, which will tell the ByteStream that a full packet has been read.
								int socketType = 0;
								socklen_t socketTypeLength = sizeof( socketType );
								int result = getsockopt( entry.first, SOL_SOCKET, SO_TYPE, &socketType, &socketTypeLength );
								if( result < 0 ) {
									int errorNumber = errno;
									if( logger.system( Logger::IO ) ) {
										logger.out() << "IoSystem::main: Could not call 'getsockopt' on handle #" << entry.first << "; error number: " << errorNumber << "." << std::endl;
									}
									break;
								}
								if( socketType == SOCK_DGRAM ) {
									entry.second->endPacket();
								}
							}

							// Try to clean up the ByteStream if it has gotten too large.
							entry.second->reset();
						}
					}
				}
			}
			// }}}
			
			// Output event {{{
			if( ( event->events & EPOLLOUT ) == EPOLLOUT ) {
				const auto& entryIterator = outMap.find( event->data.fd );
				if( entryIterator == outMap.end() ) {
					// This file descriptor is no longer in use.
				} else {
					const auto& entry = *entryIterator;
					if( logger.system( Logger::IO ) ) {
						logger.out() << "IoSystem::main: handle #" << entry.first << " is available for writing." << std::endl;
					}

					const unsigned int BUFFER_SIZE = 8192;
					char buffer[ BUFFER_SIZE ];
					while( entry.second->remainingReadLength() > 0 ) {
						unsigned int bufferSize = BUFFER_SIZE < entry.second->remainingReadLength() ? BUFFER_SIZE : entry.second->remainingReadLength();
						unsigned int bytesCopied = 0;
						bool success = entry.second->readBytes( buffer, bufferSize, bytesCopied );
						if( ! success ) {
							if( logger.system( Logger::IO ) ) {
								logger.out() << "IoSystem::main Could not extract " << bufferSize << " bytes from the buffer." << std::endl;
							}
							break;
						}

						ssize_t bytesWritten = send( entry.first, buffer, bytesCopied, 0 );
						if( bytesWritten < 0 ) {
							int errorNumber = errno;

							if( logger.system( Logger::IO ) ) {
								logger.out() << "IoSystem::main: Could not send " << bytesCopied << " bytes to handle #" << entry.first << "." << std::endl;
							}

							if( errorNumber == EBADF ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Socket " << entry.first << " is an invalid file descriptor." << std::endl;
								}
							} else if( errorNumber == ECONNRESET ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Socket " << entry.first << " reset connection." << std::endl;
								}
							} else if( errorNumber == EDESTADDRREQ ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Peer address required." << std::endl;
								}
							} else if( errorNumber == EFAULT ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Receive buffer is out of bounds." << std::endl;
								}
							} else if( errorNumber == EINTR ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Receive call was interrupted." << std::endl;
								}
							} else if( errorNumber == EINVAL ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Invalid argument passed to 'send'." << std::endl;
								}
							} else if( errorNumber == EISCONN ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Already connected." << std::endl;
								}
							} else if( errorNumber == EMSGSIZE ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Message is too big to send atomically." << std::endl;
								}
							} else if( errorNumber == ENOBUFS ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Output queue is full." << std::endl;
								}
							} else if( errorNumber == ENOMEM ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Out of memory." << std::endl;
								}
							} else if( errorNumber == ENOTCONN ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Socket " << entry.first << " is not connected." << std::endl;
								}
							} else if( errorNumber == ENOTSOCK ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Socket " << entry.first << " is not a socket." << std::endl;
								}
							} else if( errorNumber == EOPNOTSUPP ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: Some flags are invalid." << std::endl;
								}
							} else if( errorNumber == EPIPE ) {
								if( logger.system( Logger::IO ) ) {
									logger.out() << "IoSystem::main: Error: The local end has been shut down." << std::endl;
								}
							} else {
								logger.out() << "IoSystem::main: send: Unexpected error: " << errorNumber << std::endl;
							}
								
							break;
						}

						if( logger.system( Logger::IO ) ) {
							logger.out() << "IoSystem::main: Sent " << bytesCopied << " bytes to handle #" << entry.first << "." << std::endl;
						}
					}
					entry.second->nextPacket();
						
					// Try to clean up the ByteStream if it has gotten too large.
					entry.second->reset();
				}
			}
			// }}}
		}
		lock.unlock();
	}

	close( epollHandle );
}

/**
 * This starts the the IoSystem as a thread.
 **/
void IoSystem::start() {
	if( ! isStopped ) {
		return;
	}
	
	if( logger.system( Logger::IO ) ) {
		logger.out() << "IoSystem::start: Starting I/O system..." << std::endl;
	}

	isStopped = false;
	thisThread = std::thread( std::bind( &IoSystem::main, this ) );
}

/**
 * This stops the IoSystem thread, if it's running.
 **/
void IoSystem::stop() {
	if( isStopped ) {
		return;
	}
	
	if( logger.system( Logger::IO ) ) {
		logger.out() << "IoSystem::stop: Stopping I/O system..." << std::endl;
	}

	isStopped = true;
	wakeup();
	thisThread.join();
}

} }