Setting idletimeout in Go

Question

I have a function in go which is handling connections which are coming through tcp and handled via ssh. I am trying to set an idle timeout by creating struct in the connection function.

Use case - a customer should be able to make a connection and upload/download multiple files

Reference - IdleTimeout in tcp server

Function code:

type Conn struct {
            net.Conn
            idleTimeout time.Duration
        }

func HandleConn(conn net.Conn) {
    var err error
    rAddr := conn.RemoteAddr()
    session := shortuuid.New()
    config := LoadSSHServerConfig(session)
   
    blocklistItem := blocklist.GetBlockListItem(rAddr)
    if blocklistItem.IsBlocked() {
        conn.Close()
        atomic.AddInt64(&stats.Stats.BlockedConnections, 1)
        return
    }

    
    func (c *Conn) Read(b []byte) (int, error) {
        err := c.Conn.SetReadDeadline(time.Now().Add(c.idleTimeout))
        if err != nil {
            return 0, err
        }
        return c.Conn.Read(b)
    }

    sConn, chans, reqs, err := ssh.NewServerConn(conn, config)
    if err != nil {
        if err == io.EOF {
            log.Errorw("SSH: Handshaking was terminated", log.Fields{
                "address": rAddr,
                "error":   err,
                "session": session})
        } else {
            log.Errorw("SSH: Error on handshaking", log.Fields{
                "address": rAddr,
                "error":   err,
                "session": session})
        }

        atomic.AddInt64(&stats.Stats.AuthorizationFailed, 1)
        return
    }

    log.Infow("connection accepted", log.Fields{
        "user": sConn.User(),
    })

    if user, ok := users[session]; ok {
        log.Infow("SSH: Connection accepted", log.Fields{
            "user":          user.LogFields(),
            "clientVersion": string(sConn.ClientVersion())})

        atomic.AddInt64(&stats.Stats.AuthorizationSucceeded, 1)
        
        // The incoming Request channel must be serviced.
        go ssh.DiscardRequests(reqs)

        // Key ID: sConn.Permissions.Extensions["key-id"]
        handleServerConn(user, chans)

        log.Infow("connection finished", log.Fields{"user": user.LogFields()})
        
        log.Infow("checking connections", log.Fields{
            //"cc":          Stats.AcceptedConnections,
            "cc2": &stats.Stats.AcceptedConnections})

        // Remove connection from local cache
        delete(users, session)

    } else {
        log.Infow("user not found from memory", log.Fields{"username": sConn.User()})
    }

}

This code is coming from the Listen function:

func Listen() {
   

    listener, err := net.Listen("tcp", sshListen)
    if err != nil {
        panic(err)
    }
    

    if useProxyProtocol {
        listener = &proxyproto.Listener{
            Listener:           listener,
            ProxyHeaderTimeout: time.Second * 10,
        }
    }

    
    for {
        // Once a ServerConfig has been configured, connections can be accepted.
        conn, err := listener.Accept()
        if err != nil {
            log.Errorw("SSH: Error accepting incoming connection", log.Fields{"error": err})
            atomic.AddInt64(&stats.Stats.FailedConnections, 1)
            continue
        }
    

        // Before use, a handshake must be performed on the incoming net.Conn.
        // It must be handled in a separate goroutine,
        // otherwise one user could easily block entire loop.
        // For example, user could be asked to trust server key fingerprint and hangs.
        go HandleConn(conn)
    }
}

Is that even possible to set a deadline for only the connections which have been idle for 20 secinds (no upload/downloads).

EDIT 1 : Following @LiamKelly's suggestions, I have made the changes in the code. Now the code is like

type SshProxyConn struct {
    net.Conn
    idleTimeout time.Duration
}

func (c *SshProxyConn) Read(b []byte) (int, error) {
    err := c.Conn.SetReadDeadline(time.Now().Add(c.idleTimeout))
    if err != nil {
        return 0, err
    }
    return c.Conn.Read(b)
}
func HandleConn(conn net.Conn) {
    //lines of code as above
    sshproxyconn := &SshProxyConn{nil, time.Second * 20}
    Conn, chans, reqs, err := ssh.NewServerConn(sshproxyconn, config)
    //lines of code
}

But now the issue is that SSH is not happening. I am getting the error "Connection closed" when I try to do ssh. Is it still waiting for "conn" variable in the function call?

Answer 1

Is that even possible to set a deadline for only the connections which have been idle for 20 [seconds]

Ok so first a general disclaimer, I am going to assume go-protoproxy implements the Conn interface as we would expected. Also as you hinted at before, I don't think you can put a a struct method inside another function (I also recommend renaming it something unique to prevent Conn vs net.Conn confusion).

type SshProxyConn struct {
    net.Conn
    idleTimeout time.Duration
}

func (c *SshProxyConn) Read(b []byte) (int, error) {
    err := c.Conn.SetReadDeadline(time.Now().Add(c.idleTimeout))
    if err != nil {
        return 0, err
    }
    return c.Conn.Read(b)
}


func HandleConn(conn net.Conn) {

This makes is more clear what your primary issue is, which you passed the normal net.Conn to your SSH server, not your wrapper class. So

sConn, chans, reqs, err := ssh.NewServerConn(conn, config)

should be EDIT

sshproxyconn := &SshProxyConn{conn, time.Second * 20}
Conn, chans, reqs, err := ssh.NewServerConn(sshproxyconn , config)