handshake timeout

conn.go

@@ -9,13 +9,16 @@ import (
 type serverConn struct {
 	net.Conn
 
-	idleTimeout   time.Duration
-	maxDeadline   time.Time
-	closeCanceler context.CancelFunc
+	idleTimeout       time.Duration
+	handshakeDeadline time.Time
+	maxDeadline       time.Time
+	closeCanceler     context.CancelFunc
 }
 
 func (c *serverConn) Write(p []byte) (n int, err error) {
-	c.updateDeadline()
+	if c.idleTimeout > 0 {
+		c.updateDeadline()
+	}
 	n, err = c.Conn.Write(p)
 	if _, isNetErr := err.(net.Error); isNetErr && c.closeCanceler != nil {
 		c.closeCanceler()
@@ -24,7 +27,9 @@ func (c *serverConn) Write(p []byte) (n int, err error) {
 }
 
 func (c *serverConn) Read(b []byte) (n int, err error) {
-	c.updateDeadline()
+	if c.idleTimeout > 0 {
+		c.updateDeadline()
+	}
 	n, err = c.Conn.Read(b)
 	if _, isNetErr := err.(net.Error); isNetErr && c.closeCanceler != nil {
 		c.closeCanceler()
@@ -41,15 +46,18 @@ func (c *serverConn) Close() (err error) {
 }
 
 func (c *serverConn) updateDeadline() {
-	switch {
-	case c.idleTimeout > 0:
+	deadline := c.maxDeadline
+
+	if !c.handshakeDeadline.IsZero() && (deadline.IsZero() || c.handshakeDeadline.Before(deadline)) {
+		deadline = c.handshakeDeadline
+	}
+
+	if c.idleTimeout > 0 {
 		idleDeadline := time.Now().Add(c.idleTimeout)
-		if idleDeadline.Unix() < c.maxDeadline.Unix() || c.maxDeadline.IsZero() {
-			c.Conn.SetDeadline(idleDeadline)
-			return
+		if deadline.IsZero() || idleDeadline.Before(deadline) {
+			deadline = idleDeadline
 		}
-		fallthrough
-	default:
-		c.Conn.SetDeadline(c.maxDeadline)
 	}
+
+	c.Conn.SetDeadline(deadline)
 }

server.go

@@ -50,8 +50,9 @@ type Server struct {
 
 	ConnectionFailedCallback ConnectionFailedCallback // callback to report connection failures
 
-	IdleTimeout time.Duration // connection timeout when no activity, none if empty
-	MaxTimeout  time.Duration // absolute connection timeout, none if empty
+	HandshakeTimeout time.Duration // connection timeout until successful handshake, none if empty
+	IdleTimeout      time.Duration // connection timeout when no activity, none if empty
+	MaxTimeout       time.Duration // absolute connection timeout, none if empty
 
 	// ChannelHandlers allow overriding the built-in session handlers or provide
 	// extensions to the protocol, such as tcpip forwarding. By default only the
@@ -277,6 +278,10 @@ func (srv *Server) HandleConn(newConn net.Conn) {
 	if srv.MaxTimeout > 0 {
 		conn.maxDeadline = time.Now().Add(srv.MaxTimeout)
 	}
+	if srv.HandshakeTimeout > 0 {
+		conn.handshakeDeadline = time.Now().Add(srv.HandshakeTimeout)
+	}
+	conn.updateDeadline()
 	defer conn.Close()
 	sshConn, chans, reqs, err := gossh.NewServerConn(conn, srv.config(ctx))
 	if err != nil {
@@ -285,7 +290,8 @@ func (srv *Server) HandleConn(newConn net.Conn) {
 		}
 		return
 	}
-
+	conn.handshakeDeadline = time.Time{}
+	conn.updateDeadline()
 	srv.trackConn(sshConn, true)
 	defer srv.trackConn(sshConn, false)