really finish initializing the session before returning it / giving up control of the router, in the Conn.search function used by Dial

src/yggdrasil/conn.go

@@ -90,30 +90,37 @@ func (c *Conn) setMTU(from phony.Actor, mtu uint16) {
 
 // This should never be called from the router goroutine, used in the dial functions
 func (c *Conn) search() error {
-	var sinfo *searchInfo
-	var isIn bool
-	phony.Block(&c.core.router, func() {
-		sinfo, isIn = c.core.router.searches.searches[*c.nodeID]
-	})
-	if !isIn {
-		done := make(chan struct{}, 1)
-		var sess *sessionInfo
 	var err error
-		searchCompleted := func(sinfo *sessionInfo, e error) {
-			sess = sinfo
-			err = e
-			// FIXME close can be called multiple times, do a non-blocking send instead
-			select {
-			case done <- struct{}{}:
-			default:
-			}
-		}
+	done := make(chan struct{})
 	phony.Block(&c.core.router, func() {
-			sinfo = c.core.router.searches.newIterSearch(c.nodeID, c.nodeMask, searchCompleted)
+		_, isIn := c.core.router.searches.searches[*c.nodeID]
+		if !isIn {
+			searchCompleted := func(sinfo *sessionInfo, e error) {
+				select {
+				case <-done:
+					// Somehow this was called multiple times, TODO don't let that happen
+					if sinfo != nil {
+						// Need to clean up to avoid a session leak
+						sinfo.cancel.Cancel(nil)
+					}
+				default:
+					if sinfo != nil {
+						// Finish initializing the session
+						sinfo.conn = c
+					}
+					c.session = sinfo
+					err = e
+					close(done)
+				}
+			}
+			sinfo := c.core.router.searches.newIterSearch(c.nodeID, c.nodeMask, searchCompleted)
 			sinfo.continueSearch()
+		} else {
+			err = errors.New("search already exists")
+			close(done)
+		}
 	})
 	<-done
-		c.session = sess
 	if c.session == nil && err == nil {
 		panic("search failed but returned no error")
 	}
@@ -122,13 +129,8 @@ func (c *Conn) search() error {
 		for i := range c.nodeMask {
 			c.nodeMask[i] = 0xFF
 		}
-			c.session.conn = c
 	}
 	return err
-	} else {
-		return errors.New("search already exists")
-	}
-	return nil
 }
 
 // Used in session keep-alive traffic