89 lines
2.0 KiB
Go
89 lines
2.0 KiB
Go
package protocol
|
|
|
|
import (
|
|
"bytes"
|
|
"testing"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
)
|
|
|
|
func TestPeekBodyLengthAndPrefix(t *testing.T) {
|
|
raw, err := Encode(Frame{Version: CurrentVersion, Command: CmdHeartbeat, Body: []byte("abc")})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
n, err := PeekBodyLength(raw[:HeaderSize])
|
|
if err != nil || n != 3 {
|
|
t.Fatalf("peek body %d %v", n, err)
|
|
}
|
|
if !LooksLikeFrame(raw) {
|
|
t.Fatal("looks like frame")
|
|
}
|
|
if !MaybeFramePrefix(raw[:1]) {
|
|
t.Fatal("single magic byte is a frame prefix")
|
|
}
|
|
if MaybeFramePrefix([]byte("P")) {
|
|
t.Fatal("P is not a frame prefix")
|
|
}
|
|
}
|
|
|
|
func TestCompressRoundTrip(t *testing.T) {
|
|
body := bytes.Repeat([]byte("spark-im-compress"), 40)
|
|
out, flags, err := PrepareBody(body, 32)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if flags&FlagCompress == 0 {
|
|
t.Fatal("expected compress flag")
|
|
}
|
|
got, err := DecodeBody(out, flags, 65536)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !bytes.Equal(got, body) {
|
|
t.Fatal("decompress mismatch")
|
|
}
|
|
}
|
|
|
|
func TestParsePrefixV1AndV2(t *testing.T) {
|
|
line := []byte("PROXY TCP4 203.0.113.10 192.0.2.1 443 9000\r\n")
|
|
h, n, err := ParsePrefix(append(append([]byte{}, line...), 'S', 'P'))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if h.SrcIP.String() != "203.0.113.10" || h.SrcPort != 443 || n != len(line) {
|
|
t.Fatalf("v1 %+v n=%d want %d", h, n, len(line))
|
|
}
|
|
if _, _, err := ParsePrefix([]byte("P")); !NeedMore(err) {
|
|
t.Fatalf("partial v1: %v", err)
|
|
}
|
|
v2 := buildProxyV2(t)
|
|
h2, n2, err := ParsePrefix(v2)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if h2.SrcIP.String() != "203.0.113.10" || h2.SrcPort != 443 || n2 != len(v2) {
|
|
t.Fatalf("v2 %+v n=%d", h2, n2)
|
|
}
|
|
if _, _, err := ParsePrefix([]byte("GET /")); err != errs.ErrProxyHeader {
|
|
t.Fatalf("garbage: %v", err)
|
|
}
|
|
}
|
|
|
|
func buildProxyV2(t *testing.T) []byte {
|
|
t.Helper()
|
|
buf := make([]byte, 16+12)
|
|
copy(buf, proxyV2Sig)
|
|
buf[12] = 0x21
|
|
buf[13] = 0x11
|
|
buf[14] = 0
|
|
buf[15] = 12
|
|
copy(buf[16:20], []byte{203, 0, 113, 10})
|
|
copy(buf[20:24], []byte{192, 0, 2, 1})
|
|
buf[24] = 1
|
|
buf[25] = 187 // 443
|
|
buf[26] = 0x23
|
|
buf[27] = 0x28 // 9000
|
|
return buf
|
|
}
|