Skip to content

Commit

Permalink
ssh: ssh keep alive
Browse files Browse the repository at this point in the history
  • Loading branch information
alexandrejbr authored and Alexandre Rodrigues committed Nov 27, 2024
1 parent a83f774 commit 5b99b40
Show file tree
Hide file tree
Showing 6 changed files with 159 additions and 39 deletions.
10 changes: 9 additions & 1 deletion lib/ssh/src/ssh.hrl
Original file line number Diff line number Diff line change
Expand Up @@ -1246,7 +1246,15 @@ in the User's Guide chapter.
userauth_preference,
available_host_keys,
pwdfun_user_state,
authenticated = false
authenticated = false,

%% Keep-alive
alive_interval = infinity :: non_neg_integer() | infinity,
alive_count = 0 :: non_neg_integer(),
alive_started = false :: boolean(),
last_alive_at = 0 :: non_neg_integer(),
awaiting_keepalive_response = false :: boolean(),
alive_sent_probes = 0 :: non_neg_integer()
}).

-record(alg,
Expand Down
155 changes: 125 additions & 30 deletions lib/ssh/src/ssh_connection_handler.erl
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,11 @@
-define(call_disconnectfun_and_log_cond(LogMsg, DetailedText, StateName, D),
call_disconnectfun_and_log_cond(LogMsg, DetailedText, ?MODULE, ?LINE, StateName, D)).

-define(KEEP_ALIVE_REQUEST,
{ssh_msg_global_request,"[email protected]", true,<<>>}).
-define(KEEP_ALIVE_RESPONSE_F, {ssh_msg_request_failure}).
-define(KEEP_ALIVE_RESPONSE_S, {ssh_msg_request_success}).

%%====================================================================
%% Start / stop
%%====================================================================
Expand Down Expand Up @@ -439,11 +444,18 @@ init_ssh_record(Role, Socket, Opts) ->

init_ssh_record(Role, Socket, PeerAddr, Opts) ->
AuthMethods = ?GET_OPT(auth_methods, Opts),
{AliveCount, AliveIntervalSeconds} = ?GET_OPT(alive_params, Opts),
AliveInterval = case AliveIntervalSeconds of
V when is_integer(V) -> V * 1000;
infinity -> infinity
end,
S0 = #ssh{role = Role,
opts = Opts,
userauth_supported_methods = AuthMethods,
available_host_keys = available_hkey_algorithms(Role, Opts),
random_length_padding = ?GET_OPT(max_random_length_padding, Opts)
random_length_padding = ?GET_OPT(max_random_length_padding, Opts),
alive_interval = AliveInterval,
alive_count = AliveCount
},

{Vsn, Version} = ssh_transport:versions(Role, Opts),
Expand Down Expand Up @@ -752,6 +764,11 @@ handle_event(internal, #ssh_msg_debug{} = Msg, _StateName, D) ->
debug_fun(Msg, D),
keep_state_and_data;

handle_event(_, {conn_msg, Msg}, _, D = #data{ssh_params = Ssh})
when Ssh#ssh.awaiting_keepalive_response,
(Msg =:= ?KEEP_ALIVE_RESPONSE_F orelse Msg =:= ?KEEP_ALIVE_RESPONSE_S) ->
{keep_state, D#data{ssh_params = Ssh#ssh{awaiting_keepalive_response = false}}};

handle_event(internal, {conn_msg,Msg}, StateName, #data{connection_state = Connection0,
event_queue = Qev0} = D0) ->
Role = ?role(StateName),
Expand Down Expand Up @@ -833,6 +850,21 @@ handle_event({timeout,check_data_size}, _, StateName, D0) ->
keep_state_and_data
end;

handle_event({timeout, alive}, _, StateName, D = #data{ssh_params=Ssh}) ->
{TriggerFlag, Actions} = get_next_alive_timeout(Ssh),
case TriggerFlag of
true -> % timeout occured
triggered_alive(StateName, D, Ssh, Actions);
false -> % no timeout, check later
{keep_state, D, Actions}
end;

handle_event({timeout, renegotiation_alive}, _, StateName, D) ->
Details = "Renegotiation alive timeout reached.",
{Shutdown, D1} = ?send_disconnect(?SSH_DISCONNECT_CONNECTION_LOST, Details, StateName, D),
{stop, Shutdown, D1};


handle_event({call,From}, get_alg, _, D) ->
#ssh{algorithms=Algs} = D#data.ssh_params,
{keep_state_and_data, [{reply,From,Algs}]};
Expand Down Expand Up @@ -1132,48 +1164,49 @@ handle_event(info, {Proto, Sock, NewData}, StateName,
D0 = #data{socket = Sock,
transport_protocol = Proto,
ssh_params = SshParams}) ->
D1 = reset_alive(D0),
try ssh_transport:handle_packet_part(
D0#data.decrypted_data_buffer,
<<(D0#data.encrypted_data_buffer)/binary, NewData/binary>>,
D0#data.aead_data,
D0#data.undecrypted_packet_length,
D0#data.ssh_params)
D1#data.decrypted_data_buffer,
<<(D1#data.encrypted_data_buffer)/binary, NewData/binary>>,
D1#data.aead_data,
D1#data.undecrypted_packet_length,
D1#data.ssh_params)
of
{packet_decrypted, DecryptedBytes, EncryptedDataRest, Ssh1} ->
D1 = D0#data{ssh_params =
D2 = D1#data{ssh_params =
Ssh1#ssh{recv_sequence = ssh_transport:next_seqnum(Ssh1#ssh.recv_sequence)},
decrypted_data_buffer = <<>>,
undecrypted_packet_length = undefined,
aead_data = <<>>,
encrypted_data_buffer = EncryptedDataRest},
try
ssh_message:decode(set_kex_overload_prefix(DecryptedBytes,D1))
ssh_message:decode(set_kex_overload_prefix(DecryptedBytes,D2))
of
#ssh_msg_kexinit{} = Msg ->
{keep_state, D1, [{next_event, internal, prepare_next_packet},
{keep_state, D2, [{next_event, internal, prepare_next_packet},
{next_event, internal, {Msg,DecryptedBytes}}
]};

#ssh_msg_global_request{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_request_success{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_request_failure{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_open{} = Msg -> {keep_state, D1,
#ssh_msg_global_request{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_request_success{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_request_failure{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_open{} = Msg -> {keep_state, D2,
[{{timeout, max_initial_idle_time}, cancel} |
?CONNECTION_MSG(Msg)
]};
#ssh_msg_channel_open_confirmation{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_open_failure{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_window_adjust{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_data{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_extended_data{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_eof{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_close{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_request{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_failure{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_success{} = Msg -> {keep_state, D1, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_open_confirmation{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_open_failure{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_window_adjust{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_data{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_extended_data{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_eof{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_close{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_request{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_failure{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};
#ssh_msg_channel_success{} = Msg -> {keep_state, D2, ?CONNECTION_MSG(Msg)};

Msg ->
{keep_state, D1, [{next_event, internal, prepare_next_packet},
{keep_state, D2, [{next_event, internal, prepare_next_packet},
{next_event, internal, Msg}
]}
catch
Expand All @@ -1183,15 +1216,15 @@ handle_event(info, {Proto, Sock, NewData}, StateName,
?send_disconnect(?SSH_DISCONNECT_PROTOCOL_ERROR,
io_lib:format("Bad packet: Decrypted, but can't decode~n~p:~p~n~P",
[C,E,ST,MaxLogItemLen]),
StateName, D1),
StateName, D2),
{stop, Shutdown, D}
end;

{get_more, DecryptedBytes, EncryptedDataRest, AeadData, RemainingSshPacketLen, Ssh1} ->
%% Here we know that there are not enough bytes in
%% EncryptedDataRest to use. We must wait for more.
inet:setopts(Sock, [{active, once}]),
{keep_state, D0#data{encrypted_data_buffer = EncryptedDataRest,
{keep_state, D1#data{encrypted_data_buffer = EncryptedDataRest,
decrypted_data_buffer = DecryptedBytes,
undecrypted_packet_length = RemainingSshPacketLen,
aead_data = AeadData,
Expand All @@ -1201,15 +1234,15 @@ handle_event(info, {Proto, Sock, NewData}, StateName,
{Shutdown, D} =
?send_disconnect(?SSH_DISCONNECT_PROTOCOL_ERROR,
"Bad packet: bad mac",
StateName, D0#data{ssh_params=Ssh1}),
StateName, D1#data{ssh_params=Ssh1}),
{stop, Shutdown, D};

{error, {exceeds_max_size,PacketLen}} ->
{Shutdown, D} =
?send_disconnect(?SSH_DISCONNECT_PROTOCOL_ERROR,
io_lib:format("Bad packet: Size (~p bytes) exceeds max size",
[PacketLen]),
StateName, D0),
StateName, D1),
{stop, Shutdown, D}
catch
C:E:ST ->
Expand All @@ -1218,7 +1251,7 @@ handle_event(info, {Proto, Sock, NewData}, StateName,
?send_disconnect(?SSH_DISCONNECT_PROTOCOL_ERROR,
io_lib:format("Bad packet: Couldn't decrypt~n~p:~p~n~P",
[C,E,ST,MaxLogItemLen]),
StateName, D0),
StateName, D1),
{stop, Shutdown, D}
end;

Expand Down Expand Up @@ -1746,7 +1779,10 @@ start_rekeying(Role, D0) ->
send_bytes(SshPacket, D0),
D = D0#data{ssh_params = Ssh,
key_exchange_init_msg = KeyInitMsg},
{next_state, {kexinit,Role,renegotiate}, D, {change_callback_module,ssh_fsm_kexinit}}.
{next_state, {kexinit,Role,renegotiate}, D,
[{change_callback_module,ssh_fsm_kexinit},
{{timeout, alive}, cancel},
{{timeout, renegotiation_alive}, renegotiation_alive_timeout(Ssh), none}]}.


init_renegotiate_timers(_OldState, NewState, D) ->
Expand Down Expand Up @@ -2088,6 +2124,65 @@ update_inet_buffers(Socket) ->
_:_ -> ok
end.

%%%----------------------------------------------------------------
%%% Keep-alive

%% @doc Reset the last_alive timer on #data{ssh_params=#ssh{}} record
%% @private
reset_alive(D = #data{ssh_params = Ssh}) ->
D#data{ssh_params = reset_alive_ssh_params(Ssh)}.

%% @doc Update #data.ssh_params last_alive on an incoming SSH message
%% @private
reset_alive_ssh_params(SSH = #ssh{alive_interval = AliveInterval})
when is_integer(AliveInterval) ->
Now = erlang:monotonic_time(milli_seconds),
SSH#ssh{alive_sent_probes = 0,
last_alive_at = Now};
reset_alive_ssh_params(SSH) ->
SSH.

%% @doc Returns a pair of {TriggerFlag, Actions} where trigger flag indicates that
%% the timeout has been triggered already and it is time to disconnect, and
%% Actions may contain a new timeout action to check for the timeout again.
get_next_alive_timeout(#ssh{alive_interval = AliveInterval,
last_alive_at = LastAlive})
when erlang:is_integer(AliveInterval) ->
TimeToNextAlive = AliveInterval - (erlang:monotonic_time(milli_seconds) - LastAlive),
case TimeToNextAlive of
Trigger when Trigger =< 0 ->
%% Already it is time to disconnect, or to ping
{true, [{{timeout, alive}, AliveInterval, none}]};
TimeToNextAlive ->
{false, [{{timeout, alive}, TimeToNextAlive, none}]}
end;
get_next_alive_timeout(_) ->
{false, []}.

triggered_alive(StateName, D0 = #data{},
#ssh{alive_count = Count,
alive_sent_probes = SentProbesCount}, _Actions)
when SentProbesCount >= Count ->
%% Max probes count reached (equal to `alive_count`), we disconnect
Details = "Alive timeout triggered",
{Shutdown, D} = ?send_disconnect(?SSH_DISCONNECT_CONNECTION_LOST, Details, StateName, D0),
{stop, Shutdown, D};

triggered_alive(_StateName, Data, _Ssh = #ssh{alive_sent_probes = SentProbes}, Actions) ->
Data1 = send_msg(?KEEP_ALIVE_REQUEST, Data),
Ssh = Data1#data.ssh_params,
Now = erlang:monotonic_time(milli_seconds),
Ssh1 = Ssh#ssh{alive_sent_probes = SentProbes + 1,
awaiting_keepalive_response = true,
last_alive_at = Now},
{keep_state, Data1#data{ssh_params = Ssh1}, Actions}.

renegotiation_alive_timeout(#ssh{alive_interval = infinity}) ->
infinity;
renegotiation_alive_timeout(#ssh{alive_interval = Interval, alive_count = Count}) ->
Interval * Count.


%%%################################################################
%%%#
%%%# Tracing
Expand Down
4 changes: 3 additions & 1 deletion lib/ssh/src/ssh_fsm_kexinit.erl
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,9 @@ handle_event(internal, #ssh_msg_newkeys{} = Msg, {new_keys,Role,renegotiate}, D)
{ok, Ssh} = ssh_transport:handle_new_keys(Msg, D#data.ssh_params),
%% {ok, ExtInfo, Ssh} = ssh_transport:ext_info_message(Ssh1),
%% ssh_connection_handler:send_bytes(ExtInfo, D),
{next_state, {ext_info,Role,renegotiate}, D#data{ssh_params=Ssh}};
{next_state, {ext_info,Role,renegotiate}, D#data{ssh_params=Ssh},
[{{timeout, alive}, Ssh#ssh.alive_interval, none},
{{timeout, renegotiation_alive}, cancel}]};


%%% ######## {ext_info, client|server, init|renegotiate} ####
Expand Down
5 changes: 4 additions & 1 deletion lib/ssh/src/ssh_fsm_userauth_client.erl
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,10 @@ handle_event(internal, #ssh_msg_userauth_success{}, {userauth,client}, D0=#data{
ssh_auth:ssh_msg_userauth_result(success),
ssh_connection_handler:handshake(ssh_connected, D0),
D = D0#data{ssh_params=Ssh#ssh{authenticated = true}},
{next_state, {connected,client}, D, {change_callback_module,ssh_connection_handler}};
{next_state, {connected,client}, D,
[{{timeout, alive}, Ssh#ssh.alive_interval, none},
{change_callback_module,ssh_connection_handler}]};



%%---- userauth failure response to clientfrom the server
Expand Down
15 changes: 9 additions & 6 deletions lib/ssh/src/ssh_fsm_userauth_server.erl
Original file line number Diff line number Diff line change
Expand Up @@ -89,9 +89,9 @@ handle_event(internal,
{authorized, User, {Reply, Ssh1}} ->
D = connected_state(Reply, Ssh1, User, Method, D0),
{next_state, {connected,server}, D,
[set_max_initial_idle_timeout(D),
start_alive(D, [set_max_initial_idle_timeout(D),
{change_callback_module,ssh_connection_handler}
]};
])};
{not_authorized, {User, Reason}, {Reply, Ssh}} when Method == "keyboard-interactive" ->
retry_fun(User, Reason, D0),
D = ssh_connection_handler:send_msg(Reply, D0#data{ssh_params = Ssh}),
Expand Down Expand Up @@ -125,9 +125,9 @@ handle_event(internal, #ssh_msg_userauth_info_response{} = Msg, {userauth_keyboa
{authorized, User, {Reply, Ssh1}} ->
D = connected_state(Reply, Ssh1, User, "keyboard-interactive", D0),
{next_state, {connected,server}, D,
[set_max_initial_idle_timeout(D),
start_alive(D, [set_max_initial_idle_timeout(D),
{change_callback_module,ssh_connection_handler}
]};
])};
{not_authorized, {User, Reason}, {Reply, Ssh}} ->
retry_fun(User, Reason, D0),
D = ssh_connection_handler:send_msg(Reply, D0#data{ssh_params = Ssh}),
Expand All @@ -143,9 +143,9 @@ handle_event(internal, #ssh_msg_userauth_info_response{} = Msg, {userauth_keyboa
ssh_auth:handle_userauth_info_response({extra,Msg}, D0#data.ssh_params),
D = connected_state(Reply, Ssh1, User, "keyboard-interactive", D0),
{next_state, {connected,server}, D,
[set_max_initial_idle_timeout(D),
start_alive(D, [set_max_initial_idle_timeout(D),
{change_callback_module,ssh_connection_handler}
]
])
};


Expand Down Expand Up @@ -213,3 +213,6 @@ retry_fun(User, Reason, #data{ssh_params = #ssh{opts = Opts,
ok
end.

start_alive(#data{ssh_params = #ssh{alive_interval = AliveInterval}},
Actions) ->
[{{timeout, alive}, AliveInterval, none} | Actions].
9 changes: 9 additions & 0 deletions lib/ssh/src/ssh_options.erl
Original file line number Diff line number Diff line change
Expand Up @@ -852,6 +852,15 @@ default(common) ->
class => user_option
},

alive_params =>
#{default => {0, infinity},
chk => fun({AliveCount, AliveIntervalSeconds}) ->
check_pos_integer(AliveCount) andalso
check_timeout(AliveIntervalSeconds)
end,
class => user_option
},

%%%%% Undocumented
transport =>
#{default => ?DEFAULT_TRANSPORT,
Expand Down

0 comments on commit 5b99b40

Please sign in to comment.