hysteria/extras/outbounds/speedtest/server.go

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)
}