diff --git a/client.go b/client.go index 12c1693..fd5b9ca 100644 --- a/client.go +++ b/client.go @@ -42,6 +42,7 @@ type TunaSessionClient struct { sessions map[string]*ncp.Session sessionConns map[string]map[string]net.Conn sharedKeys map[string]*[sharedKeySize]byte + connCount map[string]int isClosed bool } @@ -62,6 +63,7 @@ func NewTunaSessionClient(clientAccount *nkn.Account, m *nkn.MultiClient, wallet sessions: make(map[string]*ncp.Session), sessionConns: make(map[string]map[string]net.Conn), sharedKeys: make(map[string]*[sharedKeySize]byte), + connCount: make(map[string]int), } return c, nil @@ -302,7 +304,23 @@ func (c *TunaSessionClient) listenNet(i int) { } } + c.Lock() + c.connCount[sessionKey]++ + c.Unlock() + c.handleConn(conn, sess, i) + + c.Lock() + c.connCount[sessionKey]-- + shouldClose := c.connCount[sessionKey] == 0 + if shouldClose { + delete(c.connCount, sessionKey) + } + c.Unlock() + + if shouldClose { + sess.Close() + } }(conn) } } @@ -460,8 +478,24 @@ func (c *TunaSessionClient) DialWithConfig(remoteAddr string, config *DialConfig for i := 0; i < len(pubAddrs.Addrs); i++ { if conn, ok := conns[connID(i)]; ok { go func(conn net.Conn, i int) { - defer conn.Close() + c.Lock() + c.connCount[sessionKey]++ + c.Unlock() + c.handleConn(conn, sess, i) + conn.Close() + + c.Lock() + c.connCount[sessionKey]-- + shouldClose := c.connCount[sessionKey] == 0 + if shouldClose { + delete(c.connCount, sessionKey) + } + c.Unlock() + + if shouldClose { + sess.Close() + } }(conn, i) } } @@ -566,7 +600,7 @@ func (c *TunaSessionClient) newSession(remoteAddr string, sessionID []byte, conn if writeTimeout > 0 { err := conn.SetWriteDeadline(time.Now().Add(writeTimeout)) if err != nil { - return err + return ncp.ErrConnClosed } } buf, err := c.encode(buf, remoteAddr) @@ -580,7 +614,7 @@ func (c *TunaSessionClient) newSession(remoteAddr string, sessionID []byte, conn if writeTimeout > 0 { err = conn.SetWriteDeadline(zeroTime) if err != nil { - return err + return ncp.ErrConnClosed } } return nil diff --git a/go.mod b/go.mod index 7b5ff96..3471d83 100644 --- a/go.mod +++ b/go.mod @@ -4,9 +4,9 @@ go 1.13 require ( github.com/imdario/mergo v0.3.9 - github.com/nknorg/ncp-go v1.0.1 + github.com/nknorg/ncp-go v1.0.2 github.com/nknorg/nkn-sdk-go v1.3.0 github.com/nknorg/nkn/v2 v2.0.2 - github.com/nknorg/tuna v0.0.0-20200729232930-42c0c6547187 + github.com/nknorg/tuna v0.0.0-20200803091628-f86c69d40928 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 ) diff --git a/go.sum b/go.sum index 9865d17..c326895 100644 --- a/go.sum +++ b/go.sum @@ -74,8 +74,8 @@ github.com/nknorg/go-nat v1.0.1/go.mod h1:dblX1Ac2j08rTUGs5CKCAfjHGN5eDFhbeqt2rc github.com/nknorg/go-portscanner v0.0.0-20181002101859-8493ef01db79 h1:AT5YRim0696HFnGSvFg0uOmSuo1X11roGBQhKB25PHE= github.com/nknorg/go-portscanner v0.0.0-20181002101859-8493ef01db79/go.mod h1:lTev9/EZpc4U3w/sp95T8UPeaav/0Hk+JjykVDuPaQs= github.com/nknorg/ncp-go v1.0.0/go.mod h1:LfiwrFlKZAAoMwhCBA959jSzVD8E9RGOH9nyfXkQSfY= -github.com/nknorg/ncp-go v1.0.1 h1:oJKmcaAcVC1oCJiJ4tDvot7kved8uUGBYXldDHxc8vM= -github.com/nknorg/ncp-go v1.0.1/go.mod h1:LfiwrFlKZAAoMwhCBA959jSzVD8E9RGOH9nyfXkQSfY= +github.com/nknorg/ncp-go v1.0.2 h1:nZRdfXtbrVEBmhxCfEMu20XZIzk8Wztx4OE47drzIWw= +github.com/nknorg/ncp-go v1.0.2/go.mod h1:LfiwrFlKZAAoMwhCBA959jSzVD8E9RGOH9nyfXkQSfY= github.com/nknorg/nkn-sdk-go v1.3.0 h1:iT/MqUx7VyqVo3ivz3VIBbB+VvaByld0YrQiVIlfdCI= github.com/nknorg/nkn-sdk-go v1.3.0/go.mod h1:SeHZ25YcKIXDauhPw/v4XZjprWd+DtLDPNs3/U9/5qE= github.com/nknorg/nkn/v2 v2.0.0-20200703083743-f83cde273fe2/go.mod h1:LxKZTpn5VZTAbjsACSFL5nlQAQ4+WQKyIpVZUYopXm4= @@ -85,8 +85,8 @@ github.com/nknorg/nnet v0.0.0-20200521002812-357d1b11179f h1:JvDKr6D9LdnsdScND+8 github.com/nknorg/nnet v0.0.0-20200521002812-357d1b11179f/go.mod h1:4DHEQEMhlRGIKGSyhATdjeusdqaHafDatadtpeHBpvI= github.com/nknorg/portmapper v0.0.0-20200114081049-1c03cdccc283 h1:uS3/DvxCbi0zZau66ggQAgEjyGmql2mj77UQFgumq1I= github.com/nknorg/portmapper v0.0.0-20200114081049-1c03cdccc283/go.mod h1:dL4PQJ4670oTO6LqvkjrBQEkD+iMiOYjlKRBBw55Csg= -github.com/nknorg/tuna v0.0.0-20200729232930-42c0c6547187 h1:fQyuj5rkyQr8OzG5HdmgAhH+/fq7I3XU7c9vMe89Z6k= -github.com/nknorg/tuna v0.0.0-20200729232930-42c0c6547187/go.mod h1:GhSzMlgyFy9SvT9TCc/UUEGmMiLxPkyeiWaiM0IWT1U= +github.com/nknorg/tuna v0.0.0-20200803091628-f86c69d40928 h1:bO40M16Zu4pFxjIyc2PG4oTfu2IUeZcWFVEST1VYizA= +github.com/nknorg/tuna v0.0.0-20200803091628-f86c69d40928/go.mod h1:GhSzMlgyFy9SvT9TCc/UUEGmMiLxPkyeiWaiM0IWT1U= github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=