feat(websocket): add enhanced proxy and plugin hook

apisix-master-0.rockspec

@@ -32,6 +32,7 @@ description = {
 
 dependencies = {
     "lua-resty-ctxdump = 0.1-0",
+    "api7-lua-resty-websocket = 0.1.0-0",
     "api7-lua-resty-redis-connector = 0.13.0",
     "lyaml = 6.2.8-1",
     "api7-lua-resty-dns-client = 7.1.2-0",

apisix/balancer.lua

@@ -243,10 +243,30 @@ local function parse_server_for_upstream_host(picked_server, upstream_scheme)
 end
 
 
+-- reports a connection outcome (get_last_failure()-shaped state/code) for the
+-- node ctx.balancer_ip/balancer_port currently point at
+local function report_failure(ctx, checker, up_conf, state, code)
+    local host = up_conf.checks and up_conf.checks.active and up_conf.checks.active.host
+    local port = up_conf.checks and up_conf.checks.active and up_conf.checks.active.port
+    if state == "failed" then
+        if code == 504 then
+            checker:report_timeout(ctx.balancer_ip, port or ctx.balancer_port, host)
+        else
+            checker:report_tcp_failure(ctx.balancer_ip, port or ctx.balancer_port, host)
+        end
+    else
+        checker:report_http_status(ctx.balancer_ip, port or ctx.balancer_port, host, code)
+    end
+end
+
+
 -- pick_server will be called:
 -- 1. in the access phase so that we can set headers according to the picked server
 -- 2. each time we need to retry upstream
-local function pick_server(route, ctx)
+--
+-- prev_failure, when given, overrides get_last_failure() for callers outside
+-- balancer_by_lua* that already know their own connection's outcome.
+local function pick_server(route, ctx, prev_failure)
     local up_conf = ctx.upstream_conf
 
     local nodes_count = #up_conf.nodes
@@ -283,18 +303,13 @@ local function pick_server(route, ctx)
         end
 
         if checker then
-            local state, code = get_last_failure()
-            local host = up_conf.checks and up_conf.checks.active and up_conf.checks.active.host
-            local port = up_conf.checks and up_conf.checks.active and up_conf.checks.active.port
-            if state == "failed" then
-                if code == 504 then
-                    checker:report_timeout(ctx.balancer_ip, port or ctx.balancer_port, host)
-                else
-                    checker:report_tcp_failure(ctx.balancer_ip, port or ctx.balancer_port, host)
-                end
+            local state, code
+            if prev_failure then
+                state, code = prev_failure.state, prev_failure.code
             else
-                checker:report_http_status(ctx.balancer_ip, port or ctx.balancer_port, host, code)
+                state, code = get_last_failure()
             end
+            report_failure(ctx, checker, up_conf, state, code)
         end
     end
 
@@ -371,6 +386,17 @@ end
 _M.pick_server = pick_server
 
 
+-- reports a final failure with no next node to pick_server() for
+function _M.report_failure(ctx, prev_failure)
+    local checker = ctx.up_checker
+    if not checker then
+        return
+    end
+
+    report_failure(ctx, checker, ctx.upstream_conf, prev_failure.state, prev_failure.code)
+end
+
+
 -- Keyed by the `ca_certs` array itself: a config update always rebuilds that
 -- table, so a stale digest can never outlive the certificates it was made from.
 local ca_certs_digest_cache = core.lrucache.new({

apisix/cli/ngx_tpl.lua

@@ -1179,6 +1179,16 @@ http {
         }
         {% end %}
 
+        location @websocket_pass {
+            content_by_lua_block {
+                apisix.websocket_content_phase()
+            }
+
+            log_by_lua_block {
+                apisix.websocket_log_phase()
+            }
+        }
+
         {% if enabled_plugins["proxy-mirror"] then %}
         location = /proxy_mirror {
             internal;

apisix/core.lua

@@ -66,4 +66,5 @@ return {
     event       = require("apisix.core.event"),
     env         = require("apisix.core.env"),
     data_encryption = require("apisix.core.data_encryption"),
+    websocket   = require("apisix.core.websocket"),
 }

apisix/core/websocket.lua

@@ -0,0 +1,86 @@
+--
+-- Licensed to the Apache Software Foundation (ASF) under one or more
+-- contributor license agreements.  See the NOTICE file distributed with
+-- this work for additional information regarding copyright ownership.
+-- The ASF licenses this file to You under the Apache License, Version 2.0
+-- (the "License"); you may not use this file except in compliance with
+-- the License.  You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+--
+local ngx = ngx
+local tostring = tostring
+
+local ROLE_CLIENT = "client"
+local ROLE_UPSTREAM = "upstream"
+local CTX_KEY_CLIENT = "websocket_client"
+local CTX_KEY_UPSTREAM = "websocket_upstream"
+
+-- ngx.ctx is per-request and can only be accessed from within a request
+-- context, so it must be fetched inside each wrapped function, not cached
+-- as a module-level upvalue at require() time (which also runs during
+-- init_by_lua, before any request exists).
+
+local function wrap_stash_frame(key)
+  return function(frame)
+    ngx.ctx[key] = frame
+  end
+end
+
+
+local function wrap_get_frame(key)
+  return function()
+    return ngx.ctx[key]
+  end
+end
+
+
+local function wrap_set_frame_data(key)
+  return function(data)
+    ngx.ctx[key].payload = data
+  end
+end
+
+
+local function wrap_set_status(key)
+  return function(status)
+    ngx.ctx[key].code = status
+  end
+end
+
+
+local _M = {
+  ROLE_CLIENT = ROLE_CLIENT,
+  ROLE_UPSTREAM = ROLE_UPSTREAM,
+  [ROLE_CLIENT] = {
+    stash_frame = wrap_stash_frame(CTX_KEY_CLIENT),
+    get_frame = wrap_get_frame(CTX_KEY_CLIENT),
+    set_frame_data = wrap_set_frame_data(CTX_KEY_CLIENT),
+    set_status = wrap_set_status(CTX_KEY_CLIENT),
+    --drop_frame = wrap_drop_frame
+  },
+  [ROLE_UPSTREAM] = {
+    stash_frame = wrap_stash_frame(CTX_KEY_UPSTREAM),
+    get_frame = wrap_get_frame(CTX_KEY_UPSTREAM),
+    set_frame_data = wrap_set_frame_data(CTX_KEY_UPSTREAM),
+    set_status = wrap_set_status(CTX_KEY_UPSTREAM),
+  },
+}
+
+function _M.get_role(role)
+  if role == ROLE_CLIENT or role == "client" then
+    return _M[ROLE_CLIENT]
+  elseif role == ROLE_UPSTREAM or role == "upstream" then
+    return _M[ROLE_UPSTREAM]
+  else
+    return nil, "invalid role: " .. tostring(role)
+  end
+end
+
+return _M