%% Copyright 2026 Benoit Chesneau %% %% Licensed 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. %% @doc WebSocket session state machine. %% %% Owns a single stream handle (supplied by the embedder through a %% `ws_transport' callback), drives the frame codec, auto-responds to %% control frames, and dispatches inbound messages to the user handler %% module. %% %% States: %% open — handshake complete, frames flowing in both directions. %% closing — a close frame was sent or received; the other direction %% may still finish. %% closed — both directions closed; the process is about to exit. %% %% Inbound-byte delivery is pull-based: after consuming a chunk the %% session calls `TransportMod:activate/1' to request the next one. The %% transport signals fresh bytes with `{ws_data, Handle, Bin}', peer %% close with `{ws_closed, Handle}', and socket-level errors with %% `{ws_error, Handle, Reason}'. -module(ws_session). -behaviour(gen_statem). -export([start_link/1, start/1]). -export([activate/1]). -export([send/2, close/2, close/3, stop/1]). -export([init/1, callback_mode/0, terminate/3, code_change/4]). -export([ready_wait/3, open/3, closing/3]). -include("../include/ws.hrl"). -record(st, { transport_mod :: module(), handle :: term(), role :: client | server, parser :: ws_frame:parser(), handler_mod :: module(), handler_state :: term(), close_sent = false :: boolean(), close_received = false :: boolean(), %% Close info received from the peer, reported to the handler's %% terminate/2 as {remote, Code, Reason} (or `remote' for a bare %% close frame with no status code). peer_close = undefined :: undefined | remote | {ws_close:code(), binary()}, %% permessage-deflate state once negotiated (RFC 7692): zlib %% streams plus per-direction context-takeover policy. deflate = undefined :: undefined | #{inflate := zlib:zstream(), deflate := zlib:zstream(), inflate_takeover := takeover | no_takeover, deflate_takeover := takeover | no_takeover, max_inflate := pos_integer() | infinity}, idle_timeout = ?WS_DEFAULT_IDLE_TIMEOUT :: timeout(), close_timeout = ?WS_DEFAULT_CLOSE_TIMEOUT :: timeout() }). -type start_opts() :: #{ transport := module(), handle := term(), role := client | server, handler := module(), handler_opts => term(), req => map(), parser_opts => map(), %% Negotiated permessage-deflate parameters (from %% ws_deflate:negotiate/2 or ws_deflate:parse_server_response/1). %% Enables compression on both directions of the session. deflate => ws_deflate:negotiated(), idle_timeout => timeout(), close_timeout => timeout() }. -export_type([start_opts/0]). %% --------------------------------------------------------------------- %% API -spec start_link(start_opts()) -> {ok, pid()} | {error, term()}. start_link(Opts) -> gen_statem:start_link(?MODULE, Opts, []). -spec start(start_opts()) -> {ok, pid()} | {error, term()}. start(Opts) -> gen_statem:start(?MODULE, Opts, []). %% @doc Tell the session the embedder has finished transport setup %% (ownership transfer and friends) — the session may now read bytes. -spec activate(pid()) -> ok. activate(Pid) -> gen_statem:cast(Pid, activate). -spec send(pid(), ws_frame:frame() | [ws_frame:frame()]) -> ok. send(Pid, FrameOrFrames) -> gen_statem:cast(Pid, {send, frames(FrameOrFrames)}). -spec close(pid(), ws_close:code()) -> ok. close(Pid, Code) -> close(Pid, Code, <<>>). -spec close(pid(), ws_close:code(), iodata()) -> ok. close(Pid, Code, Reason) -> gen_statem:cast(Pid, {close, Code, iolist_to_binary(Reason)}). -spec stop(pid()) -> ok. stop(Pid) -> gen_statem:stop(Pid). %% --------------------------------------------------------------------- %% gen_statem callbacks callback_mode() -> state_functions. init(#{transport := TM, handle := H, role := Role, handler := HMod} = Opts) -> ParserOpts = maps:get(parser_opts, Opts, #{}), Deflate = case maps:get(deflate, Opts, undefined) of undefined -> undefined; Negotiated -> init_deflate_state(Negotiated, Role, ParserOpts) end, ParserOpts1 = case Deflate of undefined -> ParserOpts; _ -> ParserOpts#{compress => true} end, Parser = ws_frame:init_parser(ParserOpts1#{role => Role}), HOpts = maps:get(handler_opts, Opts, #{}), Req = maps:get(req, Opts, #{}), St0 = #st{transport_mod = TM, handle = H, role = Role, parser = Parser, handler_mod = HMod, deflate = Deflate, idle_timeout = maps:get(idle_timeout, Opts, ?WS_DEFAULT_IDLE_TIMEOUT), close_timeout = maps:get(close_timeout, Opts, ?WS_DEFAULT_CLOSE_TIMEOUT)}, case HMod:init(Req, HOpts) of {ok, HState} -> post_init(St0#st{handler_state = HState}, []); {reply, Frames, HState} -> post_init(St0#st{handler_state = HState}, frames(Frames)); {stop, Reason} -> _ = TM:close(H), {stop, Reason} end. post_init(St, Frames) -> case send_frames(Frames, St) of {ok, St2} -> %% Defer activation until the embedder calls %% `ws_session:activate/1' — otherwise an earlier %% `inet:setopts' can deliver bytes to the original socket %% owner before `controlling_process/2' has run. {ok, ready_wait, St2}; {error, Reason} -> {stop, Reason} end. %% --- ready_wait state ------------------------------------------------- %% %% Frames can be queued via `send' while we wait; inbound bytes cannot %% arrive yet because the transport has not been activated. ready_wait(cast, activate, St = #st{transport_mod = TM, handle = H}) -> ok = TM:activate(H), {next_state, open, St, idle_action(St)}; ready_wait(cast, {send, Frames}, St) -> case send_frames(Frames, St) of {ok, St2} -> {keep_state, St2}; {error, Reason} -> {stop, Reason, St} end; ready_wait(cast, {close, Code, Reason}, St) -> initiate_close(Code, Reason, St); ready_wait(_EventType, _Msg, _St) -> keep_state_and_data. %% --- open state ------------------------------------------------------- open(info, Msg, St = #st{transport_mod = TM, handle = H, parser = P, handler_mod = HMod, handler_state = HState}) -> case TM:classify(Msg, H) of {ws_data, H, Bin} -> case ws_frame:parse(P, Bin) of {ok, Messages, P2} -> dispatch_messages(Messages, St#st{parser = P2}, open); {error, Reason, P2} -> abort_with_close(close_code_for(Reason), Reason, St#st{parser = P2}) end; {ws_closed, H} -> {stop, normal, St}; {ws_error, H, Reason} -> {stop, {transport_error, Reason}, St}; ignore -> case maybe_handler_info(HMod, Msg, HState) of passthrough -> keep_state_and_data; {ok, HState2} -> {keep_state, St#st{handler_state = HState2}}; {reply, Frames, HState2} -> case send_frames(frames(Frames), St#st{handler_state = HState2}) of {ok, St2} -> {keep_state, St2}; {error, Reason} -> {stop, Reason, St} end; {stop, Reason, HState2} -> {stop, Reason, St#st{handler_state = HState2}} end end; open(cast, {send, Frames}, St) -> case send_frames(Frames, St) of {ok, St2} -> {keep_state, St2}; {error, Reason} -> {stop, Reason, St} end; open(cast, {close, Code, Reason}, St) -> initiate_close(Code, Reason, St); open({timeout, idle}, idle, St) -> %% No inbound frame within `idle_timeout': close politely (1001). initiate_close(?WS_CLOSE_GOING_AWAY, <<>>, St). %% --- closing state ---------------------------------------------------- %% %% We have either sent or received a close frame. Keep draining until %% the other direction completes, then exit. closing(info, Msg, St = #st{transport_mod = TM, handle = H, parser = P}) -> case TM:classify(Msg, H) of {ws_data, H, Bin} -> case ws_frame:parse(P, Bin) of {ok, Messages, P2} -> dispatch_messages(Messages, St#st{parser = P2}, closing); {error, _Reason, P2} -> finish_close(St#st{parser = P2}) end; {ws_closed, H} -> finish_close(St); {ws_error, H, _R} -> finish_close(St); ignore -> keep_state_and_data end; closing(state_timeout, close, St) -> %% Peer never completed the close handshake within `close_timeout'. finish_close(St); closing(cast, {send, _}, _St) -> keep_state_and_data; closing(cast, {close, _C, _R}, _St) -> keep_state_and_data; closing(_Et, _Msg, _St) -> keep_state_and_data. %% --- shared dispatching ---------------------------------------------- dispatch_messages([], St, State) -> reactivate(St, State); dispatch_messages([Msg | Rest], St, State) -> case handle_message(Msg, St, State) of {continue, St2, State2} -> dispatch_messages(Rest, St2, State2); {abort, Code, Why, St2} -> abort_with_close(Code, Why, St2); {stop, Reason, St2} -> {stop, Reason, St2} end. handle_message(close, St, _State) -> handle_close_received(?WS_CLOSE_NORMAL, <<>>, St#st{peer_close = remote}); handle_message({close, Code, Reason}, St, _State) -> handle_close_received(Code, Reason, St#st{peer_close = {Code, Reason}}); handle_message({compressed, Kind, Payload}, St = #st{deflate = #{inflate := Z, inflate_takeover := Takeover, max_inflate := Max}}, State) -> Inflated = try ws_deflate:inflate(Z, Takeover, Payload, Max) catch _:_ -> {error, bad_deflate} end, case {Inflated, Kind} of {{ok, Data}, text} -> case ws_frame:valid_utf8(Data) of true -> invoke_handler({text, Data}, St, State); false -> {abort, ?WS_CLOSE_INVALID_UTF8, invalid_utf8, St} end; {{ok, Data}, binary} -> invoke_handler({binary, Data}, St, State); {{error, {inflate_too_big, _}}, _} -> {abort, ?WS_CLOSE_TOO_BIG, message_too_big, St}; {{error, _}, _} -> {abort, ?WS_CLOSE_PROTOCOL_ERROR, bad_deflate, St} end; handle_message({ping, Payload}, St = #st{}, State) -> case send_frames([{pong, Payload}], St) of {ok, St2} -> %% also notify the handler (optional — keeps parity with cowboy) invoke_handler({ping, Payload}, St2, State); {error, Reason} -> {stop, Reason, St} end; handle_message({pong, _} = M, St, State) -> invoke_handler(M, St, State); handle_message({text, _} = M, St, State) -> invoke_handler(M, St, State); handle_message({binary, _} = M, St, State) -> invoke_handler(M, St, State). invoke_handler(Msg, St = #st{handler_mod = HMod, handler_state = HState}, _State) -> case HMod:handle_in(Msg, HState) of {ok, HState2} -> {continue, St#st{handler_state = HState2}, _State}; {reply, Frames, HState2} -> case send_frames(frames(Frames), St#st{handler_state = HState2}) of {ok, St2} -> {continue, St2, _State}; {error, Reason} -> {stop, Reason, St} end; {stop, Reason, HState2} -> {stop, Reason, St#st{handler_state = HState2}} end. %% Peer sent us a close frame. Echo one back (once) and switch to %% closing; let the peer close the TCP side. handle_close_received(_Code, _Reason, St = #st{close_sent = true}) -> St2 = St#st{close_received = true}, finish_close(St2); handle_close_received(Code, _Reason, St) -> Echo = close_echo(Code), case send_frames([Echo], St) of {ok, St2} -> {stop, normal, St2#st{close_sent = true, close_received = true}}; {error, Reason} -> {stop, Reason, St} end. close_echo(Code) -> case ws_close:valid_on_wire(Code) of true -> {close, Code, <<>>}; false -> {close, ?WS_CLOSE_PROTOCOL_ERROR, <<>>} end. initiate_close(Code, Reason, St = #st{close_sent = false}) -> case send_frames([{close, Code, Reason}], St#st{close_sent = true}) of {ok, St2} -> {next_state, closing, St2, closing_actions(St2)}; {error, R} -> {stop, R, St} end; initiate_close(_C, _R, St) -> {keep_state, St}. %% Idle timer: a generic named timer so it survives unrelated events and %% is reset only when we choose (on inbound activity). `[]' disables it. idle_action(#st{idle_timeout = infinity}) -> []; idle_action(#st{idle_timeout = T}) -> [{{timeout, idle}, T, idle}]. %% Entering `closing': cancel the idle timer and arm the close-handshake %% timeout so a silent peer cannot pin the session open forever. closing_actions(#st{close_timeout = infinity}) -> [{{timeout, idle}, cancel}]; closing_actions(#st{close_timeout = T}) -> [{{timeout, idle}, cancel}, {state_timeout, T, close}]. abort_with_close(Code, Reason, St) -> _ = send_frames([{close, Code, io_code_reason(Reason)}], St), _ = (St#st.transport_mod):close(St#st.handle), {stop, normal, St#st{close_sent = true}}. io_code_reason(invalid_utf8) -> <<"invalid utf-8">>; io_code_reason(message_too_big) -> <<"message too big">>; io_code_reason(bad_close_code) -> <<"bad close code">>; io_code_reason(_) -> <<"protocol error">>. close_code_for(invalid_utf8) -> ?WS_CLOSE_INVALID_UTF8; close_code_for(message_too_big) -> ?WS_CLOSE_TOO_BIG; close_code_for(_) -> ?WS_CLOSE_PROTOCOL_ERROR. reactivate(St = #st{transport_mod = TM, handle = H}, StateName) -> _ = TM:activate(H), case StateName of open -> {keep_state, St, idle_action(St)}; closing -> {keep_state, St} end. finish_close(St) -> _ = (St#st.transport_mod):close(St#st.handle), {stop, normal, St}. maybe_handler_info(HMod, Msg, HState) -> case erlang:function_exported(HMod, handle_info, 2) of true -> HMod:handle_info(Msg, HState); false -> passthrough end. frames(L) when is_list(L) -> L; frames(F) -> [F]. send_frames([], St) -> {ok, St}; send_frames([F | Rest], St = #st{transport_mod = TM, handle = H, role = Role, deflate = Deflate}) -> Iolist = encode_out(F, Role, Deflate), case TM:send(H, Iolist) of ok -> send_frames(Rest, St); {error, _} = E -> E end. %% With permessage-deflate negotiated, data frames go out compressed %% (RSV1 set); control frames are never compressed (RFC 7692 §6). encode_out({Kind, Payload}, Role, #{deflate := Z, deflate_takeover := Takeover}) when Kind =:= text; Kind =:= binary -> Compressed = ws_deflate:deflate(Z, Takeover, Payload), ws_frame:encode_compressed({Kind, Compressed}, Role); encode_out(F, Role, _Deflate) -> ws_frame:encode(F, Role). %% Which context-takeover policy applies to each direction depends on %% the role: the `client_*' parameters govern client-to-server frames, %% the `server_*' parameters the reverse. init_deflate_state(Negotiated, Role, ParserOpts) -> {InTakeover, OutTakeover} = case Role of server -> {maps:get(client_context_takeover, Negotiated), maps:get(server_context_takeover, Negotiated)}; client -> {maps:get(server_context_takeover, Negotiated), maps:get(client_context_takeover, Negotiated)} end, #{inflate => ws_deflate:init_inflate(Negotiated, Role), deflate => ws_deflate:init_deflate(Negotiated, Role), inflate_takeover => InTakeover, deflate_takeover => OutTakeover, %% Inflated size is bounded by the same limit the parser applies %% to uncompressed messages, so a deflate bomb cannot bypass it. max_inflate => maps:get(max_message, ParserOpts, ?WS_DEFAULT_MAX_MESSAGE_SIZE)}. terminate(Reason, _StateName, #st{handler_mod = HMod, handler_state = HS, peer_close = PeerClose}) -> %% Surface the peer's close frame to the handler: {remote, Code, %% Reason} when it carried a status code, `remote' for a bare %% close. Other shutdowns pass the raw reason through. HReason = case PeerClose of undefined -> Reason; remote -> remote; {Code, Bin} -> {remote, Code, Bin} end, _ = case erlang:function_exported(HMod, terminate, 2) of true -> HMod:terminate(HReason, HS); false -> ok end, ok. code_change(_OldVsn, State, Data, _Extra) -> {ok, State, Data}.