LinkPort/internal/protocol/sqlserverproxy/wire.go

package sqlserverproxy

import (
	"encoding/binary"
	"errors"
	"io"
	"net"
	"serviceport/internal/protocol/transport"
	"unicode/utf16"
)

const maxMessage = 16 << 20

func readMessage(r io.Reader) (byte, []byte, error) {
	return readWithFlags(r, nil)
}
func readWithFlags(r io.Reader, flags *byte) (byte, []byte, error) {
	var result []byte
	var kind byte
	packets := 0
	for {
		var header [8]byte
		if _, err := io.ReadFull(r, header[:]); err != nil {
			return 0, nil, err
		}
		packets++
		if packets > 65536 {
			return 0, nil, errors.New("too many TDS fragments")
		}
		if flags != nil && packets == 1 {
			*flags = header[1] & 0x18
		}
		n := int(binary.BigEndian.Uint16(header[2:])) - 8
		if n < 0 || n > 32759 || len(result)+n > maxMessage || (len(result) > 0 && header[0] != kind) {
			return 0, nil, errors.New("invalid TDS message")
		}
		kind = header[0]
		start := len(result)
		result = append(result, make([]byte, n)...)
		if _, err := io.ReadFull(r, result[start:]); err != nil {
			return 0, nil, err
		}
		if header[1]&1 != 0 {
			return kind, result, nil
		}
	}
}
func writeMessage(w io.Writer, kind byte, data []byte) error {
	return writeWithFlags(w, kind, data, 0)
}
func writeWithFlags(w io.Writer, kind byte, data []byte, flags byte) error {
	for id := byte(1); ; id++ {
		n := min(len(data), 4088)
		header := []byte{kind, 0, 0, 0, 0, 0, id, 0}
		if id == 1 {
			header[1] = flags
		}
		if n == len(data) {
			header[1] |= 1
		}
		binary.BigEndian.PutUint16(header[2:], uint16(n+8))
		if err := transport.Write(w, append(header, data[:n]...)); err != nil {
			return err
		}
		data = data[n:]
		if len(data) == 0 {
			return nil
		}
	}
}
func prelogin(encryption byte) []byte {
	return []byte{0, 0, 16, 0, 6, 1, 0, 22, 0, 1, 4, 0, 23, 0, 1, 255, 15, 0, 0, 0, 0, 0, encryption, 0}
}
func encryptionOption(data []byte) (byte, error) {
	for pos := 0; pos < len(data); pos += 5 {
		if data[pos] == 255 {
			break
		}
		if pos+5 > len(data) {
			break
		}
		off, n := int(binary.BigEndian.Uint16(data[pos+1:])), int(binary.BigEndian.Uint16(data[pos+3:]))
		if off+n > len(data) {
			break
		}
		if data[pos] == 1 && n == 1 {
			return data[off], nil
		}
	}
	return 0, errors.New("invalid TDS prelogin")
}

type handshakeConn struct {
	net.Conn
	framed  bool
	pending []byte
}

func (c *handshakeConn) Read(b []byte) (int, error) {
	if !c.framed && len(c.pending) == 0 {
		return c.Conn.Read(b)
	}
	if len(c.pending) == 0 {
		kind, data, err := readMessage(c.Conn)
		if err != nil {
			return 0, err
		}
		if kind != 18 && kind != 4 {
			return 0, errors.New("invalid TDS TLS frame")
		}
		c.pending = data
	}
	n := copy(b, c.pending)
	c.pending = c.pending[n:]
	return n, nil
}
func (c *handshakeConn) Write(b []byte) (int, error) {
	if !c.framed {
		return c.Conn.Write(b)
	}
	if err := writeMessage(c.Conn, 18, b); err != nil {
		return 0, err
	}
	return len(b), nil
}
func textBytes(s string) []byte {
	units := utf16.Encode([]rune(s))
	out := make([]byte, len(units)*2)
	for i, v := range units {
		binary.LittleEndian.PutUint16(out[i*2:], v)
	}
	return out
}
func decodeText(b []byte) string {
	units := make([]uint16, len(b)/2)
	for i := range units {
		units[i] = binary.LittleEndian.Uint16(b[i*2:])
	}
	return string(utf16.Decode(units))
}
func loginField(data []byte, pos int, password bool) (string, error) {
	if len(data) < 94 || pos+4 > len(data) {
		return "", errors.New("short TDS login")
	}
	off, n := int(binary.LittleEndian.Uint16(data[pos:])), int(binary.LittleEndian.Uint16(data[pos+2:]))*2
	if off < 94 && n != 0 || off+n > len(data) {
		return "", errors.New("invalid TDS login field")
	}
	b := append([]byte(nil), data[off:off+n]...)
	if password {
		for i, v := range b {
			v ^= 0xa5
			b[i] = v<<4 | v>>4
		}
	}
	return decodeText(b), nil
}
func upstreamLogin(original []byte, user, password, host string) ([]byte, error) {
	if len(original) < 94 {
		return nil, errors.New("short TDS login")
	}
	// Rebuild the supported SQL-password login fields; never forward client authentication extensions.
	out := append([]byte(nil), original[:94]...)
	out[25] &^= 0x80
	out[27] &^= 0x11
	for i := 36; i < 94; i++ {
		out[i] = 0
	}
	for _, pos := range []int{36, 40, 44, 48, 52, 60, 64, 68} {
		value, err := loginField(original, pos, pos == 44)
		if err != nil {
			return nil, err
		}
		switch pos {
		case 40:
			value = user
		case 44:
			value = password
		case 52:
			value = host
		}
		b := textBytes(value)
		if len(b) > 8192 || len(out)+len(b) > 65535 {
			return nil, errors.New("TDS login field too long")
		}
		binary.LittleEndian.PutUint16(out[pos:], uint16(len(out)))
		binary.LittleEndian.PutUint16(out[pos+2:], uint16(len(b)/2))
		if pos == 44 {
			for i, v := range b {
				b[i] = (v<<4 | v>>4) ^ 0xa5
			}
		}
		out = append(out, b...)
	}
	binary.LittleEndian.PutUint32(out, uint32(len(out)))
	return out, nil
}