Skip to content

feat(websocket): add enhanced proxy and plugin hook - #13939

Open
bzp2010 wants to merge 13 commits into
apache:masterfrom
bzp2010:bzp/feat-websocket-proxy-enhanced
Open

bzp2010 wants to merge 13 commits into
apache:masterfrom
bzp2010:bzp/feat-websocket-proxy-enhanced

Conversation

@bzp2010

@bzp2010 bzp2010 commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

Description

APISIX already supports WebSocket via enable_websocket on a Route/Service, but that path is a plain protocol upgrade: after the 101 handshake, nginx's own proxy_pass forwards the raw TCP stream, and no plugin phase ever sees an individual frame.

This PR adds a second, independent way to proxy WebSocket: a new Upstream scheme: ws/wss. Instead of delegating to proxy_pass, APISIX opens a real bidirectional WebSocket connection to the upstream itself (via resty.websocket.proxy) and exposes each frame to plugins through four new phases, so a plugin can observe or rewrite frames in flight - for logging, redaction, protocol translation, etc. enable_websocket and scheme: ws/wss don't combine; the former is ignored when the latter is set, since the connection never reaches the proxy_pass path it configures.

What's added

  • New Upstream scheme: ws and wss, alongside the existing http/https/grpc/grpcs/tcp/udp/tls (apisix/schema_def.lua, new @websocket_pass nginx location in apisix/cli/ngx_tpl.lua).
  • Four new plugin phases, mirroring the regular HTTP request/response lifecycle but per WebSocket frame instead of per request (apisix/init.lua):
    • ws_handshake - once, before APISIX connects to the upstream (in place of access).
    • ws_client_frame / ws_upstream_frame - once per frame, before it's forwarded onward.
    • ws_close - once, when the connection ends (in place of log; it runs first, then http_log_phase() still runs afterwards so passive health check reporting, after_balance, and the api_ctx/plugins/route-record tablepool releases all still happen for a WebSocket route the same way they do for an HTTP one).
  • core.websocket (apisix/core/websocket.lua): the API a plugin uses inside ws_client_frame/ws_upstream_frame to read and rewrite the frame currently in flight - core.websocket.client/.upstream (or core.websocket.get_role(role)) expose get_frame() (type, payload, last, code) and set_frame_data(payload).
  • example-plugin now implements all four phases as a runnable reference (appends -client/-upstream to text frames it sees, in each direction).
  • Retry across upstream nodes: since this path never goes through nginx's own upstream/balancer_by_lua* machinery, it can't call ngx.balancer.get_last_failure() to report a failed attempt before picking the next node. apisix.balancer.pick_server() gained an optional third prev_failure argument ({state, code}, shaped like get_last_failure()'s return value) so a caller outside balancer_by_lua* can report the outcome it already knows from its own connection attempt instead. websocket_content_phase uses this to retry up to upstream.retries times, mapping a connect error to a timeout (504) or a TCP failure (599) outcome. If every node fails before the handshake to the client completes, the client gets a real 502; a failure after the handshake (101 already sent) instead surfaces as the connection dropping with close code 1006, since there's no HTTP response left to send at that point.
  • Tests: t/node/websocket-proxy.t + .spec.mts (Jest/TS, following the t/plugin/lago.t-style --exec: pnpm test ... bootstrap), covering frame rewriting, binary frames, fragmentation/aggregate_fragments, ping/pong, both directions of a clean close and of an abrupt disconnect, retry (refused/timeout/multi-node/least_conn), passive health check marking a node unhealthy, URI/query-string forwarding (including with proxy-rewrite), and concurrent-connection isolation. New fixtures added to t/lib/server.lua, and t/APISIX.pm gained the same @websocket_pass location the real nginx template has, since the test-nginx harness maintains its own copy rather than rendering ngx_tpl.lua.
  • Docs: scheme documented as accepting ws/wss, with a note contrasting it against enable_websocket (admin-api.md, en+zh); the four new phases documented in plugin-develop.md's "extra phase" section with a core.websocket usage example; a pointer to that section added to terminology/plugin.md's phase lifecycle overview.
  • Dependency: api7-lua-resty-websocket 0.1.0-0 (apisix-master-0.rockspec), replacing stock openresty/lua-resty-websocket for resty.websocket.client/.server/.protocol, and providing resty.websocket.proxy (ported with additional hardening: on_frame callback errors are caught instead of crashing the worker, a fragment-count-limit bug is fixed, upstream TLS verification now defaults on, and thread cleanup on a failed ngx.thread.wait no longer leaks).

Known limitations (not addressed in this PR)

  • resty.websocket.proxy also accepts client_max_frame_size/client_max_fragments/upstream_max_frame_size/upstream_max_fragments, which bound how much a fragmented message gets buffered before aggregate_fragments joins it. These are deliberately left unset for now (see the comment in websocket_content_phase) - the plan is to expose them through a dedicated, dynamically configurable plugin later, not Upstream schema fields.

Demo

local bp_manager_mod = require("apisix.utils.batch-processor-manager")
local core           = require("apisix.core")

local ipairs = ipairs
local ngx    = ngx

local plugin_name = "ws-duration-logger"
local batch_processor_manager = bp_manager_mod.new("ws duration logger", plugin_name)

local schema = {
    type = "object",
    properties = {},
}

local metadata_schema = {
    type = "object",
    properties = {},
}

local _M = {
    version = 0.1,
    priority = 0,
    name = plugin_name,
    schema = batch_processor_manager:wrap_schema(schema),
    metadata_schema = batch_processor_manager:wrap_metadata_schema(metadata_schema),
}


function _M.check_schema(conf, schema_type)
    if schema_type == core.schema.TYPE_METADATA then
        return core.schema.check(metadata_schema, conf)
    end
    return core.schema.check(schema, conf)
end


-- Keyed by connection id, so this tracks every WebSocket connection
-- currently in flight across this worker, not just the one ws_close is
-- currently being called for. A plugin that only ever needs "this
-- connection's" data could stash it on ctx instead; this table is here to
-- demonstrate the pattern for a plugin that needs to see across connections
-- (for example, an in-worker count of currently open sessions).
local session_start = {}


function _M.ws_handshake(conf, ctx)
    local conn_id = core.id.gen_uuid_v4()
    -- ws_close only gets ctx back, not the frame data ws_handshake saw, so
    -- stash the id it needs to look its own session back up in session_start
    ctx.ws_duration_conn_id = conn_id

    session_start[conn_id] = {
        start_time = ngx.now(),
        user_id = core.request.header(ctx, "user_id"),
    }
end


function _M.ws_close(conf, ctx)
    local conn_id = ctx.ws_duration_conn_id
    if not conn_id then
        -- ws_close without a matching ws_handshake shouldn't happen, but
        -- don't report a bogus duration if it somehow does
        return
    end

    local session = session_start[conn_id]
    session_start[conn_id] = nil
    if not session then
        return
    end

    local entry = {
        connection_id = conn_id,
        user_id = session.user_id,
        duration = ngx.now() - session.start_time,
    }

    if batch_processor_manager:add_entry(conf, entry) then
        return
    end

    local func = function(entries)
        for _, e in ipairs(entries) do
            ngx.log(ngx.ERR, core.json.encode(e))
        end
        return true
    end

    batch_processor_manager:add_entry_to_new_processor(conf, entry, ctx, func)
end


return _M

Checklist

  • I have explained the need for this PR and the problem it solves
  • I have explained the changes or the new features added to this PR
  • I have added tests corresponding to this change
  • I have updated the documentation to reflect this change
  • I have verified that this change is backward compatible (If not, please discuss on the APISIX mailing list first)

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Core timeout, TLS, header forwarding, default-port, retry, and health-check semantics are currently incomplete.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds APISIX-managed ws/wss proxying with frame-level plugin hooks, retries, documentation, and integration tests.

Changes:

  • Introduces WebSocket proxy phases and frame mutation APIs.
  • Adds upstream selection, retry reporting, and nginx routing.
  • Adds dependency, fixtures, tests, and documentation.
File summaries
File Description
apisix/init.lua Implements WebSocket proxy lifecycle.
apisix/core/websocket.lua Provides frame access and mutation APIs.
apisix/core.lua Exports the WebSocket core module.
apisix/balancer.lua Accepts externally reported failures.
apisix/schema_def.lua Adds ws and wss schemes.
apisix/cli/ngx_tpl.lua Adds the WebSocket nginx location.
apisix/plugins/example-plugin.lua Demonstrates the new phases.
apisix-master-0.rockspec Adds the WebSocket dependency.
t/APISIX.pm Mirrors the nginx test location.
t/lib/server.lua Adds WebSocket test backends.
t/node/websocket-proxy.t Runs the integration suite.
t/node/websocket-proxy.spec.mts Tests proxy behavior and retries.
t/package.json Adds Node WebSocket packages.
t/pnpm-lock.yaml Locks the new packages.
docs/en/latest/admin-api.md Documents the new schemes.
docs/zh/latest/admin-api.md Adds Chinese API documentation.
docs/en/latest/plugin-develop.md Documents frame hooks and APIs.
docs/en/latest/terminology/plugin.md Extends the plugin phase overview.
Review details

Files not reviewed (1)

  • t/pnpm-lock.yaml: Generated file

Suppressed comments (8)

apisix/init.lua:1012

  • The pinned proxy constructs both WebSocket endpoints with resty.websocket.* defaults, whose maximum receive payload is 65,535 bytes. Because these options do not raise that limit, any unfragmented frame larger than 64 KiB closes the enhanced proxy even though the existing transparent WebSocket path accepts it. Update the dependency/proxy API to configure protocol-level receive limits consistently with the intended frame-limit policy.
    local ok, proxy, err = pcall(ws_proxy.new, {
        aggregate_fragments = true,
        recv_timeout = recv_timeout_ms,

apisix/init.lua:1008

  • Only connect and receive timeouts are passed into this transport; up_timeout.send is never applied. Frame writes consequently inherit whichever connect/read timeout was last set on the socket instead of the configured send timeout. The proxy API needs separate read/send timeout handling so the Upstream timeout contract remains effective.
    local connect_timeout_ms = up_timeout and up_timeout.connect and up_timeout.connect * 1000
    local recv_timeout_ms = up_timeout and up_timeout.read and up_timeout.read * 1000

apisix/init.lua:1081

  • Only host and server_name are supplied to the new upstream handshake, so client Cookie, Authorization, Origin, and Sec-WebSocket-Protocol headers—and header rewrites performed by plugins—are dropped. The proxy also synthesizes the downstream handshake independently instead of forwarding the upstream-selected subprotocol. This breaks authenticated and subprotocol-based WebSockets; propagate a filtered current request-header set and the negotiated response headers.
        ok, connect_err = proxy:connect(endpoint, {
            host = server.upstream_host,
            server_name = server.domain,
        })

apisix/init.lua:1079

  • Hard-coding the handshake Host to server.upstream_host makes every WebSocket upstream behave like pass_host: node. The default pass mode no longer forwards the client's Host, and pass_host: rewrite ignores upstream_host; retries also need to recompute node mode for the newly selected server.
            host = server.upstream_host,

apisix/init.lua:1081

  • The pinned proxy defaults WSS verification on and honors TLS policy only through connect options, but this call passes none of the Upstream tls fields. Consequently tls.verify: false is ignored, custom ca_certs cannot be used, and configured client certificates never reach a WSS upstream. Pass the complete TLS policy through, or reject unsupported WSS TLS configurations rather than silently changing their semantics.
        ok, connect_err = proxy:connect(endpoint, {
            host = server.upstream_host,
            server_name = server.domain,
        })

apisix/init.lua:1100

  • This maps every non-timeout handshake failure to a TCP failure. An upstream that responds normally with 400/401/403 produces an unexpected HTTP response code error here, so APISIX retries the request against other nodes and penalizes a healthy node as tcp_failures. Inspect the proxy client's response status and report an HTTP outcome (and avoid retrying non-retryable client/authentication responses).
        if connect_err and str_find(connect_err, "timeout", 1, true) then
            prev_failure = {state = "failed", code = 504}
        else
            prev_failure = {state = "failed", code = 599}

apisix/init.lua:1115

  • The last failed connection attempt is never reported as a timeout/TCP failure because prev_failure is only consumed while picking another server. With retries: 0 no failure is reported at all, and with retries the final node is omitted, so passive tcp_failures/timeouts thresholds can remain inaccurate. Report the final known failure before returning 502.
    if not ok then
        return core.response.exit(502)

apisix/schema_def.lua:512

  • The schema description remains stale after extending the enum, so generated schema/help output still says L7 supports only the four older schemes.
            description = "The scheme of the upstream." ..
                " For L7 proxy, it can be one of grpc/grpcs/http/https." ..
                " For L4 proxy, it can be one of tcp/tls/udp." ..
  • Files reviewed: 17/18 changed files
  • Comments generated: 6
  • Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread apisix/init.lua Outdated
Comment thread apisix/init.lua
Comment thread apisix/schema_def.lua
Comment thread docs/en/latest/plugin-develop.md Outdated
Comment thread docs/en/latest/terminology/plugin.md Outdated
Comment thread t/node/websocket-proxy.spec.mts Outdated
@bzp2010
bzp2010 marked this pull request as ready for review September 15, 2026 00:25
nic-6443
nic-6443 previously approved these changes Sep 15, 2026

@shreemaan-abhishek shreemaan-abhishek left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Comment thread apisix/init.lua
end
end

local ok, proxy, err = pcall(ws_proxy.new, {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Single frames larger than 65535 bytes kill the connection on this path.

ws_proxy.new() builds resty.websocket.client and resty.websocket.server without options (see the TODO at resty/websocket/proxy.lua:110), so both sides use the library default max_payload_len = 65535. protocol.lua checks that per frame in recv_frame(), before aggregate_fragments joins anything, so leaving client_max_frame_size & co. unset does not avoid it. An oversized frame marks the socket fatal and the forwarder tears down the whole connection.

enable_websocket has no such limit: after the 101, nginx relays raw bytes and never parses frames. Verified locally against the same upstream (client -> APISIX -> echo server that reports the received length):

frame size scheme: ws enable_websocket
65535 echoed echoed
65536 connection closed (failed to receive the first 2 bytes: closed) echoed
1 MiB - echoed

Most client libraries do not fragment by default, so an app sending a large JSON message or a file chunk breaks after switching to ws/wss. resty.websocket.proxy needs to pass max_payload_len (or max_recv_len/max_send_len) through to both constructors, with a default that does not cap ordinary traffic. A test with a >64K frame in each direction would cover it.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes sense, but it won’t be fixed in this PR because:

  1. I still need to modify the library so that this parameter can be passed through to lua-resty-websocket, which requires additional work.

  2. Even if we complete the work mentioned above, we won’t be able to modify it solely through configuration options in the upstream. I’m inclined to use the existing kafka-proxy plugin model by adding a websocket-proxy plugin to allow for custom configuration. These configurations are written to ctx via the plugin and affect the proxy path.

Therefore, there will be another PR following this one to upgrade dependencies and introduce the new plugin.

Comment thread apisix/init.lua

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we previously faced issues with running jest suite via test::nginx, ref: #12904, #12836, #12868.

The root cause of the issue was never found, thus I am not aligned to moving forward with running tests via this method.

@bzp2010 bzp2010 Sep 16, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@shreemaan-abhishek

I don't agree that these testing errors should be attributed to Jest itself.

If you look at these issues and PRs, it’s not hard to see that they’re all related to Lago tests. There have also been tests related to mcp-bridge in the past.
What they have in common is that they all require downloading resources, such as container images (docker compose up) and npm dependencies (pnpx) during test execution (rather than beforehand), which relies on network communication. As is well known, GitHub CI’s network is unstable, and operations like pulling container images often result in errors or timeouts.

The test::nginx test shell always assumes that tests can complete within an ideal timeframe, almost never considering that these tests may depend on external resources and could fail because of them.
And to be honest, our implementation of the custom --exec extension isn’t perfect either; it relies on mechanisms like ngx.log to report and output errors, which can sometimes cause it to exceed the maximum string length limit. In particular, when timeouts occur due to external resources, Jest may hang and fail to output logs.


However, I do agree with the view that we should no longer rely on the test::nginx shell to launch OpenResty in order to run JavaScript-based tests.

I will consider creating capabilities for JS testing that are consistent with test::nginx for managing the OpenResty lifecycle, including startup, reload, and restart operations. This way, testing will be entirely confined within the JS runtime. Essentially, this involves migrating similar functionality from a legacy Perl codebase (which is also a scripting language and relies on a runtime) to a JS implementation. At the same time, Jest will be replaced with the faster and more modern vitest.
This will begin shortly after this PR is completed, so I think it’s acceptable to tolerate this test running in this less formal manner for the time being. For reference, it does not involve the failure scenarios mentioned above; dependencies are always pulled in advance, and no additional containers need to be downloaded. Therefore, timeouts rarely occur.

@bzp2010
bzp2010 dismissed stale reviews from shreemaan-abhishek and nic-6443 via 25d9628 September 17, 2026 12:27
AlinsRan
AlinsRan previously approved these changes Sep 18, 2026
nic-6443
nic-6443 previously approved these changes Sep 18, 2026
@bzp2010
bzp2010 dismissed stale reviews from nic-6443 and AlinsRan via cc84fd8 September 18, 2026 07:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants