181 lines
3.5 KiB
Go
181 lines
3.5 KiB
Go
package protocol
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/binary"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"strings"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
)
|
|
|
|
const maxProxyLine = 108
|
|
|
|
var (
|
|
errNeedMore = errors.New("proxy header incomplete")
|
|
proxyV2Sig = []byte{0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A}
|
|
proxyV1Magic = []byte("PROXY ")
|
|
)
|
|
|
|
type Header struct {
|
|
SrcIP net.IP
|
|
SrcPort int
|
|
}
|
|
|
|
func NeedMore(err error) bool {
|
|
return errors.Is(err, errNeedMore)
|
|
}
|
|
|
|
// ParsePrefix consumes a PROXY v1/v2 header from buf. n is bytes to discard.
|
|
func ParsePrefix(buf []byte) (h Header, n int, err error) {
|
|
if len(buf) == 0 {
|
|
return Header{}, 0, errNeedMore
|
|
}
|
|
if isPartialPrefix(buf, proxyV1Magic) || isPartialPrefix(buf, proxyV2Sig) {
|
|
return Header{}, 0, errNeedMore
|
|
}
|
|
if hasPrefix(buf, proxyV1Magic) {
|
|
return parseV1(buf)
|
|
}
|
|
if hasPrefix(buf, proxyV2Sig) {
|
|
return parseV2(buf)
|
|
}
|
|
return Header{}, 0, errs.ErrProxyHeader
|
|
}
|
|
|
|
func isPartialPrefix(buf, magic []byte) bool {
|
|
if len(buf) >= len(magic) {
|
|
return false
|
|
}
|
|
return hasPrefix(magic, buf)
|
|
}
|
|
|
|
func parseV1(buf []byte) (Header, int, error) {
|
|
idx := indexCRLF(buf)
|
|
if idx < 0 {
|
|
if len(buf) > maxProxyLine {
|
|
return Header{}, 0, errs.ErrProxyHeader
|
|
}
|
|
return Header{}, 0, errNeedMore
|
|
}
|
|
line := string(buf[:idx])
|
|
h, err := parseV1Line(line)
|
|
if err != nil {
|
|
return Header{}, 0, err
|
|
}
|
|
return h, idx + 2, nil
|
|
}
|
|
|
|
func parseV1Line(line string) (Header, error) {
|
|
fields := strings.Fields(strings.TrimSpace(line))
|
|
if len(fields) < 2 || fields[0] != "PROXY" {
|
|
return Header{}, errs.ErrProxyHeader
|
|
}
|
|
if fields[1] == "UNKNOWN" {
|
|
return Header{}, nil
|
|
}
|
|
if len(fields) < 6 {
|
|
return Header{}, errs.ErrProxyHeader
|
|
}
|
|
ip := net.ParseIP(fields[2])
|
|
if ip == nil {
|
|
return Header{}, errs.ErrProxyHeader
|
|
}
|
|
port, ok := atoi(fields[4])
|
|
if !ok {
|
|
return Header{}, errs.ErrProxyHeader
|
|
}
|
|
return Header{SrcIP: ip, SrcPort: port}, nil
|
|
}
|
|
|
|
func parseV2(buf []byte) (Header, int, error) {
|
|
if len(buf) < 16 {
|
|
return Header{}, 0, errNeedMore
|
|
}
|
|
addrLen := int(binary.BigEndian.Uint16(buf[14:16]))
|
|
total := 16 + addrLen
|
|
if len(buf) < total {
|
|
return Header{}, 0, errNeedMore
|
|
}
|
|
h := Header{}
|
|
fam := buf[13] >> 4
|
|
if fam == 1 && addrLen >= 12 { // IPv4
|
|
h.SrcIP = net.IPv4(buf[16], buf[17], buf[18], buf[19])
|
|
h.SrcPort = int(binary.BigEndian.Uint16(buf[24:26]))
|
|
}
|
|
if fam == 2 && addrLen >= 36 { // IPv6
|
|
h.SrcIP = net.IP(append([]byte(nil), buf[16:32]...))
|
|
h.SrcPort = int(binary.BigEndian.Uint16(buf[48:50]))
|
|
}
|
|
return h, total, nil
|
|
}
|
|
|
|
func ReadHeader(r *bufio.Reader) (Header, error) {
|
|
var acc []byte
|
|
for {
|
|
b, err := r.ReadByte()
|
|
if err != nil {
|
|
return Header{}, err
|
|
}
|
|
acc = append(acc, b)
|
|
if len(acc) > maxProxyLine+16 {
|
|
return Header{}, errs.ErrProxyHeader
|
|
}
|
|
h, n, err := ParsePrefix(acc)
|
|
if NeedMore(err) {
|
|
continue
|
|
}
|
|
if err != nil {
|
|
return Header{}, err
|
|
}
|
|
if n < len(acc) {
|
|
for i := len(acc) - 1; i >= n; i-- {
|
|
_ = r.UnreadByte()
|
|
}
|
|
}
|
|
return h, nil
|
|
}
|
|
}
|
|
|
|
func Discard(r io.Reader) error {
|
|
_, err := io.Copy(io.Discard, r)
|
|
return err
|
|
}
|
|
|
|
func hasPrefix(b, pre []byte) bool {
|
|
if len(b) < len(pre) {
|
|
return false
|
|
}
|
|
for i := range pre {
|
|
if b[i] != pre[i] {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func indexCRLF(b []byte) int {
|
|
for i := 0; i+1 < len(b); i++ {
|
|
if b[i] == '\r' && b[i+1] == '\n' {
|
|
return i
|
|
}
|
|
}
|
|
return -1
|
|
}
|
|
|
|
func atoi(s string) (int, bool) {
|
|
n := 0
|
|
if s == "" {
|
|
return 0, false
|
|
}
|
|
for _, c := range s {
|
|
if c < '0' || c > '9' {
|
|
return 0, false
|
|
}
|
|
n = n*10 + int(c-'0')
|
|
}
|
|
return n, true
|
|
}
|