LinkPort/internal/protocol/mysqlproxy/compression_test.go

package mysqlproxy

import (
	"bytes"
	"compress/zlib"
	"context"
	"encoding/binary"
	"github.com/go-mysql-org/go-mysql/mysql"
	"github.com/go-mysql-org/go-mysql/server"
	"io"
	"linkport/internal/model"
	"net"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// A test server transport implementing the compressed framing independently of
// go-mysql's client compressor (the server library itself has no compression).
type compressedBackendConn struct {
	net.Conn
	active, greeted bool
	pending         *bytes.Reader
	sequence        byte
}

func (c *compressedBackendConn) Read(p []byte) (int, error) {
	if !c.active {
		return c.Conn.Read(p)
	}
	if c.pending != nil && c.pending.Len() > 0 {
		return c.pending.Read(p)
	}
	var h [7]byte
	if _, err := io.ReadFull(c.Conn, h[:]); err != nil {
		return 0, err
	}
	n := int(h[0]) | int(h[1])<<8 | int(h[2])<<16
	raw := make([]byte, n)
	if _, err := io.ReadFull(c.Conn, raw); err != nil {
		return 0, err
	}
	c.sequence = h[3] + 1
	if h[4] != 0 || h[5] != 0 || h[6] != 0 {
		r, err := zlib.NewReader(bytes.NewReader(raw))
		if err != nil {
			return 0, err
		}
		raw, err = io.ReadAll(r)
		r.Close()
		if err != nil {
			return 0, err
		}
	}
	c.pending = bytes.NewReader(raw)
	return c.pending.Read(p)
}
func (c *compressedBackendConn) Write(p []byte) (int, error) {
	if !c.greeted {
		c.greeted = true
		p = bytes.Clone(p)
		end := bytes.IndexByte(p[5:], 0)
		offset := 5 + end + 1 + 4 + 8 + 1
		caps := binary.LittleEndian.Uint16(p[offset:])
		binary.LittleEndian.PutUint16(p[offset:], caps|uint16(mysql.CLIENT_COMPRESS))
	}
	if !c.active {
		return c.Conn.Write(p)
	}
	var out bytes.Buffer
	w := zlib.NewWriter(&out)
	w.Write(p)
	w.Close()
	compressed := out.Bytes()
	h := []byte{byte(len(compressed)), byte(len(compressed) >> 8), byte(len(compressed) >> 16), c.sequence, byte(len(p)), byte(len(p) >> 8), byte(len(p) >> 16)}
	c.sequence++
	if _, err := c.Conn.Write(append(h, compressed...)); err != nil {
		return 0, err
	}
	return len(p), nil
}

type trafficConn struct {
	net.Conn
	written *atomic.Int64
}

func (c *trafficConn) Write(p []byte) (int, error) {
	n, e := c.Conn.Write(p)
	c.written.Add(int64(n))
	return n, e
}

func TestMySQLCompressionNegotiationTrafficAndFallback(t *testing.T) {
	for _, supported := range []bool{false, true} {
		t.Run(map[bool]string{false: "fallback", true: "zlib"}[supported], func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			var queries atomic.Int32
			var written atomic.Int64
			dial := func(context.Context, string, string) (net.Conn, error) {
				a, b := net.Pipe()
				go func() {
					defer b.Close()
					stop := context.AfterFunc(ctx, func() { b.Close() })
					defer stop()
					var raw net.Conn = b
					var compressed *compressedBackendConn
					if supported {
						compressed = &compressedBackendConn{Conn: b}
						raw = compressed
					}
					srv := server.NewServer("5.7.44-test", 45, mysql.AUTH_NATIVE_PASSWORD, nil, nil)
					c, err := newNativeConn(srv, raw, "target", "secret", &fakeBackend{queries: &queries})
					if err != nil {
						return
					}
					if supported {
						compressed.active = true
					}
					for !c.Closed() {
						if c.HandleCommand() != nil {
							return
						}
					}
				}()
				return &trafficConn{a, &written}, nil
			}
			p := model.Proxy{TargetHost: "db.invalid", TargetPort: 3306, TargetUser: "target", TargetPassword: "secret", TargetCompression: "zlib"}
			c, err := open(ctx, p, dial, "")
			if err != nil {
				t.Fatal(err)
			}
			defer c.Close()
			expected := uint8(mysql.MYSQL_COMPRESS_NONE)
			if supported {
				expected = mysql.MYSQL_COMPRESS_ZLIB
			}
			if c.Compression != expected {
				t.Fatalf("negotiated %d, want %d", c.Compression, expected)
			}
			query := "SELECT 42 /*" + strings.Repeat("compressible payload ", 4000) + "*/"
			before := written.Load()
			for i := 0; i < 3; i++ {
				result, err := c.Execute(query)
				if err != nil {
					t.Fatal(err)
				}
				v, err := result.GetInt(0, 0)
				if err != nil || v != 42 {
					t.Fatal("response changed")
				}
				result.Close()
			}
			if supported && written.Load()-before >= int64(len(query)) {
				t.Fatal("compression did not reduce upstream bytes")
			}
			statement, err := c.Prepare("SELECT ?")
			if err != nil {
				t.Fatal(err)
			}
			result, err := statement.Execute(42)
			if err != nil {
				t.Fatal(err)
			}
			result.Close()
			statement.Close()
		})
	}
}