From 5417986f171e67ed8f735a92bb3eb44f0e88149b Mon Sep 17 00:00:00 2001
From: eileencodes <eileencodes@gmail.com>
Date: Tue, 4 Oct 2022 16:33:45 -0400
Subject: [PATCH] Implement caching_sha2_password for mysql 8+

This PR implements caching_sha2_password for mysql 8. Note that we have
chosen on purpose to only implement the path where TSL is used. We will
not be implementing the non-TSL path.

Co-authored-by: John Hawthorn <john@hawthorn.email>
Aaron Patterson (tenderlove) <tenderlove@ruby-lang.org>
---
 inc/trilogy/protocol.h |  3 +++
 src/client.c           | 54 ++++++++++++++++++++++++++++++++++++++++++
 src/protocol.c         | 11 ++++++++-
 3 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/inc/trilogy/protocol.h b/inc/trilogy/protocol.h
index 57a38a03..5b66af87 100644
--- a/inc/trilogy/protocol.h
+++ b/inc/trilogy/protocol.h
@@ -367,6 +367,7 @@ typedef enum {
 // Typical response packet types
 typedef enum {
     TRILOGY_PACKET_OK = 0x0,
+    TRILOGY_PACKET_AUTH_MORE_DATA = 0x01,
     TRILOGY_PACKET_EOF = 0xfe,
     TRILOGY_PACKET_ERR = 0xff,
     TRILOGY_PACKET_UNKNOWN
@@ -1035,4 +1036,6 @@ int trilogy_parse_stmt_ok_packet(const uint8_t *buff, size_t len, trilogy_stmt_o
 int trilogy_parse_stmt_row_packet(const uint8_t *buff, size_t len, trilogy_column_packet_t *columns,
                                   uint64_t column_count, trilogy_binary_value_t *out_values);
 
+int trilogy_build_auth_clear_password(trilogy_builder_t *builder, const char *pass, size_t pass_len);
+
 #endif
diff --git a/src/client.c b/src/client.c
index be467d5d..3df8bf5b 100644
--- a/src/client.c
+++ b/src/client.c
@@ -403,6 +403,9 @@ void trilogy_auth_clear_password(trilogy_conn_t *conn)
     }
 }
 
+#define FAST_AUTH_OK 3
+#define FAST_AUTH_FAIL 4
+
 int trilogy_auth_recv(trilogy_conn_t *conn, trilogy_handshake_t *handshake)
 {
     int rc = read_packet(conn);
@@ -416,6 +419,57 @@ int trilogy_auth_recv(trilogy_conn_t *conn, trilogy_handshake_t *handshake)
         trilogy_auth_clear_password(conn);
         return read_ok_packet(conn);
 
+    case TRILOGY_PACKET_AUTH_MORE_DATA: {
+        uint8_t byte = conn->packet_buffer.buff[1];
+        switch (byte) {
+            case FAST_AUTH_OK:
+                break;
+            case FAST_AUTH_FAIL:
+                {
+                    trilogy_builder_t builder;
+                    int err = begin_command_phase(&builder, conn, conn->packet_parser.sequence_number);
+
+                    if (err < 0) {
+                        return err;
+                    }
+
+                    err = trilogy_build_auth_clear_password(&builder, conn->socket->opts.password, conn->socket->opts.password_len);
+
+                    if (err < 0) {
+                        return err;
+                    }
+
+                    int rc = begin_write(conn);
+
+                    while (rc == TRILOGY_AGAIN) {
+                        rc = trilogy_sock_wait_write(conn->socket);
+                        if (rc != TRILOGY_OK) {
+                            return rc;
+                        }
+
+                        rc = trilogy_flush_writes(conn);
+                    }
+
+                    break;
+                }
+            default:
+                return TRILOGY_UNEXPECTED_PACKET;
+        }
+        while (1) {
+            rc = read_packet(conn);
+
+            if (rc == TRILOGY_OK) {
+                return rc;
+            }
+            else if (rc == TRILOGY_AGAIN) {
+                rc = trilogy_sock_wait_read(conn->socket);
+            }
+
+            if (rc != TRILOGY_OK) {
+                return rc;
+            }
+        }
+    }
     case TRILOGY_PACKET_ERR:
         trilogy_auth_clear_password(conn);
         return read_err_packet(conn);
diff --git a/src/protocol.c b/src/protocol.c
index 232bb97f..49fe5c64 100644
--- a/src/protocol.c
+++ b/src/protocol.c
@@ -593,7 +593,16 @@ int trilogy_build_auth_packet(trilogy_builder_t *builder, const char *user, cons
 
     trilogy_builder_finalize(builder);
 
-    return TRILOGY_OK;
+fail:
+    return rc;
+}
+
+int trilogy_build_auth_clear_password(trilogy_builder_t *builder, const char *pass, size_t pass_len) {
+    int rc = TRILOGY_OK;
+
+    CHECKED(trilogy_builder_write_buffer(builder, pass, pass_len));
+    CHECKED(trilogy_builder_write_uint8(builder, 0));
+    trilogy_builder_finalize(builder);
 
 fail:
     return rc;