tun2: cleanup

lib/tun2/server.go

@@ -244,6 +244,26 @@ func (c *Connection) OpenStream() (net.Conn, error) {
 	return stream, c.conn.SetDeadline(time.Time{})
 }
 
+// Close destroys resouces specific to the connection.
+func (c *Connection) Close() error {
+	err := c.controlStream.Close()
+	if err != nil {
+		return err
+	}
+
+	err = c.session.Close()
+	if err != nil {
+		return err
+	}
+
+	err = c.conn.Close()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (s *Server) HandleConn(c net.Conn, isKCP bool) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -369,8 +389,17 @@ func (s *Server) HandleConn(c net.Conn, isKCP bool) {
 				cancel()
 			}
 		case <-ctx.Done():
+			s.RemoveConn(auth, connection)
+			connection.Close()
+
+			return
+		}
+	}
+}
+
+func (s *Server) RemoveConn(auth *Auth, connection *Connection) {
 	s.connlock.Lock()
-			delete(s.conns, c)
+	delete(s.conns, connection.conn)
 	s.connlock.Unlock()
 
 	var conns []*Connection
@@ -379,7 +408,7 @@ func (s *Server) HandleConn(c net.Conn, isKCP bool) {
 	if ok {
 		conns, ok = val.([]*Connection)
 		if !ok {
-					ln.Error(err, connection.F(), ln.F{
+			ln.Error(errors.New("fundamental assertion is not met"), connection.F(), ln.F{
 				"action": "looking_up_for_disconnect_removal",
 			})
 			return
@@ -402,14 +431,6 @@ func (s *Server) HandleConn(c net.Conn, isKCP bool) {
 	ln.Log(connection.F(), ln.F{
 		"action": "client_disconnecting",
 	})
-
-			controlStream.Close()
-			session.Close()
-			c.Close()
-
-			return
-		}
-	}
 }
 
 func (s *Server) RoundTrip(req *http.Request) (*http.Response, error) {