Ailur
/
httpserver
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 2 Wiki Activity

Added support for throttling connections

This commit is contained in:
Tracker-Friendly 2024-09-02 15:58:27 +01:00
parent b48d45e3cf
commit d8264ae2b7
3 changed files with 87 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
httpserver/httpserver
View File

Binary file not shown.

29
httpserver/main.go
View File

@ -30,11 +30,12 @@ func handleErr(err error, errCode int) {
func main() {
if len(os.Args) < 2 {
err, errCode := httpserver.StartServer("8000", "./", "0.0.0.0", "1.0")
err, errCode := httpserver.StartServer("8000", "./", "0.0.0.0", "1.0", -1)
handleErr(err, errCode)
} else {
args := os.Args[1:]
protocolVer, path, address, port := "1.0", "./", "0.0.0.0", "8000"
throttleRate := int64(-1)
for num, arg := range args {
if strings.Contains(arg, "-") {
if strings.Contains(arg, "-p") || strings.Contains(arg, "--protocol") {
@ -42,7 +43,7 @@ func main() {
protocolVer = args[num+1]
args = append(args[:num+1], args[num+2:]...)
} else {
fmt.Println("usage: httpserver [-h] [--cgi] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [port]\nhttpserver: error: argument -p/--protocol: expected one argument")
fmt.Println("usage: httpserver [-h] [--cgi] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]\nhttpserver: error: argument -p/--protocol: expected one argument")
os.Exit(2)
}
} else if strings.Contains(arg, "-d") || strings.Contains(arg, "--directory") {
@ -50,7 +51,7 @@ func main() {
path = args[num+1]
args = append(args[:num+1], args[num+2:]...)
} else {
fmt.Println("usage: httpserver [-h] [--cgi] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [port]\nhttpserver: error: argument -d/--directory: expected one argument")
fmt.Println("usage: httpserver [-h] [--cgi] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]\nhttpserver: error: argument -d/--directory: expected one argument")
os.Exit(2)
}
} else if strings.Contains(arg, "-b") || strings.Contains(arg, "--bind") {
@ -58,14 +59,26 @@ func main() {
address = args[num+1]
args = append(args[:num+1], args[num+2:]...)
} else {
fmt.Println("usage: httpserver [-h] [--cgi] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [port]\nhttpserver: error: argument -b/--bind: expected one argument")
fmt.Println("usage: httpserver [-h] [--cgi] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]\nhttpserver: error: argument -b/--bind: expected one argument")
os.Exit(2)
}
} else if strings.Contains(arg, "-t") || strings.Contains(arg, "--throttle") {
if len(args) > num+1 {
var err error
throttleRate, err = strconv.ParseInt(args[num+1], 10, 64)
if err != nil {
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]\nhttpserver: error: argument -t/--throttle: invalid int value: '" + args[num+1] + "'")
os.Exit(2)
}
} else {
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]\nhttpserver: error: argument -t/--throttle: expected one argument")
os.Exit(2)
}
} else if strings.Contains(arg, "-h") || strings.Contains(arg, "--help") {
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [port]\n\npositional arguments:\n port bind to this port (default: 8000)\n\noptions:\n -h, --help show this help message and exit\n -b ADDRESS, --bind ADDRESS\n bind to this address (default: all interfaces)\n -d DIRECTORY, --directory DIRECTORY\n serve this directory (default: current directory)\n -p VERSION, --protocol VERSION\n conform to this HTTP version (default: HTTP/1.0)")
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]\n\npositional arguments:\n port bind to this port (default: 8000)\n\noptions:\n -h, --help show this help message and exit\n -b ADDRESS, --bind ADDRESS\n bind to this address (default: all interfaces)\n -d DIRECTORY, --directory DIRECTORY\n serve this directory (default: current directory)\n -p VERSION, --protocol VERSION\n conform to this HTTP version (default: HTTP/1.0)")
os.Exit(0)
} else {
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [port]")
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]")
fmt.Println("httpserver: error: unrecognized arguments: " + arg)
os.Exit(2)
}
@ -74,7 +87,7 @@ func main() {
if args[num] == arg {
_, err := strconv.Atoi(arg)
if err != nil {
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [port]")
fmt.Println("usage: httpserver [-h] [-b ADDRESS] [-d DIRECTORY] [-p VERSION] [-t RATE] [port]")
fmt.Println("httpserver: error: argument port: invalid int value: '" + arg + "'")
os.Exit(2)
} else {
@ -84,7 +97,7 @@ func main() {
}
}
}
err, errCode := httpserver.StartServer(port, path, address, protocolVer)
err, errCode := httpserver.StartServer(port, path, address, protocolVer, throttleRate)
handleErr(err, errCode)
}
}

69
main.go
View File

@ -14,7 +14,7 @@ import (
// 1234567. Very clever, Google.
var timeLayout = "02/Jan/2006 15:04:05"
func StartServer(port string, path string, address string, protocolVer string) (error, int) {
func StartServer(port string, path string, address string, protocolVer string, throttleRate int64) (error, int) {
var httpServer *http.Server
addressPort := address + ":" + port
fileServer := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@ -29,11 +29,18 @@ func StartServer(port string, path string, address string, protocolVer string) (
fmt.Println(ip + " - - [" + time.Now().Format(timeLayout) + "] \"" + r.Method + " " + r.URL.Path + " " + r.Proto + "\" " + "200" + " -")
})
var throttledFileServer http.Handler
if throttleRate != -1 {
throttledFileServer = ThrottleMiddleware(throttleRate)(fileServer)
} else {
throttledFileServer = fileServer
}
fmt.Println("Serving HTTP on", address, "port", port, "(http://"+address+":"+port+"/) ...")
if protocolVer == "2.0" || protocolVer == "2" {
httpServer = &http.Server{Addr: addressPort, Handler: fileServer, TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler))}
httpServer = &http.Server{Addr: addressPort, Handler: throttledFileServer, TLSNextProto: make(map[string]func(*http.Server, *tls.Conn, http.Handler))}
} else {
httpServer = &http.Server{Addr: addressPort, Handler: fileServer}
httpServer = &http.Server{Addr: addressPort, Handler: throttledFileServer}
}
err := httpServer.ListenAndServe()
if err != nil {
@ -45,3 +52,59 @@ func StartServer(port string, path string, address string, protocolVer string) (
}
return nil, 0
}
func ThrottleMiddleware(rate int64, burstSize ...int64) func(http.Handler) http.Handler {
defaultBurstSize := int64(128)
if len(burstSize) > 0 && burstSize[0] > 0 {
defaultBurstSize = burstSize[0]
}
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
tw := &ThrottledResponseWriter{
writer: w,
rate: rate,
burstSize: defaultBurstSize,
}
next.ServeHTTP(tw, r)
})
}
}
type ThrottledResponseWriter struct {
writer http.ResponseWriter
rate int64
burstSize int64
}
func (tw *ThrottledResponseWriter) Write(p []byte) (int, error) {
if len(p) == 0 {
return 0, nil
}
var written int
for written < len(p) {
chunkSize := tw.burstSize
if int64(len(p)-written) < chunkSize {
chunkSize = int64(len(p) - written)
}
n, err := tw.writer.Write(p[written : written+int(chunkSize)])
if err != nil {
return written + n, err
}
written += n
time.Sleep(time.Duration(chunkSize*8*int64(time.Second)/tw.rate) * time.Nanosecond)
}
return written, nil
}
func (tw *ThrottledResponseWriter) Header() http.Header {
return tw.writer.Header()
}
func (tw *ThrottledResponseWriter) WriteHeader(statusCode int) {
tw.writer.WriteHeader(statusCode)
}