/
rack_socket.rb
162 lines (141 loc) · 5.72 KB
/
rack_socket.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
# This file is covered by the Ruby license. See COPYING for more details.
# Copyright (C) 2009-2010, Apple Inc. All rights reserved.
framework 'Foundation'
require 'CTParser'
require 'stringio'
CTParser # Making sure the Objective-C class is pre-loaded
module ControlTower
class RackSocket
VERSION = [1,0].freeze
def initialize(host, port, server, concurrent)
@app = server.app
@socket = TCPServer.new(host, port)
@socket.listen(50)
@status = :closed # Start closed and give the server time to start
if concurrent
@multithread = true
@request_queue = Dispatch::Queue.concurrent
puts "Caution! Wake turbulance from heavy aircraft landing on parallel runway.\n(Parallel Request Action ENABLED!)"
else
@multithread = false
@request_queue = Dispatch::Queue.new('com.apple.ControlTower.rack_socket_queue')
end
@request_group = Dispatch::Group.new
end
def open
@status = :open
while (@status == :open)
connection = @socket.accept
@request_queue.async(@request_group) do
env = { 'rack.errors' => $stderr,
'rack.multiprocess' => false,
'rack.multithread' => @multithread,
'rack.run_once' => false,
'rack.version' => VERSION }
begin
request_data = parse!(connection, env)
if request_data
request_data['REMOTE_ADDR'] = connection.addr[3]
status, headers, body = @app.call(request_data)
# Unless somebody's already set it for us (or we don't need it), set the Content-Length
unless (status == -1 ||
(status >= 100 and status <= 199) ||
status == 204 ||
status == 304 ||
headers.has_key?('Content-Length'))
headers['Content-Length'] = if body.respond_to?(:each)
size = 0
body.each { |x| size += x.bytesize }
size
else
body.bytesize
end
end
# TODO -- We don't handle keep-alive connections yet
headers['Connection'] = 'close'
resp = "HTTP/1.1 #{status}\r\n"
headers.each do |header, value|
resp << "#{header}: #{value}\r\n"
end
resp << "\r\n"
# Start writing the response
connection.write resp
# Finish writing out the body
if body.respond_to?(:each)
body.each do |chunk|
connection.write chunk
end
else
connection.write body
end
else
$stderr.puts "Error: No request data received!"
end
rescue EOFError, Errno::ECONNRESET, Errno::EPIPE, Errno::EINVAL
$stderr.puts "Error: Connection terminated!"
rescue Object => e
$stderr.puts "Error: Problem transmitting data -- #{e.inspect}"
$stderr.puts e.backtrace.join("\n")
ensure
# We should clean up after our tempfile, if we used one.
input = env['rack.input']
unlink input if input.class == Tempfile
connection.close rescue nil
end
end
end
end
def close
@status = :close
# You get 30 seconds to empty the request queue and get outa here!
Dispatch::Source.timer(30, 0, 1, Dispatch::Queue.concurrent) do
$stderr.puts "Timed out waiting for connections to close"
exit 1
end
@request_group.wait
@socket.close
end
private
def parse!(connection, env)
parser = Thread.current[:http_parser] ||= CTParser.new
parser.reset
data = NSMutableData.alloc.init
data.increaseLengthBy(1) # add sentinel
parsing_headers = true # Parse headers first
nread = 0
content_length = 0
content_uploaded = 0
connection_handle = NSFileHandle.alloc.initWithFileDescriptor(connection.fileno)
while (parsing_headers || content_uploaded < content_length) do
# Read the availableData on the socket and give up if there's nothing
incoming_bytes = connection_handle.availableData
return nil if incoming_bytes.length == 0
# Until the headers are done being parsed, we'll parse them
if parsing_headers
data.setLength(data.length - 1) # Remove sentinel
data.appendData(incoming_bytes)
data.increaseLengthBy(1) # Add sentinel
nread = parser.parseData(data, forEnvironment: env, startingAt: nread)
if parser.finished == 1
parsing_headers = false # We're done, now on to receiving the body
content_length = env['CONTENT_LENGTH'].to_i
content_uploaded = env['rack.input'].length
end
else # Done parsing headers, now just collect request body:
content_uploaded += incoming_bytes.length
env['rack.input'].appendData(incoming_bytes)
end
end
if content_length > 1024 * 1024
body_file = Tempfile.new('control-tower-request-body-')
NSFileHandle.alloc.initWithFileDescriptor(body_file.fileno).writeData(env['rack.input'])
body_file.rewind
env['rack.input'] = body_file
else
env['rack.input'] = StringIO.new(NSString.alloc.initWithData(env['rack.input'], encoding: NSASCIIStringEncoding))
end
# Returning what we've got...
return env
end
end
end