glider-ssh/conn.go
Jeff Lindsay 48c9603bfc server: timeouts and context canceling on closed connections (#46)
Signed-off-by: Jeff Lindsay <progrium@gmail.com>
2017-07-24 16:25:45 -05:00

51 lines
948 B
Go

package ssh
import (
"context"
"net"
"time"
)
type serverConn struct {
net.Conn
idleTimeout time.Duration
maxDeadline time.Time
closeCanceler context.CancelFunc
}
func (c *serverConn) Write(p []byte) (n int, err error) {
c.updateDeadline()
n, err = c.Conn.Write(p)
if _, isNetErr := err.(net.Error); isNetErr && c.closeCanceler != nil {
c.closeCanceler()
}
return
}
func (c *serverConn) Read(b []byte) (n int, err error) {
c.updateDeadline()
n, err = c.Conn.Read(b)
if _, isNetErr := err.(net.Error); isNetErr && c.closeCanceler != nil {
c.closeCanceler()
}
return
}
func (c *serverConn) Close() (err error) {
err = c.Conn.Close()
if c.closeCanceler != nil {
c.closeCanceler()
}
return
}
func (c *serverConn) updateDeadline() {
idleDeadline := time.Now().Add(c.idleTimeout)
if idleDeadline.Unix() < c.maxDeadline.Unix() {
c.Conn.SetDeadline(idleDeadline)
} else {
c.Conn.SetDeadline(c.maxDeadline)
}
}