Write a Postgres proxy. Day 1.

RediSQL, SQL steroids for Redis. Is a very fast in-memory SQL engine. Its main features are:

1. Speed, up to 130,000 insert per second
2. Familiarity, it support standard SQL, no weird dialects
3. Simplicity, it is very easy to operate and to use with binding for any language.

Code on github: RedBeardLab/rediSQL

In this series of post we are writing a postgres proxy that accept connections made using the postgres (PG) protocol and forward them to RediSQL.

Motivation and introduction of the project are here.

My hope for this project is to distill the knowledge I am getting from this work and help other that are interested in exploring the PG protocol.

Intro

This post is about the first day of this project so it is mostly introduction of the references used during this work and a little bit of code.

In this day we quickly reach the stage where we are able to receive a query from psql (the CLI tool for PG).

We will start the post showing the PG references that are most useful and the Python references of the asyncio module we used.

Then we will explore very quickly the few lines of code that I ended up writing.

The last section will explore the error I made during this day, trivial errors but that where a big time sink anyway.

The references

Before to start this work I questioned if I should implement this proxy for Postgres or for MySQL.

To choose I explored at the documentation of both projects and both are quite good. However, the documentation for PG looked simpler to follow and more linear and I just decide to go for PG.

The main documentation for this project is the Chapter 50 of the PG documentation.

In particular the following sections are of extreme interest:

• The Message Flow section explains what is the flow of messages between PG and the client. It helps in understanding what message we should expect from the client and what message we are required to send as a server.
• The Message Data Types section simply explains how to read and interpretate the section “Message Formats”.
• The Message Formats section goes into the details and enumerate the format of each kind of message. As an example we discover that usually each message start with a single letter that identify the type of message (like R is used for authentication related messages or that queries start with Q), then 4 bytes (an Int32) indicate the length of the whole message and finally the body of the message itself.

Python “ASYNCIO”

While I would like to merge this project in the main RediSQL rust codebase, I am a strong believer that starting the project in Python is a good idea. I will gain the knowledge necessary to successfully re-write the software in Rust while having already faced most of the implementation difficulties in a language that allows very fast iteration. Moreover, it will be just impossible to merge the Python code base into the RediSQL rust codebase, so I will just be forced to re-write it.

While I am not looking for performance I still opted to work with asyncio, mostly because it was a long time I didn’t do any big work in Python and I wanted to get a pulse of the available tools. Moreover I hoped that it would be closer to what I would find in Rust with Tokio, but it seems to me that the two models are not very similar.

On the Python side I keep referring to the Callback Based API for [asyncio](https://docs.python.org/3.5/library/asyncio-protocol.html#transports-and-protocols-callback-based-api).

The API is very simple, you simply sub-classed the asyncio.Protocol class and implemented three callbacks:

1. connection_made for when a new connection is created to the server.
2. data_received for when a new packed of data arrives to the server.
3. connection_lost for when we loose connection with the client.

As you can imagine all the logic is in the data_receivedcallback, and it will be more complex than a standard web-server. Indeed HTTP is a stateless protocol, everything is simpler if the protocol is stateless, each request does not depends on the previous one.

The PG protocol is stateful, it means that we need to store and use information from previous messages. As an example, a client, before to send its queries, needs to send an handshake and to authenticate. This means that our server will have at least two state, an “initial” state where each connection start and a “ready” state where a connection end ups only if it completed the handshake and authenticate.

The code

Finally here the code of this first day of code. The code is mostly boilerplate copied from the Python documentation but it is already enough to accept a connection from psql and receive the first query.

To test our progresses we started the Python server and, at the same time we execute psql giving as input a file with few SQL statements to execute.

psql -f goal.sql -h localhost -p 8888

The workflow of the day

Other than boilerplate code, the interesting part of the code are the definition of the magic number that identifies the messages:

SSLRequestCode = b'\x04\xd2\x16\x2f' # == hex(80877103)
StartupMessageCode = b'\x00\x03\x00\x00' # == hex(196608)

NoSSL = b'\x4E' # == 'N'

AuthenticationOk = b'\x52\x00\x00\x00\x08\x00\x00\x00\x00'
AuthenticationCleartextPassword = b'\x52\x00\x00\x00\x08\x00\x00\x00\x03'

ReadyForQuery = b'\x5A\x00\x00\x00\x05\x49' # == Z0005I , the last I stand for Idle 

And the logic to reply to the client:

    def _reply(self, data):
        if self.state == "initial" and data[4:8] == SSLRequestCode:
            self.transport.write(NoSSL)
        elif self.state == "initial" and data[4:8] == StartupMessageCode:
            # we don't require a password
            self.transport.write(AuthenticationOk)
            # good to go for the first query!
            self.transport.write(ReadyForQuery)
        return

Let’s explore how we get to this few lines of code.

My discovering process

Scanning quickly the documentation it could seems like the first message to expect is the StartupMessage, however, the first message sent by psql is the SSLRequest message, and this took quite a while to figure out.

The SSLRequest message is recognized because it contains the magic number 80877103 which we encode in the python code as b'\x04\xd2\x16\x2f'

Since we don’t yet support SSL we simply respond to the SSLRequest with N(encoded as b'\x4E') to let know to the client that we are not going to use SSL. At this point, the client, can either drop the connection or decide to accept a non-encrypted connection and send the StartupMessage in plain text.

Also for the StartupMessage there is a magic number (196608) which we encoded as b'\x00\x03\x00\x00'.

Along with the magic number, the StartupMessage contains information like the user who is starting the connection, what database the user is trying to connect and other information. At the moment we ignore all those information.

After the StartupMessage the server requires authentication, in our case we don’t care about authentication just yet and we just send the AuthenticationOk message.

The next step is a little tricky.

We just send a message to the client, the AuthenticationOk message and so I would expect the client to send the server something back.

Wrong!

Now, the server need to be proactive and tell the client that it is ok to start sending queries. We need to send two messages, one after the other to the client.

Indeed you can see in the code that we immediately send the ReadyForQuery message.

At this point our time is over for this day, however we can clearly see from the log that the next message received by the server is the first query of our file!

Success!

Errors made during this day

During this coding section I wasted a lot of time because I didn’t read the documentation with enough care.

Indeed I was expecting the StartupMessage as first message and not the SSLRequest. I spend a lot of time trying to fit the StartupMessage into the SSLRequest, maybe I was reading the message with the wrong endianess? Maybe there was “garbage” from the protocol layer?

Nah! I am just reading the wrong message.

Another time sink was me reading the wrong column in the ASCII table trying to use the decimal, instead of the hexadecimal, encoding. All the messages start with a letter, in our cases we needed the N for rejecting the SSL, and the R for the AuthenticationOk message and finally the Z for the ReadyForQuery message. As an example the N is 78 in decimal and 4E in hexadecimal. I was trying to encode N as b'\x78' instead of b'\x4E'.

Conclusion

I hope your enjoyed the post.

I will keep publishing about this topic on this blog, so if you are interested feel free to follow me on twitter or subscribe to the mail list just below.

All post of this serie: Writing a Postgres proxy.