From 1520e7ac724d57947f86ece6916e77d84fbf3056 Mon Sep 17 00:00:00 2001 From: Anthonin Bonnefoy Date: Fri, 17 Jan 2025 15:12:48 +0100 Subject: Add PgProto test module to send message on a raw socket Add a new PostgreSQL::Test::PgProto class that provides helper functions to send PostgreSQL protocol message to a socket. This initial version provides functions to send a startup packet, send a simple query, read the content of a response and consume responsed until the socket is closed. This is useful to test edge cases like sessions stuck in a specific socket state. --- src/test/perl/Makefile | 2 + src/test/perl/PostgreSQL/Test/PgProto.pm | 219 +++++++++++++++++++++++ src/test/perl/meson.build | 1 + 3 files changed, 222 insertions(+) create mode 100644 src/test/perl/PostgreSQL/Test/PgProto.pm diff --git a/src/test/perl/Makefile b/src/test/perl/Makefile index d82fb67540e..d8e9fb6b00d 100644 --- a/src/test/perl/Makefile +++ b/src/test/perl/Makefile @@ -25,6 +25,7 @@ install: all installdirs $(INSTALL_DATA) $(srcdir)/PostgreSQL/Test/Kerberos.pm '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/Kerberos.pm' $(INSTALL_DATA) $(srcdir)/PostgreSQL/Test/Cluster.pm '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/Cluster.pm' $(INSTALL_DATA) $(srcdir)/PostgreSQL/Test/BackgroundPsql.pm '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/BackgroundPsql.pm' + $(INSTALL_DATA) $(srcdir)/PostgreSQL/Test/PgProto.pm '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/PgProto.pm' $(INSTALL_DATA) $(srcdir)/PostgreSQL/Test/AdjustUpgrade.pm '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/AdjustUpgrade.pm' $(INSTALL_DATA) $(srcdir)/PostgreSQL/Version.pm '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Version.pm' @@ -35,6 +36,7 @@ uninstall: rm -f '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/Kerberos.pm' rm -f '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/Cluster.pm' rm -f '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/BackgroundPsql.pm' + rm -f '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/PgProto.pm' rm -f '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Test/AdjustUpgrade.pm' rm -f '$(DESTDIR)$(pgxsdir)/$(subdir)/PostgreSQL/Version.pm' diff --git a/src/test/perl/PostgreSQL/Test/PgProto.pm b/src/test/perl/PostgreSQL/Test/PgProto.pm new file mode 100644 index 00000000000..0fc7b46a190 --- /dev/null +++ b/src/test/perl/PostgreSQL/Test/PgProto.pm @@ -0,0 +1,219 @@ + +# Copyright (c) 2025, PostgreSQL Global Development Group + +=pod + +=head1 NAME + +PostgreSQL::Test::PgProto - class for manipulating PG protocol over a raw socket + +=head1 SYNOPSIS + + use PostgreSQL::Test::PgProto; + + my $node = PostgreSQL::Test::PgProto->new('mynode'); + + # Create a data directory with initdb + $node->init(); + + # Start the PostgreSQL server + $node->start(); + + # Get raw socket to the node + my $sock = $node->raw_connect(); + + # Create pgproto from the raw socket + my $pgproto = PostgreSQL::Test::PgProto->new($sock); + + # Send startup packet + my %parameters = ( user => "postgres", database => "postgres", application_name => "app" ); + $pgproto->send_startup_message(\%parameters); + + # Read startup sequence until Ready For Query is reached + $pgproto->read_until_message('Z'); + + # Send a simple query + $pgproto->send_simple_query('SELECT 1'); + +=head1 DESCRIPTION + +PostgreSQL::Test::PgProto contains functionality for sending and reading +PostgreSQL protocol messages over a raw socket. + +=cut + +package PostgreSQL::Test::PgProto; +use Test::More; + +use strict; +use warnings FATAL => 'all'; + +=pod + +=head1 METHODS + +=over + +=item PostgreSQL::Test::PgProto->new($sock) + +Builds a new object of class C. + +=cut + +sub new +{ + my $class = shift; + my ($sock) = @_; + + my $pgproto = { + 'sock' => $sock + }; + bless $pgproto, $class; + return $pgproto; +} + +=pod + +=item $pgproto->send_startup_message + +Send a startup message with the provided parameters. Database parameters need to +be passed as a reference to a hash. The server's response is not consumed. + +=cut + +sub send_startup_message +{ + my $self = shift; + my $parameters = shift; + + # Startup packet contains: + # Packet length: 4 bytes + # Major proto version: 2 bytes + # Minor proto version: 2 bytes + # Multiple Parameters: + # Parameter Name (null terminated string) + # Parameter Value (null terminated string) + # Ending null character + my $pack_template = "Nnn(Z*Z*)" . keys(%{$parameters}) . 'x'; + # Packet length, proto and final null character + my $total_length = 9; + + for(keys %{$parameters}){ + my $key_length = length($_) + 1; + my $value_length = length($parameters->{$_}) + 1; + $total_length += $key_length + $value_length; + } + + my $startup_packet = pack($pack_template, $total_length, 3, 0, %{$parameters}); + $self->{sock}->send($startup_packet); +} + +=pod + +=item $pgproto->send_simple_query + +Send a simple query message to the server. The response is not consumed. + +=cut + +sub send_simple_query +{ + my ($self, $query) = @_; + + # Query message contains: + # Message type 'Q' (1 byte) + # Message length not including message type (4 bytes) + # Null terminated string + my $query_packet = pack("CNZ*", ord('Q'), length($query) + 5, $query); + note "Sending following simple query through raw_tcp: $query"; + $self->{sock}->send($query_packet); +} + +=pod + +=item $pgproto->read_session_pid + +Returns the pid of the session. The session needs to be in a ready for query +state. All results will be consumed and will leave the session in a ready for +query state. + +=cut + +sub read_session_pid +{ + my ($self) = @_; + + $self->send_simple_query("select pg_backend_pid()"); + my $data_row = $self->read_until_message('D'); + # We should have only one field and one column with the + # pid representing the rest of the payload outside of the + # data row header + my ($field_count, $column_length, $pid) = unpack("nNA*", $data_row); + note "raw_tcp has pid $pid"; + # Consume until Ready for query is reached + $self->read_until_message('Z'); + return $pid; +} + +=pod + +=item $pgproto->wait_until_closed + +Block and read all responses until the socket is terminated by the server. + +=cut + +sub wait_until_closed +{ + my ($self) = @_; + my $received = ""; + + while (1) { + $self->{sock}->recv($received, 64*1024); + if ($received eq "") { + # Closed socket was detected + return; + } + } +} + +=pod + +=item $pgproto->read_until_message + +Read all messages from the server until C is found. The message's +payload will be returned. + +=cut + +sub read_until_message +{ + my ($self, $message_type) = @_; + + note "Reading until message of type $message_type is found"; + while (1) + { + my $header = ""; + $self->{sock}->recv($header, 5); + my ($type, $length) = unpack("CN", $header); + my $type_char = chr($type); + note "Reading message of type $type_char and length $length"; + + # Need to remove message's length from the payload's length + $length -= 4; + if ($length < 0) { + diag("read_until_message Unexpected payload length $length"); + return; + } + my $payload = ""; + $self->{sock}->recv($payload, $length); + + if ($type_char eq $message_type) { + # We've found the desired message type + note "Found expected message type $message_type"; + return $payload; + } + } +} + +1; diff --git a/src/test/perl/meson.build b/src/test/perl/meson.build index 58e30f15f9d..04639a11d95 100644 --- a/src/test/perl/meson.build +++ b/src/test/perl/meson.build @@ -13,5 +13,6 @@ install_data( 'PostgreSQL/Test/Kerberos.pm', 'PostgreSQL/Test/Cluster.pm', 'PostgreSQL/Test/BackgroundPsql.pm', + 'PostgreSQL/Test/PgProto.pm', 'PostgreSQL/Test/AdjustUpgrade.pm', install_dir: dir_pgxs / 'src/test/perl/PostgreSQL/Test') -- 2.39.5 (Apple Git-154)