From: Sergey Matveev <stargrave@stargrave.org>
Date: Tue, 5 Oct 2021 16:08:31 +0000 (+0300)
Subject: R/W deadlines on connections
X-Git-Url: http://www.git.stargrave.org/?p=godlighty.git;a=commitdiff_plain;h=5b7eb71cfe22abb59c481d3cee3a719ff6d88095

R/W deadlines on connections
---

diff --git a/cmd/godlighty/main.go b/cmd/godlighty/main.go
index 4198753..e15c165 100644
--- a/cmd/godlighty/main.go
+++ b/cmd/godlighty/main.go
@@ -40,7 +40,10 @@ import (
 
 const MaxConns = 128
 
-var GracefulTime = 10 * time.Second
+var (
+	GracefulTime = 10 * time.Second
+	RWTimeout    = 30 * time.Second
+)
 
 func main() {
 	bind := flag.String("bind", "[::]:80", "Address to bind and listen on")
@@ -57,7 +60,7 @@ func main() {
 	shutdown := make(chan os.Signal)
 	signal.Notify(shutdown, syscall.SIGTERM, syscall.SIGINT, syscall.SIGHUP)
 	exitErr := make(chan error)
-	l, err := net.Listen("tcp", *bind)
+	l, err := godlighty.DeadlinedListen("tcp", *bind, RWTimeout, RWTimeout)
 	if err != nil {
 		log.Fatalln(err)
 	}
@@ -101,7 +104,11 @@ func main() {
 		}
 	}()
 
-	srv := http.Server{Handler: godlighty.MainHandler}
+	srv := http.Server{
+		Handler:           godlighty.MainHandler,
+		ReadHeaderTimeout: RWTimeout,
+		IdleTimeout:       time.Minute,
+	}
 	go func() {
 		<-shutdown
 		log.Println("shutting down")
diff --git a/conn.go b/conn.go
new file mode 100644
index 0000000..81ec846
--- /dev/null
+++ b/conn.go
@@ -0,0 +1,79 @@
+// Deadlined socket read/write. https://github.com/golang/go/issues/16100
+
+package godlighty
+
+import (
+	"net"
+	"time"
+)
+
+type DeadlinedConn struct {
+	net.Conn
+	ReadTimeout              time.Duration
+	WriteTimeout             time.Duration
+	ReadThreshold            int64
+	WriteThreshold           int64
+	BytesReadFromDeadline    int64
+	BytesWrittenFromDeadline int64
+}
+
+func (c *DeadlinedConn) Read(b []byte) (n int, err error) {
+	if c.BytesReadFromDeadline > c.ReadThreshold {
+		c.BytesReadFromDeadline = 0
+		err = c.Conn.SetDeadline(time.Now().Add(c.ReadTimeout))
+		if err != nil {
+			return
+		}
+	}
+	n, err = c.Conn.Read(b)
+	c.BytesReadFromDeadline += int64(n)
+	return
+}
+
+func (c *DeadlinedConn) Write(b []byte) (n int, err error) {
+	if c.BytesWrittenFromDeadline > c.WriteThreshold {
+		c.BytesWrittenFromDeadline = 0
+		err = c.Conn.SetWriteDeadline(time.Now().Add(c.WriteTimeout))
+		if err != nil {
+			return
+		}
+	}
+	n, err = c.Conn.Write(b)
+	c.BytesWrittenFromDeadline += int64(n)
+	return
+}
+
+type DeadlinedListener struct {
+	net.Listener
+	ReadTimeout  time.Duration
+	WriteTimeout time.Duration
+}
+
+func (l *DeadlinedListener) Accept() (net.Conn, error) {
+	c, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return &DeadlinedConn{
+		Conn:           c,
+		ReadTimeout:    l.ReadTimeout,
+		WriteTimeout:   l.WriteTimeout,
+		ReadThreshold:  int64((l.ReadTimeout * 1024) / time.Second),
+		WriteThreshold: int64((l.WriteTimeout * 1024) / time.Second),
+	}, nil
+}
+
+func DeadlinedListen(
+	network, addr string,
+	readTimeout, writeTimeout time.Duration,
+) (net.Listener, error) {
+	l, err := net.Listen(network, addr)
+	if err != nil {
+		return nil, err
+	}
+	return &DeadlinedListener{
+		Listener:     l,
+		ReadTimeout:  readTimeout,
+		WriteTimeout: writeTimeout,
+	}, nil
+}