<?php
/**
 * Simple lock server daemon that accepts lock/unlock requests.
 *
 * This code should not require MediaWiki setup or PHP files.
 *
 * This program 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.
 *
 * This program 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 this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 * http://www.gnu.org/copyleft/gpl.html
 *
 * @file
 * @ingroup LockManager Maintenance
 */

if ( php_sapi_name() !== 'cli' ) {
	die( "This is not a valid entry point.\n" );
}
error_reporting( E_ALL );

// Run the server...
set_time_limit( 0 );
LockServerDaemon::init(
	getopt( '', array(
		'address:', 'port:', 'authKey:',
		'lockTimeout::', 'maxClients::', 'maxBacklog::', 'maxLocks::',
	) )
)->main();

/**
 * Simple lock server daemon that accepts lock/unlock requests
 */
class LockServerDaemon {
	/** @var resource */
	protected $sock; // socket to listen/accept on
	/** @var Array */
	protected $sessions = array(); // (session => resource)
	/** @var Array */
	protected $deadSessions = array(); // (session => UNIX timestamp)

	/** @var LockHolder */
	protected $lockHolder;

	protected $address; // string IP address
	protected $port; // integer
	protected $authKey; // string key
	protected $lockTimeout; // integer number of seconds
	protected $maxBacklog; // integer
	protected $maxClients; // integer

	protected $startTime; // integer UNIX timestamp
	protected $ticks = 0; // integer counter

	/* @var LockServerDaemon */
	protected static $instance = null;

	/**
	 * @params $config Array
	 * @return LockServerDaemon
	 */
	public static function init( array $config ) {
		if ( self::$instance ) {
			throw new Exception( 'LockServer already initialized.' );
		}
		foreach ( array( 'address', 'port', 'authKey' ) as $par ) {
			if ( !isset( $config[$par] ) ) {
				die( "Usage: php LockServerDaemon.php " .
					"--address <address> --port <port> --authkey <key> " .
					"[--lockTimeout <seconds>] " .
					"[--maxLocks <integer>] [--maxClients <integer>] [--maxBacklog <integer>]"
				);
			}
		}
		self::$instance = new self( $config );
		return self::$instance;
	}

	/**
	 * @params $config Array
	 */
	protected function __construct( array $config ) {
		// Required parameters...
		$this->address = $config['address'];
		$this->port = $config['port'];
		$this->authKey = $config['authKey'];
		// Parameters with defaults...
		$this->lockTimeout = isset( $config['lockTimeout'] )
			? (int)$config['lockTimeout']
			: 60;
		$this->maxClients = isset( $config['maxClients'] )
			? (int)$config['maxClients']
			: 1000; // less than default FD_SETSIZE
		$this->maxBacklog = isset( $config['maxBacklog'] )
			? (int)$config['maxBacklog']
			: 100;
		$maxLocks = isset( $config['maxLocks'] )
			? (int)$config['maxLocks']
			: 10000;

		$this->lockHolder = new LockHolder( $maxLocks );
	}

	/**
	 * @return void
	 */
	protected function setupServerSocket() {
		if ( !function_exists( 'socket_create' ) ) {
			throw new Exception( "PHP sockets extension missing from PHP CLI mode." );
		}
		$sock = socket_create( AF_INET, SOCK_STREAM, SOL_TCP );
		if ( $sock === false ) {
			throw new Exception( "socket_create(): " . socket_strerror( socket_last_error() ) );
		}
		socket_set_option( $sock, SOL_SOCKET, SO_REUSEADDR, 1 ); // bypass 2MLS
		socket_set_nonblock( $sock ); // don't block on accept()
		if ( socket_bind( $sock, $this->address, $this->port ) === false ) {
			throw new Exception( "socket_bind(): " .
				socket_strerror( socket_last_error( $sock ) ) );
		} elseif ( socket_listen( $sock, $this->maxBacklog ) === false ) {
			throw new Exception( "socket_listen(): " .
				socket_strerror( socket_last_error( $sock ) ) );
		}
		$this->sock = $sock;
		$this->startTime = time();
	}

	/**
	 * Entry-point function that listens to the server socket, accepts
	 * new clients, and recieves/responds to requests to lock resources.
	 */
	public function main() {
		$this->setupServerSocket(); // setup listening socket
		$socketArray = new SocketArray(); // sockets being serviced
		$socketArray->addSocket( $this->sock ); // add listening socket
		do {
			list( $read, $write ) = $socketArray->socketsForSelect();
			if ( socket_select( $read, $write, $except = NULL, NULL ) < 1 ) {
				continue; // wait
			}
			// Check if there is a client trying to connect...
			if ( in_array( $this->sock, $read ) && $socketArray->size() < $this->maxClients ) {
				$newSock = socket_accept( $this->sock );
				if ( $newSock ) {
					socket_set_option( $newSock, SOL_SOCKET, SO_KEEPALIVE, 1 );
					socket_set_nonblock( $newSock ); // don't block on read()/write()
					$socketArray->addSocket( $newSock );
				}
			}
			// Loop through all the clients that have data to read...
			foreach ( $read as $read_sock ) {
				if ( $read_sock === $this->sock ) {
					continue; // skip listening socket
				}
				// Avoids PHP_NORMAL_READ per https://bugs.php.net/bug.php?id=33471
				$data = socket_read( $read_sock, 65535 );
				// Check if the client is disconnected
				if ( $data === false || $data === '' ) {
					$socketArray->closeSocket( $read_sock );
					$this->recordDeadSocket( $read_sock ); // remove session
				// Check if we reached the end of a message
				} elseif ( substr( $data, -1 ) === "\n" ) {
					// Newline is the last char (given ping-pong message usage)
					$cmd = $socketArray->readRcvBuffer( $read_sock ) . $data;
					// Perform the requested command...
					$response = $this->doCommand( rtrim( $cmd ), $read_sock );
					// Send the response to the client...
					$socketArray->appendSndBuffer( $read_sock, $response . "\n" );
				// Otherwise, we just have more message data to append
				} elseif ( !$socketArray->appendRcvBuffer( $read_sock, $data ) ) {
					$socketArray->closeSocket( $read_sock ); // too big
					$this->recordDeadSocket( $read_sock ); // remove session
				}
			}
			// Loop through all the clients that have data to write...
			foreach ( $write as $write_sock ) {
				$bytes = socket_write( $write_sock, $socketArray->readSndBuffer( $write_sock ) );
				// Check if the client is disconnected
				if ( $bytes === false ) {
					$socketArray->closeSocket( $write_sock );
					$this->recordDeadSocket( $write_sock ); // remove session
				// Otherwise, truncate these bytes from the start of the write buffer
				} else {
					$socketArray->consumeSndBuffer( $write_sock, $bytes );
				}
			}
			// Prune dead locks every few socket events...
			if ( ++$this->ticks >= 9 ) {
				$this->ticks = 0;
				$this->purgeExpiredLocks();
			}
		} while ( true );
	}

	/**
	 * @param $data string
	 * @param $sourceSock resource
	 * @return string
	 */
	protected function doCommand( $data, $sourceSock ) {
		$cmdArr = $this->getCommand( $data );
		if ( is_string( $cmdArr ) ) {
			return $cmdArr; // error
		}
		list( $function, $session, $type, $resources ) = $cmdArr;
		// On first command, track the session => sock correspondence
		if ( !isset( $this->sessions[$session] ) ) {
			$this->sessions[$session] = $sourceSock;
			unset( $this->deadSessions[$session] ); // renew if dead
		}
		if ( $function === 'ACQUIRE' ) {
			return $this->lockHolder->lock( $session, $type, $resources );
		} elseif ( $function === 'RELEASE' ) {
			return $this->lockHolder->unlock( $session, $type, $resources );
		} elseif ( $function === 'RELEASE_ALL' ) {
			return $this->lockHolder->release( $session );
		} elseif ( $function === 'STAT' ) {
			return $this->stat();
		}
		return 'INTERNAL_ERROR';
	}

	/**
	 * @param $data string
	 * @return Array
	 */
	protected function getCommand( $data ) {
		$m = explode( ':', $data ); // <session, key, command, type, values>
		if ( count( $m ) == 5 ) {
			list( $session, $key, $command, $type, $values ) = $m;
			if ( sha1( $session . $command . $type . $values . $this->authKey ) !== $key ) {
				return 'BAD_KEY';
			} elseif ( strlen( $session ) !== 32 ) {
				return 'BAD_SESSION';
			}
			$values = explode( '|', $values );
			if ( $command === 'ACQUIRE' ) {
				$needsLockArgs = true;
			} elseif ( $command === 'RELEASE' ) {
				$needsLockArgs = true;
			} elseif ( $command === 'RELEASE_ALL' ) {
				$needsLockArgs = false;
			} elseif ( $command === 'STAT' ) {
				$needsLockArgs = false;
			} else {
				return 'BAD_COMMAND';
			}
			if ( $needsLockArgs ) {
				if ( $type !== 'SH' && $type !== 'EX' ) {
					return 'BAD_TYPE';
				}
				foreach ( $values as $value ) {
					if ( strlen( $value ) !== 31 ) {
						return 'BAD_FORMAT';
					}
				}
			}
			return array( $command, $session, $type, $values );
		}
		return 'BAD_FORMAT';
	}

	/**
	 * Remove a socket's corresponding session from tracking and
	 * store it in the dead session tracking if it still has locks.
	 *
	 * @param $socket resource
	 * @return bool
	 */
	protected function recordDeadSocket( $socket ) {
		$session = array_search( $socket, $this->sessions );
		if ( $session !== false ) {
			unset( $this->sessions[$session] );
			// Record recently killed sessions that still have locks
			if ( $this->lockHolder->sessionHasLocks( $session ) ) {
				$this->deadSessions[$session] = time();
			}
			return true;
		}
		return false;
	}

	/**
	 * Clear locks for sessions that have been dead for a while
	 *
	 * @return integer Number of sessions purged
	 */
	protected function purgeExpiredLocks() {
		$count = 0;
		$now = time();
		foreach ( $this->deadSessions as $session => $timestamp ) {
			if ( ( $now - $timestamp ) > $this->lockTimeout ) {
				$this->lockHolder->release( $session );
				unset( $this->deadSessions[$session] );
				++$count;
			}
		}
		return $count;
	}

	/**
	 * Get the current timestamp and memory usage
	 *
	 * @return string
	 */
	protected function stat() {
		return ( time() - $this->startTime ) . ':' . memory_get_usage();
	}
}

/**
 * LockServerDaemon helper class that keeps track socket states
 */
class SocketArray {
	/* @var Array */
	protected $clients = array(); // array of client sockets
	/* @var Array */
	protected $rBuffers = array(); // corresponding socket read buffers
	/* @var Array */
	protected $wBuffers = array(); // corresponding socket write buffers

	const BUFFER_SIZE = 65535;

	/**
	 * @return Array (list of sockets to read, list of sockets to write)
	 */
	public function socketsForSelect() {
		$rSockets = array();
		$wSockets = array();
		foreach ( $this->clients as $key => $socket ) {
			if ( $this->wBuffers[$key] !== '' ) {
				$wSockets[] = $socket; // wait for writing to unblock
			} else {
				$rSockets[] = $socket; // wait for reading to unblock
			}
		}
		return array( $rSockets, $wSockets );
	}

	/**
	 * @return integer Number of client sockets
	 */
	public function size() {
		return count( $this->clients );
	}

	/**
	 * @param $sock resource
	 * @return bool
	 */
	public function addSocket( $sock ) {
		$this->clients[] = $sock;
		$this->rBuffers[] = '';
		$this->wBuffers[] = '';
		return true;
	}

	/**
	 * @param $sock resource
	 * @return bool
	 */
	public function closeSocket( $sock ) {
		$key = array_search( $sock, $this->clients );
		if ( $key === false ) {
			return false;
		}
		socket_close( $sock );
		unset( $this->clients[$key] );
		unset( $this->rBuffers[$key] );
		unset( $this->wBuffers[$key] );
		return true;
	}

	/**
	 * @param $sock resource
	 * @param $data string
	 * @return bool
	 */
	public function appendRcvBuffer( $sock, $data ) {
		$key = array_search( $sock, $this->clients );
		if ( $key === false ) {
			return false;
		} elseif ( ( strlen( $this->rBuffers[$key] ) + strlen( $data ) ) > self::BUFFER_SIZE ) {
			return false;
		}
		$this->rBuffers[$key] .= $data;
		return true;
	}

	/**
	 * @param $sock resource
	 * @return string|bool
	 */
	public function readRcvBuffer( $sock ) {
		$key = array_search( $sock, $this->clients );
		if ( $key === false ) {
			return false;
		}
		$data = $this->rBuffers[$key];
		$this->rBuffers[$key] = ''; // consume data
		return $data;
	}

	/**
	 * @param $sock resource
	 * @param $data string
	 * @return bool
	 */
	public function appendSndBuffer( $sock, $data ) {
		$key = array_search( $sock, $this->clients );
		if ( $key === false ) {
			return false;
		} elseif ( ( strlen( $this->wBuffers[$key] ) + strlen( $data ) ) > self::BUFFER_SIZE ) {
			return false;
		}
		$this->wBuffers[$key] .= $data;
		return true;
	}

	/**
	 * @param $sock resource
	 * @return bool
	 */
	public function readSndBuffer( $sock ) {
		$key = array_search( $sock, $this->clients );
		if ( $key === false ) {
			return false;
		}
		return $this->wBuffers[$key];
	}

	/**
	 * @param $sock resource
	 * @param $bytes integer
	 * @return bool
	 */
	public function consumeSndBuffer( $sock, $bytes ) {
		$key = array_search( $sock, $this->clients );
		if ( $key === false ) {
			return false;
		}
		$this->wBuffers[$key] = (string)substr( $this->wBuffers[$key], $bytes );
		return true;
	}
}

/**
 * LockServerDaemon helper class that keeps track of the locks
 */
class LockHolder {
	/** @var Array */
	protected $shLocks = array(); // (key => session => 1)
	/** @var Array */
	protected $exLocks = array(); // (key => session)

	/** @var Array */
	protected $sessionIndexSh = array(); // (session => key => 1)
	/** @var Array */
	protected $sessionIndexEx = array(); // (session => key => 1)
	protected $lockCount = 0; // integer

	protected $maxLocks; // integer

	/**
	 * @params $maxLocks integer Maximum number of locks to allow
	 */
	public function __construct( $maxLocks ) {
		$this->maxLocks = $maxLocks;
	}

	/**
	 * @param $session string
	 * @return bool
	 */
	public function sessionHasLocks( $session ) {
		return isset( $this->sessionIndexSh[$session] )
			|| isset( $this->sessionIndexEx[$session] );
	}

	/**
	 * @param $session string
	 * @param $type string
	 * @param $keys Array
	 * @return string
	 */
	public function lock( $session, $type, array $keys ) {
		if ( ( $this->lockCount + count( $keys ) ) > $this->maxLocks ) {
			return 'TOO_MANY_LOCKS';
		}
		if ( $type === 'SH' ) {
			// Check if any keys are already write-locked...
			foreach ( $keys as $key ) {
				if ( isset( $this->exLocks[$key] ) && $this->exLocks[$key] !== $session ) {
					return 'CANT_ACQUIRE';
				}
			}
			// Acquire the read-locks...
			foreach ( $keys as $key ) {
				$this->set_sh_lock( $key, $session );
			}
			return 'ACQUIRED';
		} elseif ( $type === 'EX' ) {
			// Check if any keys are already read-locked or write-locked...
			foreach ( $keys as $key ) {
				if ( isset( $this->exLocks[$key] ) && $this->exLocks[$key] !== $session ) {
					return 'CANT_ACQUIRE';
				}
				if ( isset( $this->shLocks[$key] ) ) {
					foreach ( $this->shLocks[$key] as $otherSession => $x ) {
						if ( $otherSession !== $session ) {
							return 'CANT_ACQUIRE';
						}
					}
				}
			}
			// Acquire the write-locks...
			foreach ( $keys as $key ) {
				$this->set_ex_lock( $key, $session );
			}
			return 'ACQUIRED';
		}
		return 'INTERNAL_ERROR';
	}

	/**
	 * @param $session string
	 * @param $type string
	 * @param $keys Array
	 * @return string
	 */
	public function unlock( $session, $type, array $keys ) {
		if ( $type === 'SH' ) {
			foreach ( $keys as $key ) {
				$this->unset_sh_lock( $key, $session );
			}
			return 'RELEASED';
		} elseif ( $type === 'EX' ) {
			foreach ( $keys as $key ) {
				$this->unset_ex_lock( $key, $session );
			}
			return 'RELEASED';
		}
		return 'INTERNAL_ERROR';
	}

	/**
	 * @param $session string
	 * @return string
	 */
	public function release( $session ) {
		if ( isset( $this->sessionIndexSh[$session] ) ) {
			foreach ( $this->sessionIndexSh[$session] as $key => $x ) {
				$this->unset_sh_lock( $key, $session );
			}
		}
		if ( isset( $this->sessionIndexEx[$session] ) ) {
			foreach ( $this->sessionIndexEx[$session] as $key => $x ) {
				$this->unset_ex_lock( $key, $session );
			}
		}
		return 'RELEASED_ALL';
	}

	/**
	 * @param $key string
	 * @param $session string
	 * @return void
	 */
	protected function set_sh_lock( $key, $session ) {
		if ( !isset( $this->shLocks[$key][$session] ) ) {
			$this->shLocks[$key][$session] = 1;
			$this->sessionIndexSh[$session][$key] = 1;
			++$this->lockCount; // we are adding a lock
		}
	}

	/**
	 * @param $key string
	 * @param $session string
	 * @return void
	 */
	protected function set_ex_lock( $key, $session ) {
		if ( !isset( $this->exLocks[$key][$session] ) ) {
			$this->exLocks[$key] = $session;
			$this->sessionIndexEx[$session][$key] = 1;
			++$this->lockCount; // we are adding a lock
		}
	}

	/**
	 * @param $key string
	 * @param $session string
	 * @return void
	 */
	protected function unset_sh_lock( $key, $session ) {
		if ( isset( $this->shLocks[$key][$session] ) ) {
			unset( $this->shLocks[$key][$session] );
			if ( !count( $this->shLocks[$key] ) ) {
				unset( $this->shLocks[$key] );
			}
			unset( $this->sessionIndexSh[$session][$key] );
			if ( !count( $this->sessionIndexSh[$session] ) ) {
				unset( $this->sessionIndexSh[$session] );
			}
			--$this->lockCount;
		}
	}

	/**
	 * @param $key string
	 * @param $session string
	 * @return void
	 */
	protected function unset_ex_lock( $key, $session ) {
		if ( isset( $this->exLocks[$key] ) && $this->exLocks[$key] === $session ) {
			unset( $this->exLocks[$key] );
			unset( $this->sessionIndexEx[$session][$key] );
			if ( !count( $this->sessionIndexEx[$session] ) ) {
				unset( $this->sessionIndexEx[$session] );
			}
			--$this->lockCount;
		}
	}
}