github-fast-envd/github-fast-env/github-fast-envd.rb

#!/usr/bin/env ruby

require "base64"
require "io/console"
require "socket"

require "/data/github/current/config/environment"
require "irb"
require "irb/completion"

class ClientError < StandardError
	def initialize(message)
		super(message)
	end
end

class ClientScriptError < StandardError
	def initialize(source)
		super(source.message)
		@source = source
	end

	def source
		@source
	end
end

$original_stdin = $stdin.dup
$original_stdout = $stdout.dup
$original_stderr = $stderr.dup

control_socket_path = "/tmp/github-fast-envd.sock"

if File.exist?(control_socket_path) and File.socket?(control_socket_path)
	File.unlink(control_socket_path)
end

control_server = UNIXServer.new(control_socket_path)
File.chmod 0700, control_socket_path

$original_stderr.puts "serving control socket"

connection_id = 0

def open_pipe(path)
	if File.exist?(path) and File.pipe?(path)
		File.unlink(path)
	end

	File.mkfifo(path, mode = 0600)

	path
end

def read_command(control_socket)
	ready_ios = IO.select([control_socket], [], [], 10)

	if not ready_ios
		log "error", "timeout while communicating with github-fast-envd"
		exit 1
	end

	response = control_socket.readline.strip.split(" ", 2)

	if response.empty?
		log "error", "malformed response from github-fast-envd"
		exit 1
	end

	command = response[0]

	if command == "error"
		if response.length < 2
			log "error", "malformed error response from github-fast-envd"
			exit 1
		end

		error_message = response[1]

		log "error", "#{error_message}"
		exit 1
	end

	arguments = response[1].nil? ? [] : response[1].split(" ")

	return command, arguments
end

def set_up_named_pipes(control_socket, connection_id)
	pipe_base_path = "/tmp/github-fast-envd.#{connection_id}"
	pipe_base_path_encoded = Base64.encode64(pipe_base_path).delete("\n")

	stdin = open_pipe("#{pipe_base_path}.stdin")
	stdout = open_pipe("#{pipe_base_path}.stdout")
	stderr = open_pipe("#{pipe_base_path}.stderr")

	# TODO: support script arguments
	control_socket.puts "named-pipes #{pipe_base_path_encoded}"

	stdin = File.open(stdin, "r")
	stdin.sync = true
	stdout = File.open(stdout, "w")
	stdout.sync = true
	stderr = File.open(stderr, "w")
	stderr.sync = true

	$original_stderr.puts "  set up named pipes"

	control_socket.puts "ready"

	response = control_socket.readline.strip

	if response != "ready"
		raise ClientError.new "invalid command"
		Kernel.exit!
	end

	$stdin.reopen(stdin)
	$stdout.reopen(stdout)
	$stderr.reopen(stderr)
end

def set_up_pseudoterminal(control_socket, pseudoterminal_path)
	pseudoterminal_io = File.open(pseudoterminal_path, File::RDWR | File::NOCTTY)

	$original_stderr.puts "  connecting to pseudoterminal #{pseudoterminal_path}"

	$stdin.reopen(pseudoterminal_io)
	$stdout.reopen(pseudoterminal_io)
	$stderr.reopen(pseudoterminal_io)

	$original_stderr.puts "  connected to pseudoterminal #{pseudoterminal_path}"

	control_socket.puts "ready"
end

while true
	control_socket = control_server.accept
	$original_stderr.puts "- new connection"

	begin
		command, arguments = read_command(control_socket)

		if command != "new"
			raise ClientError.new "unexpected command"
		end

		if arguments.empty?
			raise ClientError.new "malformed command"
		end

		protocol_version = arguments[0]

		if protocol_version != "v1"
			raise ClientError.new "unsupported protocol version (#{protocol_version})"
		end

		mode = arguments[1]

		if mode == "named-pipes"
			if arguments.length < 3
				raise ClientError.new "malformed command"
			end
		elsif mode == "pseudoterminal"
			if arguments.length < 4
				raise ClientError.new "malformed command"
			end

			pseudoterminal_path = Base64.decode64(arguments[2])
		else
			raise ClientError.new "unknown mode (#{mode})"
		end

		script_path = Base64.decode64(arguments.last)

		connection_id += 1

		child_process = fork {
			process_id = Process.pid
			control_socket.puts "pid #{process_id}"

			exit_code = "unknown"

			if mode == "named-pipes"
				set_up_named_pipes(control_socket, connection_id)
			else
				set_up_pseudoterminal(control_socket, pseudoterminal_path)
			end

			$original_stderr.puts "  executing script #{script_path} (#{process_id})"

			begin
				begin
					load script_path, true
				rescue SystemExit => error
					$original_stderr.puts "  exit code: #{error.status}"
					exit_code = error.status
				rescue StandardError => error
					$stdin = $original_stdin
					$stdout = $original_stdout
					$stderr = $original_stderr

					raise ClientScriptError.new error
				end
			rescue ClientScriptError => error
				encoded_error_output = Base64.encode64(error.source.full_message).delete("\n")
				$original_stderr.puts "  error executing script, ignoring request"
				begin
					control_socket.puts "script_error #{encoded_error_output}"
				rescue
				end
			rescue ClientError => error
				$original_stderr.puts "  error communicating with client, ignoring request (#{error})"
				begin
					control_socket.puts "error #{error}"
				rescue
				end
			rescue StandardError => error
				$original_stderr.puts "  error, ignoring request (#{error})"
				begin
					control_socket.puts "error internal server error"
				rescue
				end
			end

			begin
				control_socket.puts "done #{exit_code}"
			rescue
			end
			control_socket.close

			$original_stderr.puts "  finished handling request (#{process_id})"

			Kernel.exit!
		}

		Process.detach(child_process)
	rescue ClientError => error
		$original_stderr.puts "  error communicating with client, ignoring request (#{error})"
		begin
			control_socket.puts "error #{error}"
		rescue
		end
	rescue StandardError => error
		$original_stderr.puts "  error, ignoring request (#{error})"
		begin
			control_socket.puts "error internal server error"
		rescue
		end
	end

	control_socket.close
end