147 lines
3.9 KiB
Go
147 lines
3.9 KiB
Go
// Copyright (c) Tailscale Inc & AUTHORS
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
package speedtest
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"time"
|
|
)
|
|
|
|
// Serve starts up the server on a given host and port pair. It starts to listen for
|
|
// connections and handles each one in a goroutine. Because it runs in an infinite loop,
|
|
// this function only returns if any of the speedtests return with errors, or if the
|
|
// listener is closed.
|
|
func Serve(l net.Listener) error {
|
|
for {
|
|
conn, err := l.Accept()
|
|
if errors.Is(err, net.ErrClosed) {
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
err = handleConnection(conn)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
// handleConnection handles the initial exchange between the server and the client.
|
|
// It reads the testconfig message into a config struct. If any errors occur with
|
|
// the testconfig (specifically, if there is a version mismatch), it will return those
|
|
// errors to the client with a configResponse. After the exchange, it will start
|
|
// the speed test.
|
|
func handleConnection(conn net.Conn) error {
|
|
defer conn.Close()
|
|
var conf config
|
|
|
|
decoder := json.NewDecoder(conn)
|
|
err := decoder.Decode(&conf)
|
|
encoder := json.NewEncoder(conn)
|
|
|
|
// Both return and encode errors that occurred before the test started.
|
|
if err != nil {
|
|
encoder.Encode(configResponse{Error: err.Error()})
|
|
return err
|
|
}
|
|
|
|
// The server should always be doing the opposite of what the client is doing.
|
|
conf.Direction.Reverse()
|
|
|
|
if conf.Version != version {
|
|
err = fmt.Errorf("version mismatch! Server is version %d, client is version %d", version, conf.Version)
|
|
encoder.Encode(configResponse{Error: err.Error()})
|
|
return err
|
|
}
|
|
|
|
// Start the test
|
|
encoder.Encode(configResponse{})
|
|
_, err = doTest(conn, conf)
|
|
return err
|
|
}
|
|
|
|
// TODO include code to detect whether the code is direct vs DERP
|
|
|
|
// doTest contains the code to run both the upload and download speedtest.
|
|
// the direction value in the config parameter determines which test to run.
|
|
func doTest(conn net.Conn, conf config) ([]Result, error) {
|
|
bufferData := make([]byte, blockSize)
|
|
|
|
intervalBytes := 0
|
|
totalBytes := 0
|
|
|
|
var currentTime time.Time
|
|
var results []Result
|
|
|
|
if conf.Direction == Download {
|
|
conn.SetReadDeadline(time.Now().Add(conf.TestDuration).Add(5 * time.Second))
|
|
} else {
|
|
_, err := rand.Read(bufferData)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
}
|
|
|
|
startTime := time.Now()
|
|
lastCalculated := startTime
|
|
|
|
SpeedTestLoop:
|
|
for {
|
|
var n int
|
|
var err error
|
|
|
|
if conf.Direction == Download {
|
|
n, err = io.ReadFull(conn, bufferData)
|
|
switch err {
|
|
case io.EOF, io.ErrUnexpectedEOF:
|
|
break SpeedTestLoop
|
|
case nil:
|
|
// successful read
|
|
default:
|
|
return nil, fmt.Errorf("unexpected error has occurred: %w", err)
|
|
}
|
|
} else {
|
|
n, err = conn.Write(bufferData)
|
|
if err != nil {
|
|
// If the write failed, there is most likely something wrong with the connection.
|
|
return nil, fmt.Errorf("upload failed: %w", err)
|
|
}
|
|
}
|
|
intervalBytes += n
|
|
|
|
currentTime = time.Now()
|
|
// checks if the current time is more or equal to the lastCalculated time plus the increment
|
|
if currentTime.Sub(lastCalculated) >= increment {
|
|
results = append(results, Result{Bytes: intervalBytes, IntervalStart: lastCalculated, IntervalEnd: currentTime, Total: false})
|
|
lastCalculated = currentTime
|
|
totalBytes += intervalBytes
|
|
intervalBytes = 0
|
|
}
|
|
|
|
if conf.Direction == Upload && currentTime.Sub(startTime) > conf.TestDuration {
|
|
break SpeedTestLoop
|
|
}
|
|
}
|
|
|
|
// get last segment
|
|
if currentTime.Sub(lastCalculated) > minInterval {
|
|
results = append(results, Result{Bytes: intervalBytes, IntervalStart: lastCalculated, IntervalEnd: currentTime, Total: false})
|
|
}
|
|
|
|
// get total
|
|
totalBytes += intervalBytes
|
|
if currentTime.Sub(startTime) > minInterval {
|
|
results = append(results, Result{Bytes: totalBytes, IntervalStart: startTime, IntervalEnd: currentTime, Total: true})
|
|
}
|
|
|
|
return results, nil
|
|
}
|