mirror of
https://github.com/apernet/hysteria.git
synced 2025-04-03 20:47:38 +03:00
feat: built-in speed test client & server
This commit is contained in:
Toby
parent
84d72ef0b3
commit
a0bd58063b
9 changed files with 1017 additions and 0 deletions
36
extras/outbounds/speedtest.go
Normal file
36
extras/outbounds/speedtest.go
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
package outbounds
|
||||
|
||||
import (
|
||||
"net"
|
||||
|
||||
"github.com/apernet/hysteria/extras/outbounds/speedtest"
|
||||
)
|
||||
|
||||
const (
|
||||
SpeedtestDest = "_SpeedTest"
|
||||
)
|
||||
|
||||
// speedtestHandler is a PluggableOutbound that handles speed test requests.
|
||||
// It's used to intercept speed test requests and return a pseudo connection that
|
||||
// implements the speed test protocol.
|
||||
type speedtestHandler struct {
|
||||
Next PluggableOutbound
|
||||
}
|
||||
|
||||
func NewSpeedtestHandler(next PluggableOutbound) PluggableOutbound {
|
||||
return &speedtestHandler{
|
||||
Next: next,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *speedtestHandler) TCP(reqAddr *AddrEx) (net.Conn, error) {
|
||||
if reqAddr.Host == SpeedtestDest {
|
||||
return speedtest.NewServerConn(), nil
|
||||
} else {
|
||||
return s.Next.TCP(reqAddr)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *speedtestHandler) UDP(reqAddr *AddrEx) (UDPConn, error) {
|
||||
return s.Next.UDP(reqAddr)
|
||||
}
|
||||
125
extras/outbounds/speedtest/client.go
Normal file
125
extras/outbounds/speedtest/client.go
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
package speedtest
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
Conn net.Conn
|
||||
}
|
||||
|
||||
// Download requests the server to send l bytes of data.
|
||||
// The callback function cb is called every second with the time since the last call,
|
||||
// and the number of bytes received in that time.
|
||||
func (c *Client) Download(l uint32, cb func(time.Duration, uint32, bool)) error {
|
||||
err := writeDownloadRequest(c.Conn, l)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ok, msg, err := readDownloadResponse(c.Conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return fmt.Errorf("server rejected download request: %s", msg)
|
||||
}
|
||||
var counter uint32
|
||||
stopChan := make(chan struct{})
|
||||
defer close(stopChan)
|
||||
// Call the callback function every second,
|
||||
// with the time since the last call and the number of bytes received in that time.
|
||||
go func() {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
t := time.Now()
|
||||
for {
|
||||
select {
|
||||
case <-stopChan:
|
||||
return
|
||||
case <-ticker.C:
|
||||
cb(time.Since(t), atomic.SwapUint32(&counter, 0), false)
|
||||
t = time.Now()
|
||||
}
|
||||
}
|
||||
}()
|
||||
buf := make([]byte, chunkSize)
|
||||
startTime := time.Now()
|
||||
remaining := l
|
||||
for remaining > 0 {
|
||||
n := remaining
|
||||
if n > chunkSize {
|
||||
n = chunkSize
|
||||
}
|
||||
rn, err := c.Conn.Read(buf[:n])
|
||||
remaining -= uint32(rn)
|
||||
atomic.AddUint32(&counter, uint32(rn))
|
||||
if err != nil && !(remaining == 0 && err == io.EOF) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// One last call to the callback function to report the total time and bytes received.
|
||||
cb(time.Since(startTime), l, true)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Upload requests the server to receive l bytes of data.
|
||||
// The callback function cb is called every second with the time since the last call,
|
||||
// and the number of bytes sent in that time.
|
||||
func (c *Client) Upload(l uint32, cb func(time.Duration, uint32, bool)) error {
|
||||
err := writeUploadRequest(c.Conn, l)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ok, msg, err := readUploadResponse(c.Conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return fmt.Errorf("server rejected upload request: %s", msg)
|
||||
}
|
||||
var counter uint32
|
||||
stopChan := make(chan struct{})
|
||||
defer close(stopChan)
|
||||
// Call the callback function every second,
|
||||
// with the time since the last call and the number of bytes sent in that time.
|
||||
go func() {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
t := time.Now()
|
||||
for {
|
||||
select {
|
||||
case <-stopChan:
|
||||
return
|
||||
case <-ticker.C:
|
||||
cb(time.Since(t), atomic.SwapUint32(&counter, 0), false)
|
||||
t = time.Now()
|
||||
}
|
||||
}
|
||||
}()
|
||||
buf := make([]byte, chunkSize)
|
||||
remaining := l
|
||||
for remaining > 0 {
|
||||
n := remaining
|
||||
if n > chunkSize {
|
||||
n = chunkSize
|
||||
}
|
||||
_, err := c.Conn.Write(buf[:n])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
remaining -= n
|
||||
atomic.AddUint32(&counter, n)
|
||||
}
|
||||
// Now we should receive the upload summary from the server.
|
||||
elapsed, received, err := readUploadSummary(c.Conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// One last call to the callback function to report the total time and bytes sent.
|
||||
cb(elapsed, received, true)
|
||||
return nil
|
||||
}
|
||||
152
extras/outbounds/speedtest/protocol.go
Normal file
152
extras/outbounds/speedtest/protocol.go
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
package speedtest
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"io"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
typeDownload = 0x1
|
||||
typeUpload = 0x2
|
||||
)
|
||||
|
||||
// DownloadRequest format:
|
||||
// 0x1 (byte)
|
||||
// Request data length (uint32 BE)
|
||||
|
||||
func readDownloadRequest(r io.Reader) (uint32, error) {
|
||||
var l uint32
|
||||
err := binary.Read(r, binary.BigEndian, &l)
|
||||
return l, err
|
||||
}
|
||||
|
||||
func writeDownloadRequest(w io.Writer, l uint32) error {
|
||||
buf := make([]byte, 5)
|
||||
buf[0] = typeDownload
|
||||
binary.BigEndian.PutUint32(buf[1:], l)
|
||||
_, err := w.Write(buf)
|
||||
return err
|
||||
}
|
||||
|
||||
// DownloadResponse format:
|
||||
// Status (byte, 0=ok, 1=error)
|
||||
// Message length (uint16 BE)
|
||||
// Message (bytes)
|
||||
|
||||
func readDownloadResponse(r io.Reader) (bool, string, error) {
|
||||
var status [1]byte
|
||||
if _, err := io.ReadFull(r, status[:]); err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
var msgLen uint16
|
||||
if err := binary.Read(r, binary.BigEndian, &msgLen); err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
// No message is fine
|
||||
if msgLen == 0 {
|
||||
return status[0] == 0, "", nil
|
||||
}
|
||||
msgBuf := make([]byte, msgLen)
|
||||
_, err := io.ReadFull(r, msgBuf)
|
||||
if err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
return status[0] == 0, string(msgBuf), nil
|
||||
}
|
||||
|
||||
func writeDownloadResponse(w io.Writer, ok bool, msg string) error {
|
||||
sz := 1 + 2 + len(msg)
|
||||
buf := make([]byte, sz)
|
||||
if ok {
|
||||
buf[0] = 0
|
||||
} else {
|
||||
buf[0] = 1
|
||||
}
|
||||
binary.BigEndian.PutUint16(buf[1:], uint16(len(msg)))
|
||||
copy(buf[3:], msg)
|
||||
_, err := w.Write(buf)
|
||||
return err
|
||||
}
|
||||
|
||||
// UploadRequest format:
|
||||
// 0x2 (byte)
|
||||
// Upload data length (uint32 BE)
|
||||
|
||||
func readUploadRequest(r io.Reader) (uint32, error) {
|
||||
var l uint32
|
||||
err := binary.Read(r, binary.BigEndian, &l)
|
||||
return l, err
|
||||
}
|
||||
|
||||
func writeUploadRequest(w io.Writer, l uint32) error {
|
||||
buf := make([]byte, 5)
|
||||
buf[0] = typeUpload
|
||||
binary.BigEndian.PutUint32(buf[1:], l)
|
||||
_, err := w.Write(buf)
|
||||
return err
|
||||
}
|
||||
|
||||
// UploadResponse format:
|
||||
// Status (byte, 0=ok, 1=error)
|
||||
// Message length (uint16 BE)
|
||||
// Message (bytes)
|
||||
|
||||
func readUploadResponse(r io.Reader) (bool, string, error) {
|
||||
var status [1]byte
|
||||
if _, err := io.ReadFull(r, status[:]); err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
var msgLen uint16
|
||||
if err := binary.Read(r, binary.BigEndian, &msgLen); err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
// No message is fine
|
||||
if msgLen == 0 {
|
||||
return status[0] == 0, "", nil
|
||||
}
|
||||
msgBuf := make([]byte, msgLen)
|
||||
_, err := io.ReadFull(r, msgBuf)
|
||||
if err != nil {
|
||||
return false, "", err
|
||||
}
|
||||
return status[0] == 0, string(msgBuf), nil
|
||||
}
|
||||
|
||||
func writeUploadResponse(w io.Writer, ok bool, msg string) error {
|
||||
sz := 1 + 2 + len(msg)
|
||||
buf := make([]byte, sz)
|
||||
if ok {
|
||||
buf[0] = 0
|
||||
} else {
|
||||
buf[0] = 1
|
||||
}
|
||||
binary.BigEndian.PutUint16(buf[1:], uint16(len(msg)))
|
||||
copy(buf[3:], msg)
|
||||
_, err := w.Write(buf)
|
||||
return err
|
||||
}
|
||||
|
||||
// UploadSummary format:
|
||||
// Duration (in milliseconds, uint32 BE)
|
||||
// Received data length (uint32 BE)
|
||||
|
||||
func readUploadSummary(r io.Reader) (time.Duration, uint32, error) {
|
||||
var duration uint32
|
||||
if err := binary.Read(r, binary.BigEndian, &duration); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
var l uint32
|
||||
if err := binary.Read(r, binary.BigEndian, &l); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
return time.Duration(duration) * time.Millisecond, l, nil
|
||||
}
|
||||
|
||||
func writeUploadSummary(w io.Writer, duration time.Duration, l uint32) error {
|
||||
buf := make([]byte, 8)
|
||||
binary.BigEndian.PutUint32(buf, uint32(duration/time.Millisecond))
|
||||
binary.BigEndian.PutUint32(buf[4:], l)
|
||||
_, err := w.Write(buf)
|
||||
return err
|
||||
}
|
||||
446
extras/outbounds/speedtest/protocol_test.go
Normal file
446
extras/outbounds/speedtest/protocol_test.go
Normal file
|
|
@ -0,0 +1,446 @@
|
|||
package speedtest
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestReadDownloadRequest(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
want uint32
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
data: []byte{0x0, 0x1, 0xBD, 0xC2},
|
||||
want: 114114,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal zero",
|
||||
data: []byte{0x0, 0x0, 0x0, 0x0},
|
||||
want: 0,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "incomplete",
|
||||
data: []byte{0x0, 0x1, 0x2},
|
||||
want: 0,
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
r := bytes.NewReader(tt.data)
|
||||
got, err := readDownloadRequest(r)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("readDownloadRequest() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("readDownloadRequest() got = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteDownloadRequest(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
l uint32
|
||||
wantW string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
l: 78909912,
|
||||
wantW: "\x01\x04\xB4\x11\xD8",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal zero",
|
||||
l: 0,
|
||||
wantW: "\x01\x00\x00\x00\x00",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := &bytes.Buffer{}
|
||||
err := writeDownloadRequest(w, tt.l)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("writeDownloadRequest() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if gotW := w.String(); gotW != tt.wantW {
|
||||
t.Errorf("writeDownloadRequest() gotW = %v, want %v", gotW, tt.wantW)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadDownloadResponse(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
want bool
|
||||
want1 string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal ok",
|
||||
data: []byte{0x0, 0x0, 0x2, 0x41, 0x42},
|
||||
want: true,
|
||||
want1: "AB",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal ok no message",
|
||||
data: []byte{0x0, 0x0, 0x0, 0x0},
|
||||
want: true,
|
||||
want1: "",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal error",
|
||||
data: []byte{0x1, 0x0, 0x2, 0x43, 0x44},
|
||||
want: false,
|
||||
want1: "CD",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "incomplete",
|
||||
data: []byte{0x0, 0x99, 0x99, 0x45, 0x46, 0x47},
|
||||
want: false,
|
||||
want1: "",
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
r := bytes.NewReader(tt.data)
|
||||
got, got1, err := readDownloadResponse(r)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("readDownloadResponse() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("readDownloadResponse() got = %v, want %v", got, tt.want)
|
||||
}
|
||||
if got1 != tt.want1 {
|
||||
t.Errorf("readDownloadResponse() got1 = %v, want %v", got1, tt.want1)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteDownloadResponse(t *testing.T) {
|
||||
type args struct {
|
||||
ok bool
|
||||
msg string
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
args args
|
||||
wantW string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal ok",
|
||||
args: args{ok: true, msg: "wahaha"},
|
||||
wantW: "\x00\x00\x06wahaha",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal error",
|
||||
args: args{ok: false, msg: "bullbull"},
|
||||
wantW: "\x01\x00\x08bullbull",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "empty ok",
|
||||
args: args{ok: true, msg: ""},
|
||||
wantW: "\x00\x00\x00",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := &bytes.Buffer{}
|
||||
err := writeDownloadResponse(w, tt.args.ok, tt.args.msg)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("writeDownloadResponse() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if gotW := w.String(); gotW != tt.wantW {
|
||||
t.Errorf("writeDownloadResponse() gotW = %v, want %v", gotW, tt.wantW)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadUploadRequest(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
want uint32
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
data: []byte{0x0, 0x0, 0x26, 0xEE},
|
||||
want: 9966,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal zero",
|
||||
data: []byte{0x0, 0x0, 0x0, 0x0},
|
||||
want: 0,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "incomplete",
|
||||
data: []byte{0x1},
|
||||
want: 0,
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
r := bytes.NewReader(tt.data)
|
||||
got, err := readUploadRequest(r)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("readUploadRequest() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("readUploadRequest() got = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteUploadRequest(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
l uint32
|
||||
wantW string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
l: 2291758882,
|
||||
wantW: "\x02\x88\x99\x77\x22",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal zero",
|
||||
l: 0,
|
||||
wantW: "\x02\x00\x00\x00\x00",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := &bytes.Buffer{}
|
||||
err := writeUploadRequest(w, tt.l)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("writeUploadRequest() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if gotW := w.String(); gotW != tt.wantW {
|
||||
t.Errorf("writeUploadRequest() gotW = %v, want %v", gotW, tt.wantW)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadUploadResponse(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
want bool
|
||||
want1 string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal ok",
|
||||
data: []byte{0x0, 0x0, 0x2, 0x41, 0x42},
|
||||
want: true,
|
||||
want1: "AB",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal ok no message",
|
||||
data: []byte{0x0, 0x0, 0x0},
|
||||
want: true,
|
||||
want1: "",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal error",
|
||||
data: []byte{0x1, 0x0, 0x2, 0x43, 0x44},
|
||||
want: false,
|
||||
want1: "CD",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "incomplete",
|
||||
data: []byte{0x0, 0x99, 0x99, 0x45, 0x46, 0x47},
|
||||
want: false,
|
||||
want1: "",
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
r := bytes.NewReader(tt.data)
|
||||
got, got1, err := readUploadResponse(r)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("readUploadResponse() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("readUploadResponse() got = %v, want %v", got, tt.want)
|
||||
}
|
||||
if got1 != tt.want1 {
|
||||
t.Errorf("readUploadResponse() got1 = %v, want %v", got1, tt.want1)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteUploadResponse(t *testing.T) {
|
||||
type args struct {
|
||||
ok bool
|
||||
msg string
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
args args
|
||||
wantW string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal ok",
|
||||
args: args{ok: true, msg: "lul"},
|
||||
wantW: "\x00\x00\x03lul",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "normal error",
|
||||
args: args{ok: false, msg: "notforu"},
|
||||
wantW: "\x01\x00\x07notforu",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "empty ok",
|
||||
args: args{ok: true, msg: ""},
|
||||
wantW: "\x00\x00\x00",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := &bytes.Buffer{}
|
||||
err := writeUploadResponse(w, tt.args.ok, tt.args.msg)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("writeUploadResponse() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if gotW := w.String(); gotW != tt.wantW {
|
||||
t.Errorf("writeUploadResponse() gotW = %v, want %v", gotW, tt.wantW)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadUploadSummary(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
want time.Duration
|
||||
want1 uint32
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
data: []byte{0x0, 0x0, 0x14, 0x6E, 0x0, 0x26, 0x25, 0xA0},
|
||||
want: 5230 * time.Millisecond,
|
||||
want1: 2500000,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "zero",
|
||||
data: []byte{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
|
||||
want: 0,
|
||||
want1: 0,
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "incomplete",
|
||||
data: []byte{0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
|
||||
want: 0,
|
||||
want1: 0,
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
r := bytes.NewReader(tt.data)
|
||||
got, got1, err := readUploadSummary(r)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("readUploadSummary() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("readUploadSummary() got = %v, want %v", got, tt.want)
|
||||
}
|
||||
if got1 != tt.want1 {
|
||||
t.Errorf("readUploadSummary() got1 = %v, want %v", got1, tt.want1)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteUploadSummary(t *testing.T) {
|
||||
type args struct {
|
||||
duration time.Duration
|
||||
l uint32
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
args args
|
||||
wantW string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
args: args{duration: 5230 * time.Millisecond, l: 2500000},
|
||||
wantW: "\x00\x00\x14\x6E\x00\x26\x25\xA0",
|
||||
wantErr: false,
|
||||
},
|
||||
{
|
||||
name: "zero",
|
||||
args: args{duration: 0, l: 0},
|
||||
wantW: "\x00\x00\x00\x00\x00\x00\x00\x00",
|
||||
wantErr: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := &bytes.Buffer{}
|
||||
err := writeUploadSummary(w, tt.args.duration, tt.args.l)
|
||||
if (err != nil) != tt.wantErr {
|
||||
t.Errorf("writeUploadSummary() error = %v, wantErr %v", err, tt.wantErr)
|
||||
return
|
||||
}
|
||||
if gotW := w.String(); gotW != tt.wantW {
|
||||
t.Errorf("writeUploadSummary() gotW = %v, want %v", gotW, tt.wantW)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
99
extras/outbounds/speedtest/server.go
Normal file
99
extras/outbounds/speedtest/server.go
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
package speedtest
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
chunkSize = 64 * 1024
|
||||
)
|
||||
|
||||
// NewServerConn creates a new "pseudo" connection that implements the speed test protocol.
|
||||
// It's called "pseudo" because it's not a real TCP connection - everything is done in memory.
|
||||
func NewServerConn() net.Conn {
|
||||
rConn, iConn := net.Pipe() // return conn & internal conn
|
||||
// Start the server logic
|
||||
go server(iConn)
|
||||
return rConn
|
||||
}
|
||||
|
||||
func server(conn net.Conn) error {
|
||||
defer conn.Close()
|
||||
// First byte determines the request type
|
||||
var typ [1]byte
|
||||
if _, err := io.ReadFull(conn, typ[:]); err != nil {
|
||||
return err
|
||||
}
|
||||
switch typ[0] {
|
||||
case typeDownload:
|
||||
return handleDownload(conn)
|
||||
case typeUpload:
|
||||
return handleUpload(conn)
|
||||
default:
|
||||
return fmt.Errorf("unknown request type: %d", typ[0])
|
||||
}
|
||||
}
|
||||
|
||||
// handleDownload reads the download request and sends the requested amount of data.
|
||||
func handleDownload(conn net.Conn) error {
|
||||
l, err := readDownloadRequest(conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = writeDownloadResponse(conn, true, "OK")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
buf := make([]byte, chunkSize)
|
||||
// Fill the buffer with random data.
|
||||
// For now, we only do it once and repeat the same data for performance reasons.
|
||||
_, err = rand.Read(buf)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
remaining := l
|
||||
for remaining > 0 {
|
||||
n := remaining
|
||||
if n > chunkSize {
|
||||
n = chunkSize
|
||||
}
|
||||
_, err := conn.Write(buf[:n])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
remaining -= n
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleUpload reads the upload request, reads & discards the requested amount of data,
|
||||
// and sends the upload summary.
|
||||
func handleUpload(conn net.Conn) error {
|
||||
l, err := readUploadRequest(conn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = writeUploadResponse(conn, true, "OK")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
buf := make([]byte, chunkSize)
|
||||
startTime := time.Now()
|
||||
remaining := l
|
||||
for remaining > 0 {
|
||||
n := remaining
|
||||
if n > chunkSize {
|
||||
n = chunkSize
|
||||
}
|
||||
rn, err := conn.Read(buf[:n])
|
||||
remaining -= uint32(rn)
|
||||
if err != nil && !(remaining == 0 && err == io.EOF) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return writeUploadSummary(conn, time.Since(startTime), l)
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue