Initial upgrade implementation

Author: bigbes

Diff for: http/server.lua

@@ -16,7 +16,7 @@ local errno = require 'errno'
 local DETACHED = 101
 
 local function errorf(fmt, ...)
-    error(string.format(fmt, ...))
+    error(string.format(fmt, ...), 3)
 end
 
 local function sprintf(fmt, ...)
@@ -35,6 +35,18 @@ local function is_callable(obj)
     return false
 end
 
+local function is_callable(obj)
+    local t_obj = type(obj)
+    if t_obj == 'function' then
+        return true
+    end
+    if t_obj == 'table' then
+        local mt = getmetatable(obj)
+        return (type(mt) == 'table' and type(mt.__call) == 'function')
+    end
+    return false
+end
+
 local function uri_escape(str)
     local res = {}
     if type(str) == 'table' then
@@ -930,7 +942,7 @@ local function url_for_httpd(httpd, name, args, query)
     end
 end
 
-local function httpd_http11_parse_request(session, request_raw)
+local function httpd_parse_request(request_raw)
     local request_parsed = lib._parse_request(request_raw)
     if request_parsed.error then
         return nil, request_parsed.error
@@ -940,6 +952,14 @@ local function httpd_http11_parse_request(session, request_raw)
        request_parsed.path:find("./", nil, true) ~= nil then
         return nil, "invalid uri"
     end
+    return request_parsed
+end
+
+local function httpd_http11_parse_request(session, request_raw)
+    local request_parsed, err = httpd_parse_request(request_raw)
+    if not request_parsed then
+        return nil, err
+    end
     request_parsed.httpd = session.server
     request_parsed.s     = session.socket
     request_parsed.peer  = session.peer
@@ -976,6 +996,30 @@ local function httpd_http11_handler(session)
         return
     end
 
+    if p.headers['upgrade'] then
+        local proto_name = p.headers['upgrade']:lower()
+        local proto = session.server.upgrades[proto_name]
+        if not proto then
+            session:write('HTTP/1.1 400 Bad Request\r\n\r\n')
+            return false
+        else
+            local ok, upgrade_ok = pcall(proto.upgrade, session, p)
+            if not ok then
+                log.error("Failed to upgrade to '%s': %s", p.headers['upgrade'],
+                          upgrade_ok)
+                session:write('HTTP/1.1 500 Internal Error\r\n\r\n')
+                return false
+            elseif not upgrade_ok then
+                -- TODO: should we close connection, or we should retry again
+                return false
+            end
+
+            session.ctx.proto   = proto.name
+            session.ctx.handler = proto.handler
+            return true
+        end
+    end
+
     if p.headers['expect'] == '100-continue' then
         session:write('HTTP/1.1 100 Continue\r\n\r\n')
     elseif p.headers['expect'] then
@@ -1180,14 +1224,36 @@ local function httpd_start(self)
     return self
 end
 
+local function httpd_register_extension(self, ext_type, opts)
+    if ext_type:lower() == 'upgrade' then
+        if not (type(opts) == 'table') then
+            errorf("Upgrade extension argument should be table")
+        end
+        if not (type(opts.name) == 'string') then
+            errorf("Upgrade extension name should be %s", 'options.name', 'string')
+        end
+        if not is_callable(opts.upgrade) then
+            errorf("Upgrade extension callback should be callable")
+        end
+        if not is_callable(opts.handler) then
+            errorf("Upgrade extension handler should be callable")
+        end
+
+        self.upgrades[opts.name:lower()] = table.copy(opts)
+    else
+        errorf('Unknown extension type: %s', ext_type)
+    end
+end
+
 local httpd_methods = {
-    stop    = httpd_stop,
-    start   = httpd_start,
-    route   = add_route,
-    match   = match_route,
-    helper  = set_helper,
-    hook    = set_hook,
-    url_for = url_for_httpd,
+    stop               = httpd_stop,
+    start              = httpd_start,
+    route              = add_route,
+    match              = match_route,
+    helper             = set_helper,
+    hook               = set_hook,
+    url_for            = url_for_httpd,
+    register_extension = httpd_register_extension,
 }
 
 local httpd_mt = {
@@ -1223,10 +1289,11 @@ local function httpd_new(host, port, options)
         is_run  = false,
         options = options,
 
-        routes  = { },
-        iroutes = { },
-        helpers = { url_for = url_for_helper, },
-        hooks   = { },
+        routes   = { },
+        iroutes  = { },
+        helpers  = { url_for = url_for_helper, },
+        hooks    = { },
+        upgrades = { },
 
         -- caches
         cache   = { tpl = {}, ctx = {}, static = {}, },
@@ -1236,6 +1303,9 @@ local function httpd_new(host, port, options)
 end
 
 return {
-    DETACHED = DETACHED,
-    new = httpd_new
+    DETACHED      = DETACHED,
+    new           = httpd_new,
+    parse_headers = httpd_parse_request,
+    uri_escape    = uri_escape,
+    uri_unescape  = uri_unescape,
 }

Diff for: test/http.test.lua

@@ -9,8 +9,10 @@ local json = require('json')
 local yaml = require 'yaml'
 local urilib = require('uri')
 
-local test = tap.test("http")
-test:plan(7)
+local socket = require('socket')
+
+local test = tap.test("http"); test:plan(8)
+
 test:test("split_uri", function(test)
     test:plan(65)
     local function check(uri, rhs)
@@ -388,4 +390,101 @@ test:test("server requests", function(test)
     httpd:stop()
 end)
 
+test:test("upgrade", function(test)
+    test:plan(4)
+
+    local log = require('log')
+
+    local httpd = cfgserv()
+    httpd:start()
+    httpd:register_extension('upgrade', {
+        name = 'exist-error-upgrade',
+        upgrade = function() error('error') end,
+        handler = function() end,
+    })
+    httpd:register_extension('upgrade', {
+        name = 'exist-fails-upgrade',
+        upgrade = function(session)
+            session:write('HTTP/1.1 426 Upgrade Required\r\n\r\n')
+            return false
+        end,
+        handler = function() end,
+    })
+
+    local switching_header = 'HTTP/1.1 101 Switching Protocols\r\n' ..
+                             'Upgrade: exist\r\n' ..
+                             'Connection: Upgrade\r\n\r\n'
+    httpd:register_extension('upgrade', {
+        name = 'exists',
+        upgrade = function(session)
+            session:write(switching_header)
+            return true
+        end,
+        handler = function(session)
+            while true do
+                local in_data = session:read(24)
+                if in_data == '' then
+                    return false
+                end
+                session:write(in_data)
+            end
+        end,
+    })
+
+    test:test("upgrade failed, no protocol", function(test)
+        test:plan(1)
+        local r = http_client.get('http://127.0.0.1:12345/abc', {
+            headers = { upgrade = 'non-existent' }
+        })
+        test:is(r.status, 400, 'Error code is 400')
+    end)
+
+    test:test("upgrade failed, error while upgrade", function(test)
+        test:plan(1)
+        local r = http_client.get('http://127.0.0.1:12345/abc', {
+            headers = { upgrade = 'exist-error-upgrade' }
+        })
+        test:is(r.status, 500, 'Error code is 500')
+    end)
+
+    test:test("upgrade failed, upgrade return false", function(test)
+        test:plan(1)
+        local r = http_client.get('http://127.0.0.1:12345/abc', {
+            headers = { upgrade = 'exist-fails-upgrade' }
+        })
+        test:is(r.status, 426, 'Error code is 426')
+    end)
+
+    local ws_get_r = "GET /abc HTTP/1.1\r\nUpgrade:exists\r\n\r\n"
+
+    test:test("upgrade success, simple tcp echo", function(test)
+        test:plan(3)
+        local sck = socket.tcp_connect('127.0.0.1', 12345)
+        sck:write(ws_get_r)
+        local data = ''
+        while true do
+            local tdata = sck:read({ delimiter = { '\n\n', '\r\n\r\n' } })
+            if not tdata or tdata == '' or
+               tdata:endswith('\r\n\r\n') or tdata:endswith('\n\n') then
+                if tdata then
+                    data = data .. tdata
+                end
+                break
+            end
+        end
+
+        if not data:endswith('\r\n\r\n') then
+            test:fail('automatic fail')
+        else
+            test:is(#data, #switching_header, 'right http upgrade len')
+        end
+
+        local msg = ('x'):rep(24)
+        sck:write(msg)
+        local res = sck:read(#msg)
+        test:is(#res, #msg, 'echo is ok')
+        test:is(res, msg, 'echo is ok')
+    end)
+end)
+
 os.exit(test:check() == true and 0 or 1)