main.go

// Usage: go run main.go MYSQL_HOST[:PORT]
//
// For example:
// $ go run main.go 127.0.0.1
// Address: 127.0.0.1:
//
// Packet:
//   sequence: 0
//   payload_length: 74
//
// Protocol::HandshakeV10:
//   protocol_version: 10
//   server_version: "8.4.11"
//   connection_id: 97093
//   character_set: 255
//   status_flags: 0x0002
//     - SERVER_STATUS_AUTOCOMMIT
//   capabilities: 0xdfffffff
//     - CLIENT_LONG_PASSWORD
//     - CLIENT_FOUND_ROWS
//     - CLIENT_LONG_FLAG
//     - CLIENT_CONNECT_WITH_DB
//     - CLIENT_NO_SCHEMA
//     - CLIENT_COMPRESS
//     - CLIENT_ODBC
//     - CLIENT_LOCAL_FILES
//     - CLIENT_IGNORE_SPACE
//     - CLIENT_PROTOCOL_41
//     - CLIENT_INTERACTIVE
//     - CLIENT_SSL
//     - CLIENT_IGNORE_SIGPIPE
//     - CLIENT_TRANSACTIONS
//     - CLIENT_RESERVED
//     - CLIENT_SECURE_CONNECTION
//     - CLIENT_MULTI_STATEMENTS
//     - CLIENT_MULTI_RESULTS
//     - CLIENT_PS_MULTI_RESULTS
//     - CLIENT_PLUGIN_AUTH
//     - CLIENT_CONNECT_ATTRS
//     - CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA
//     - CLIENT_CAN_HANDLE_EXPIRED_PASSWORDS
//     - CLIENT_SESSION_TRACK
//     - CLIENT_DEPRECATE_EOF
//     - CLIENT_OPTIONAL_RESULTSET_METADATA
//     - CLIENT_ZSTD_COMPRESSION_ALGORITHM
//     - CLIENT_QUERY_ATTRIBUTES
//     - MULTI_FACTOR_AUTHENTICATION
//     - CLIENT_SSL_VERIFY_SERVER_CERT
//     - CLIENT_REMEMBER_OPTIONS
//   auth_plugin_data_length: 21
//   auth_plugin_data_part_1: 48646a68551f0a32
//   auth_plugin_data_part_2: 5a74621c04200a5908633a1f00
//   auth_plugin_name: "caching_sha2_password"

package main

import (
	"bytes"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
	"time"

	"github.com/spf13/pflag"
)

const defaultMySQLPort = "3306"

func main() {
	if err := run(os.Args[1:], os.Stdin, os.Stdout); err != nil {
		if errors.Is(err, pflag.ErrHelp) {
			os.Exit(0)
		}
		fmt.Fprintf(os.Stderr, "dump-initial-handshake: %v\n", err)
		os.Exit(1)
	}
}

func run(args []string, stdin io.Reader, stdout io.Writer) error {
	flags := pflag.NewFlagSet("dump-initial-handshake", pflag.ContinueOnError)
	flags.SetOutput(os.Stderr)

	timeout := flags.Duration("timeout", 3*time.Second, "connect/read timeout")
	base64Input := flags.Bool("base64", false, "read a base64-encoded packet from stdin instead of connecting to HOST[:PORT]")

	if err := flags.Parse(args); err != nil {
		return err
	}
	if *base64Input {
		if flags.NArg() != 0 {
			return fmt.Errorf("usage: go run ./dump-initial-handshake --base64")
		}
		fmt.Println("TIP: You can capture a packet like this: nc -w 1 -q -1 127.0.0.1 3306 | base64")
		return dumpBase64Packet(stdin, stdout)
	}

	if flags.NArg() != 1 {
		return fmt.Errorf("usage: go run ./dump-initial-handshake [--timeout 3s] HOST[:PORT]")
	}
	if *timeout <= 0 {
		return fmt.Errorf("--timeout must be greater than zero")
	}

	addr, err := normalizeAddress(flags.Arg(0))
	if err != nil {
		return err
	}

	dialer := net.Dialer{Timeout: *timeout}
	conn, err := dialer.Dial("tcp", addr)
	if err != nil {
		return fmt.Errorf("%s: dial: %w", addr, err)
	}
	defer conn.Close()

	if err := conn.SetReadDeadline(time.Now().Add(*timeout)); err != nil {
		return fmt.Errorf("%s: set read deadline: %w", addr, err)
	}

	pkt, err := readPacket(conn)
	if err != nil {
		return fmt.Errorf("%s: read initial packet: %w", addr, err)
	}

	return dumpPacket(stdout, addr, pkt)
}

func dumpBase64Packet(stdin io.Reader, stdout io.Writer) error {
	encoded, err := io.ReadAll(stdin)
	if err != nil {
		return fmt.Errorf("read stdin: %w", err)
	}

	decoded, err := base64.StdEncoding.DecodeString(stripWhitespace(string(encoded)))
	if err != nil {
		return fmt.Errorf("decode base64 stdin: %w", err)
	}

	r := bytes.NewReader(decoded)
	pkt, err := readPacket(r)
	if err != nil {
		return fmt.Errorf("read packet from base64 stdin: %w", err)
	}
	if r.Len() != 0 {
		return fmt.Errorf("base64 stdin has %d trailing byte(s) after one packet", r.Len())
	}

	return dumpPacket(stdout, "stdin(base64)", pkt)
}

func dumpPacket(stdout io.Writer, source string, pkt *packet) error {
	handshake, err := parseHandshakePacket(pkt.payload)
	if err != nil {
		printRawPacket(stdout, source, pkt)
		return fmt.Errorf("%s: parse initial handshake: %w", source, err)
	}

	printHandshake(stdout, source, pkt, handshake)
	return nil
}

func stripWhitespace(s string) string {
	return strings.Join(strings.Fields(s), "")
}

func normalizeAddress(addr string) (string, error) {
	addr = strings.TrimSpace(addr)
	if addr == "" {
		return "", fmt.Errorf("address is required")
	}

	if strings.HasPrefix(addr, "[") {
		host, port, err := net.SplitHostPort(addr)
		if err == nil {
			return joinHostPort(host, port)
		}
		if strings.HasSuffix(addr, "]") {
			return joinHostPort(strings.TrimSuffix(strings.TrimPrefix(addr, "["), "]"), defaultMySQLPort)
		}
		return "", fmt.Errorf("invalid address %q: %w", addr, err)
	}

	host, port, err := net.SplitHostPort(addr)
	if err == nil {
		return joinHostPort(host, port)
	}

	if strings.Count(addr, ":") == 0 {
		return joinHostPort(addr, defaultMySQLPort)
	}

	if strings.Count(addr, ":") == 1 {
		host, port, _ := strings.Cut(addr, ":")
		if port == "" {
			port = defaultMySQLPort
		}
		return joinHostPort(host, port)
	}

	if ip := net.ParseIP(addr); ip != nil {
		return joinHostPort(addr, defaultMySQLPort)
	}

	return "", fmt.Errorf("address %q must be HOST[:PORT] or [IPv6][:PORT]", addr)
}

func joinHostPort(host, port string) (string, error) {
	if host == "" {
		return "", fmt.Errorf("host is required")
	}
	if port == "" {
		port = defaultMySQLPort
	}
	return net.JoinHostPort(host, port), nil
}

type packet struct {
	sequence uint8
	payload  []byte
}

func readPacket(r io.Reader) (*packet, error) {
	var header [4]byte
	if _, err := io.ReadFull(r, header[:]); err != nil {
		return nil, fmt.Errorf("packet header: %w", err)
	}

	payloadLen := int(parseInt3(header[:3]))
	payload := make([]byte, payloadLen)
	if _, err := io.ReadFull(r, payload); err != nil {
		return nil, fmt.Errorf("packet payload: %w", err)
	}

	return &packet{
		sequence: header[3],
		payload:  payload,
	}, nil
}

type protocolHandshakeV10 struct {
	protocolVersion     uint8
	serverVersion       string
	threadID            uint32
	authPluginDataPart1 []byte
	capabilities        uint32
	characterSet        uint8
	statusFlags         uint16
	authPluginDataLen   uint8
	authPluginDataPart2 []byte
	authPluginName      string
	trailingBytes       []byte
}

func parseHandshakePacket(payload []byte) (*protocolHandshakeV10, error) {
	c := cursor{buf: payload}

	protocolVersion, err := c.int1("protocol version")
	if err != nil {
		return nil, err
	}
	if protocolVersion != 10 {
		return nil, fmt.Errorf("want protocol version 10, got %d", protocolVersion)
	}

	serverVersion, err := c.nullString("server version")
	if err != nil {
		return nil, err
	}

	threadID, err := c.int4("connection id")
	if err != nil {
		return nil, err
	}

	authPluginDataPart1, err := c.fixedString("auth plugin data part 1", 8)
	if err != nil {
		return nil, err
	}

	filler, err := c.int1("filler")
	if err != nil {
		return nil, err
	}
	if filler != 0 {
		return nil, fmt.Errorf("missing filler byte after auth plugin data part 1, got 0x%02x", filler)
	}

	capabilityFlags1, err := c.int2("lower capability flags")
	if err != nil {
		return nil, err
	}

	characterSet, err := c.int1("character set")
	if err != nil {
		return nil, err
	}

	statusFlags, err := c.int2("status flags")
	if err != nil {
		return nil, err
	}

	capabilityFlags2, err := c.int2("upper capability flags")
	if err != nil {
		return nil, err
	}
	capabilities := uint32(capabilityFlags1) | uint32(capabilityFlags2)<<16

	authPluginDataLen, err := c.int1("auth plugin data length")
	if err != nil {
		return nil, err
	}

	reserved, err := c.fixedString("reserved bytes", 10)
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(reserved, make([]byte, 10)) {
		return nil, fmt.Errorf("reserved bytes are not all zero: %x", reserved)
	}

	part2Len := max(13, int(authPluginDataLen)-8)
	authPluginDataPart2, err := c.fixedString("auth plugin data part 2", part2Len)
	if err != nil {
		return nil, err
	}

	var authPluginName []byte
	if capabilities&clientPluginAuth != 0 && c.remaining() > 0 {
		authPluginName, err = c.nullString("auth plugin name")
		if err != nil {
			return nil, err
		}
	}

	return &protocolHandshakeV10{
		protocolVersion:     protocolVersion,
		serverVersion:       string(serverVersion),
		threadID:            threadID,
		authPluginDataPart1: authPluginDataPart1,
		capabilities:        capabilities,
		characterSet:        characterSet,
		statusFlags:         statusFlags,
		authPluginDataLen:   authPluginDataLen,
		authPluginDataPart2: authPluginDataPart2,
		authPluginName:      string(authPluginName),
		trailingBytes:       c.rest(),
	}, nil
}

type cursor struct {
	buf []byte
}

func (c *cursor) remaining() int {
	return len(c.buf)
}

func (c *cursor) rest() []byte {
	rest := c.buf
	c.buf = nil
	return rest
}

func (c *cursor) need(name string, n int) error {
	if len(c.buf) < n {
		return fmt.Errorf("%s: need %d byte(s), have %d", name, n, len(c.buf))
	}
	return nil
}

func (c *cursor) int1(name string) (uint8, error) {
	if err := c.need(name, 1); err != nil {
		return 0, err
	}
	value := c.buf[0]
	c.buf = c.buf[1:]
	return value, nil
}

func (c *cursor) int2(name string) (uint16, error) {
	if err := c.need(name, 2); err != nil {
		return 0, err
	}
	value := parseInt2(c.buf[:2])
	c.buf = c.buf[2:]
	return value, nil
}

func (c *cursor) int4(name string) (uint32, error) {
	if err := c.need(name, 4); err != nil {
		return 0, err
	}
	value := parseInt4(c.buf[:4])
	c.buf = c.buf[4:]
	return value, nil
}

func (c *cursor) fixedString(name string, n int) ([]byte, error) {
	if err := c.need(name, n); err != nil {
		return nil, err
	}
	value := c.buf[:n]
	c.buf = c.buf[n:]
	return value, nil
}

func (c *cursor) nullString(name string) ([]byte, error) {
	i := bytes.IndexByte(c.buf, 0)
	if i < 0 {
		return nil, fmt.Errorf("%s: missing NUL terminator", name)
	}
	value := c.buf[:i]
	c.buf = c.buf[i+1:]
	return value, nil
}

func printRawPacket(w io.Writer, addr string, pkt *packet) {
	fmt.Fprintf(w, "Address: %s\n\n", addr)
	fmt.Fprintf(w, "Packet:\n")
	fmt.Fprintf(w, "  sequence: %d\n", pkt.sequence)
	fmt.Fprintf(w, "  payload_length: %d\n", len(pkt.payload))
	fmt.Fprintf(w, "  payload_hex:\n%s", indent(hex.Dump(pkt.payload), "    "))
}

func printHandshake(w io.Writer, addr string, pkt *packet, h *protocolHandshakeV10) {
	fmt.Fprintf(w, "Address: %s\n\n", addr)
	fmt.Fprintf(w, "Packet:\n")
	fmt.Fprintf(w, "  sequence: %d\n", pkt.sequence)
	fmt.Fprintf(w, "  payload_length: %d\n\n", len(pkt.payload))

	fmt.Fprintf(w, "Protocol::HandshakeV10:\n")
	fmt.Fprintf(w, "  protocol_version: %d\n", h.protocolVersion)
	fmt.Fprintf(w, "  server_version: %q\n", h.serverVersion)
	fmt.Fprintf(w, "  connection_id: %d\n", h.threadID)
	fmt.Fprintf(w, "  character_set: %d\n", h.characterSet)
	fmt.Fprintf(w, "  status_flags: 0x%04x\n", h.statusFlags)
	printFlags(w, "    ", uint64(h.statusFlags), 16, statusFlagNames)
	fmt.Fprintf(w, "  capabilities: 0x%08x\n", h.capabilities)
	printFlags(w, "    ", uint64(h.capabilities), 32, capabilityNames)
	fmt.Fprintf(w, "  auth_plugin_data_length: %d\n", h.authPluginDataLen)
	fmt.Fprintf(w, "  auth_plugin_data_part_1: %s\n", hex.EncodeToString(h.authPluginDataPart1))
	fmt.Fprintf(w, "  auth_plugin_data_part_2: %s\n", hex.EncodeToString(h.authPluginDataPart2))
	fmt.Fprintf(w, "  auth_plugin_name: %q\n", h.authPluginName)
	if len(h.trailingBytes) > 0 {
		fmt.Fprintf(w, "  trailing_bytes: %s\n", hex.EncodeToString(h.trailingBytes))
	}
}

func printFlags(w io.Writer, indentText string, value uint64, width uint, names []flagName) {
	if value == 0 {
		fmt.Fprintf(w, "%s- (none)\n", indentText)
		return
	}

	remaining := value
	for _, name := range names {
		if value&name.flag != 0 {
			fmt.Fprintf(w, "%s- %s\n", indentText, name.name)
			remaining &^= name.flag
		}
	}

	for bit := uint(0); bit < width; bit++ {
		flag := uint64(1) << bit
		if remaining&flag != 0 {
			if width == 16 {
				fmt.Fprintf(w, "%s- 0x%04x\n", indentText, flag)
			} else {
				fmt.Fprintf(w, "%s- 0x%08x\n", indentText, flag)
			}
		}
	}
}

func indent(s, prefix string) string {
	if s == "" {
		return ""
	}
	lines := strings.SplitAfter(s, "\n")
	var b strings.Builder
	for _, line := range lines {
		if line == "" {
			continue
		}
		b.WriteString(prefix)
		b.WriteString(line)
	}
	return b.String()
}

func parseInt2(bytes []byte) uint16 {
	return uint16(bytes[0]) + uint16(bytes[1])<<8
}

func parseInt3(bytes []byte) uint32 {
	return uint32(bytes[0]) + uint32(bytes[1])<<8 + uint32(bytes[2])<<16
}

func parseInt4(bytes []byte) uint32 {
	return uint32(bytes[0]) + uint32(bytes[1])<<8 + uint32(bytes[2])<<16 + uint32(bytes[3])<<24
}

type flagName struct {
	flag uint64
	name string
}

const (
	clientLongPassword               = 1 << 0
	clientFoundRows                  = 1 << 1
	clientLongFlag                   = 1 << 2
	clientConnectWithDB              = 1 << 3
	clientNoSchema                   = 1 << 4
	clientCompress                   = 1 << 5
	clientODBC                       = 1 << 6
	clientLocalFiles                 = 1 << 7
	clientIgnoreSpace                = 1 << 8
	clientProtocol41                 = 1 << 9
	clientInteractive                = 1 << 10
	clientSSL                        = 1 << 11
	clientIgnoreSigpipe              = 1 << 12
	clientTransactions               = 1 << 13
	clientReserved                   = 1 << 14
	clientSecureConnection           = 1 << 15
	clientMultiStatements            = 1 << 16
	clientMultiResults               = 1 << 17
	clientPSMultiResults             = 1 << 18
	clientPluginAuth                 = 1 << 19
	clientConnectAttrs               = 1 << 20
	clientPluginAuthLenencClientData = 1 << 21
	clientCanHandleExpiredPasswords  = 1 << 22
	clientSessionTrack               = 1 << 23
	clientDeprecateEOF               = 1 << 24
	clientOptionalResultsetMetadata  = 1 << 25
	clientZstdCompressionAlgorithm   = 1 << 26
	clientQueryAttributes            = 1 << 27
	clientMultiFactorAuthentication  = 1 << 28
	clientCapabilityExtension        = 1 << 29
	clientSSLVerifyServerCert        = 1 << 30
	clientRememberOptions            = 1 << 31
)

var capabilityNames = []flagName{
	{clientLongPassword, "CLIENT_LONG_PASSWORD"},
	{clientFoundRows, "CLIENT_FOUND_ROWS"},
	{clientLongFlag, "CLIENT_LONG_FLAG"},
	{clientConnectWithDB, "CLIENT_CONNECT_WITH_DB"},
	{clientNoSchema, "CLIENT_NO_SCHEMA"},
	{clientCompress, "CLIENT_COMPRESS"},
	{clientODBC, "CLIENT_ODBC"},
	{clientLocalFiles, "CLIENT_LOCAL_FILES"},
	{clientIgnoreSpace, "CLIENT_IGNORE_SPACE"},
	{clientProtocol41, "CLIENT_PROTOCOL_41"},
	{clientInteractive, "CLIENT_INTERACTIVE"},
	{clientSSL, "CLIENT_SSL"},
	{clientIgnoreSigpipe, "CLIENT_IGNORE_SIGPIPE"},
	{clientTransactions, "CLIENT_TRANSACTIONS"},
	{clientReserved, "CLIENT_RESERVED"},
	{clientSecureConnection, "CLIENT_SECURE_CONNECTION"},
	{clientMultiStatements, "CLIENT_MULTI_STATEMENTS"},
	{clientMultiResults, "CLIENT_MULTI_RESULTS"},
	{clientPSMultiResults, "CLIENT_PS_MULTI_RESULTS"},
	{clientPluginAuth, "CLIENT_PLUGIN_AUTH"},
	{clientConnectAttrs, "CLIENT_CONNECT_ATTRS"},
	{clientPluginAuthLenencClientData, "CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA"},
	{clientCanHandleExpiredPasswords, "CLIENT_CAN_HANDLE_EXPIRED_PASSWORDS"},
	{clientSessionTrack, "CLIENT_SESSION_TRACK"},
	{clientDeprecateEOF, "CLIENT_DEPRECATE_EOF"},
	{clientOptionalResultsetMetadata, "CLIENT_OPTIONAL_RESULTSET_METADATA"},
	{clientZstdCompressionAlgorithm, "CLIENT_ZSTD_COMPRESSION_ALGORITHM"},
	{clientQueryAttributes, "CLIENT_QUERY_ATTRIBUTES"},
	{clientMultiFactorAuthentication, "MULTI_FACTOR_AUTHENTICATION"},
	{clientCapabilityExtension, "CLIENT_CAPABILITY_EXTENSION"},
	{clientSSLVerifyServerCert, "CLIENT_SSL_VERIFY_SERVER_CERT"},
	{clientRememberOptions, "CLIENT_REMEMBER_OPTIONS"},
}

const (
	serverStatusInTrans            = 1 << 0
	serverStatusAutocommit         = 1 << 1
	serverMoreResultsExists        = 1 << 3
	serverQueryNoGoodIndexUsed     = 1 << 4
	serverQueryNoIndexUsed         = 1 << 5
	serverStatusCursorExists       = 1 << 6
	serverStatusLastRowSent        = 1 << 7
	serverStatusDBDropped          = 1 << 8
	serverStatusNoBackslashEscapes = 1 << 9
	serverStatusMetadataChanged    = 1 << 10
	serverQueryWasSlow             = 1 << 11
	serverPSOutParams              = 1 << 12
	serverStatusInTransReadonly    = 1 << 13
	serverSessionStateChanged      = 1 << 14
)

var statusFlagNames = []flagName{
	{serverStatusInTrans, "SERVER_STATUS_IN_TRANS"},
	{serverStatusAutocommit, "SERVER_STATUS_AUTOCOMMIT"},
	{serverMoreResultsExists, "SERVER_MORE_RESULTS_EXISTS"},
	{serverQueryNoGoodIndexUsed, "SERVER_QUERY_NO_GOOD_INDEX_USED"},
	{serverQueryNoIndexUsed, "SERVER_QUERY_NO_INDEX_USED"},
	{serverStatusCursorExists, "SERVER_STATUS_CURSOR_EXISTS"},
	{serverStatusLastRowSent, "SERVER_STATUS_LAST_ROW_SENT"},
	{serverStatusDBDropped, "SERVER_STATUS_DB_DROPPED"},
	{serverStatusNoBackslashEscapes, "SERVER_STATUS_NO_BACKSLASH_ESCAPES"},
	{serverStatusMetadataChanged, "SERVER_STATUS_METADATA_CHANGED"},
	{serverQueryWasSlow, "SERVER_QUERY_WAS_SLOW"},
	{serverPSOutParams, "SERVER_PS_OUT_PARAMS"},
	{serverStatusInTransReadonly, "SERVER_STATUS_IN_TRANS_READONLY"},
	{serverSessionStateChanged, "SERVER_SESSION_STATE_CHANGED"},
}