diff --git a/test/integration/client_idle_timeout_test.exs b/test/integration/client_idle_timeout_test.exs index ce8e8e95..04d3708e 100644 --- a/test/integration/client_idle_timeout_test.exs +++ b/test/integration/client_idle_timeout_test.exs @@ -1,6 +1,8 @@ defmodule Supavisor.Integration.ClientIdleTimeoutTest do use Supavisor.DataCase, async: false + alias Supavisor.Support.ProtocolClient + @moduletag :integration @idle_timeout_ms 1000 @@ -10,7 +12,7 @@ defmodule Supavisor.Integration.ClientIdleTimeoutTest do end test "server disconnects a client that sits idle past client_idle_timeout", %{db_conf: db_conf} do - {:gen_tcp, sock} = db_conf |> idle_tenant() |> scram_connect() + sock = db_conf |> idle_tenant() |> connect() assert {:error, :closed} = :gen_tcp.recv(sock, 0, @idle_timeout_ms * 5) end @@ -48,59 +50,13 @@ defmodule Supavisor.Integration.ClientIdleTimeoutTest do tenant: tenant_id, port: Application.get_env(:supavisor, :proxy_port_transaction), user: db_conf[:username], - password: db_conf[:password], - database: to_string(db_conf[:database]) + password: db_conf[:password] } end - defp scram_connect(%{tenant: tenant, port: port, user: user, password: password, database: db}) do + defp connect(%{tenant: tenant, port: port, user: user, password: password}) do {:ok, sock} = :gen_tcp.connect(~c"127.0.0.1", port, [:binary, active: false]) - - startup = - :pgo_protocol.encode_startup_message([{"user", "#{user}.#{tenant}"}, {"database", db}]) - - :ok = :gen_tcp.send(sock, startup) - - # SASL auth request - {:ok, <>} = :gen_tcp.recv(sock, 0, 5000) - - # SCRAM client-first - nonce = :pgo_scram.get_nonce(16) - client_first = :pgo_scram.get_client_first(user, nonce) - client_first_size = :erlang.iolist_size(client_first) - sasl_initial = ["SCRAM-SHA-256", 0, <>, client_first] - :ok = :gen_tcp.send(sock, :pgo_protocol.encode_scram_response_message(sasl_initial)) - - # SCRAM server-first - {:ok, <>} = :gen_tcp.recv(sock, 0, 5000) - server_first_parts = :pgo_scram.parse_server_first(server_first, nonce) - - # SCRAM client-final - {client_final, server_proof} = - :pgo_scram.get_client_final(server_first_parts, nonce, user, password) - - :ok = :gen_tcp.send(sock, :pgo_protocol.encode_scram_response_message(client_final)) - - # SCRAM server-final + auth ok + params + ReadyForQuery - {:ok, auth_data} = :gen_tcp.recv(sock, 0, 5000) - - {[<> | _], ""} = - Supavisor.Protocol.split_pkts(auth_data) - - {:ok, ^server_proof} = :pgo_scram.parse_server_final(server_final) - recv_until_ready_for_query(sock, auth_data) - - {:gen_tcp, sock} - end - - defp recv_until_ready_for_query(sock, buf) do - {pkts, ""} = Supavisor.Protocol.split_pkts(buf) - - if Enum.any?(pkts, &match?(<>, &1)) do - :ok - else - {:ok, more} = :gen_tcp.recv(sock, 0, 5000) - recv_until_ready_for_query(sock, more) - end + ProtocolClient.authenticate(sock, "#{user}.#{tenant}", password) + sock end end diff --git a/test/integration/protocol_integration_test.exs b/test/integration/protocol_integration_test.exs index 4f1d2a2d..c89eeefe 100644 --- a/test/integration/protocol_integration_test.exs +++ b/test/integration/protocol_integration_test.exs @@ -60,38 +60,8 @@ defmodule Supavisor.Integration.ProtocolIntegrationTest do {:ok, sock} = :gen_tcp.connect(~c"127.0.0.1", port, [:binary, active: false]) - startup = :pgo_protocol.encode_startup_message([{"user", "#{user}.#{tenant}"}]) - :ok = :gen_tcp.send(sock, startup) - - # SASL auth request - {:ok, <>} = :gen_tcp.recv(sock, 0, 5000) - assert "SCRAM-SHA-256" in :pgo_protocol.decode_strings(methods_bin) - - # SCRAM client-first - nonce = :pgo_scram.get_nonce(16) - client_first = :pgo_scram.get_client_first(user, nonce) - client_first_size = :erlang.iolist_size(client_first) - sasl_initial = ["SCRAM-SHA-256", 0, <>, client_first] - :ok = :gen_tcp.send(sock, :pgo_protocol.encode_scram_response_message(sasl_initial)) - - # SCRAM server-first - {:ok, <>} = :gen_tcp.recv(sock, 0, 5000) - server_first_parts = :pgo_scram.parse_server_first(server_first, nonce) - - # SCRAM client-final - {client_final, server_proof} = - :pgo_scram.get_client_final(server_first_parts, nonce, user, password) - - :ok = :gen_tcp.send(sock, :pgo_protocol.encode_scram_response_message(client_final)) - - # SCRAM server-final + auth ok + params + ReadyForQuery - {:ok, auth_data} = :gen_tcp.recv(sock, 0, 5000) - - {[<> | _], ""} = - Supavisor.Protocol.split_pkts(auth_data) - - {:ok, ^server_proof} = :pgo_scram.parse_server_final(server_final) - ProtocolClient.recv_until_ready_for_query(sock, auth_data) + # authenticate/3 sends a user-only startup without database parameter + ProtocolClient.authenticate(sock, "#{user}.#{tenant}", password) # Verify the connection defaults to the correct database :ok = :gen_tcp.send(sock, :pgo_protocol.encode_query_message("SELECT current_database()"))