Signal when a session is closed, other tweaks

cmd/yggdrasil/main.go

@@ -10,7 +10,6 @@ import (
 	"os/signal"
 	"strings"
 	"syscall"
-	"time"
 
 	"golang.org/x/text/encoding/unicode"
 
@@ -268,19 +267,6 @@ func main() {
 	defer func() {
 		n.core.Stop()
 	}()
-	// Some stuff
-	go func() {
-		time.Sleep(time.Second * 2)
-		//session, err := n.core.Dial("nodeid", "babd4e4bccb216f77bb723c1b034b63a652060aabfe9506b51f687183e9b0fd13f438876f5a3ab21cac9c8101eb88e2613fe2a8b0724add09d7ef5a72146c31f")
-		session, err := n.core.Dial("nodeid", "9890e135604e8aa6039a909e40c629824d852042a70e51957d5b9d700195663d50552e8e869af132b4617d76f8ef00314d94cce23aa8d6b051b3b952a32a4966")
-		logger.Println(session, err)
-		b := []byte{1, 2, 3, 4, 5}
-		for {
-			logger.Println(session.Write(b))
-			logger.Println(session.Read(b))
-			time.Sleep(time.Second)
-		}
-	}()
 	// Make some nice output that tells us what our IPv6 address and subnet are.
 	// This is just logged to stdout for the user.
 	address := n.core.Address()

src/yggdrasil/conn.go

@@ -27,9 +27,6 @@ func (c *Conn) startSearch() {
 		}
 		if sinfo != nil {
 			c.session = sinfo
-			c.core.log.Println("Search from API succeeded")
-			c.core.log.Println("Pubkey:", hex.EncodeToString(sinfo.theirPermPub[:]))
-			c.core.log.Println("Coords:", sinfo.coords)
 		}
 	}
 	doSearch := func() {
@@ -69,25 +66,32 @@ func (c *Conn) Read(b []byte) (int, error) {
 		// To prevent blocking forever on a session that isn't initialised
 		return 0, errors.New("session not initialised")
 	}
-	p := <-c.session.recv
+	select {
-	defer util.PutBytes(p.Payload)
+	case p, ok := <-c.session.recv:
-	if !c.session.nonceIsOK(&p.Nonce) {
+		if !ok {
-		return 0, errors.New("packet dropped due to invalid nonce")
+			return 0, errors.New("session was closed")
+		}
+		defer util.PutBytes(p.Payload)
+		if !c.session.nonceIsOK(&p.Nonce) {
+			return 0, errors.New("packet dropped due to invalid nonce")
+		}
+		bs, isOK := crypto.BoxOpen(&c.session.sharedSesKey, p.Payload, &p.Nonce)
+		if !isOK {
+			util.PutBytes(bs)
+			return 0, errors.New("packet dropped due to decryption failure")
+		}
+		b = b[:0]
+		b = append(b, bs...)
+		c.session.updateNonce(&p.Nonce)
+		c.session.time = time.Now()
+		c.session.bytesRecvd += uint64(len(bs))
+		return len(b), nil
+	case <-c.session.closed:
+		return len(b), errors.New("session was closed")
 	}
-	bs, isOK := crypto.BoxOpen(&c.session.sharedSesKey, p.Payload, &p.Nonce)
-	if !isOK {
-		util.PutBytes(bs)
-		return 0, errors.New("packet dropped due to decryption failure")
-	}
-	b = b[:0]
-	b = append(b, bs...)
-	c.session.updateNonce(&p.Nonce)
-	c.session.time = time.Now()
-	c.session.bytesRecvd += uint64(len(bs))
-	return len(b), nil
 }
 
-func (c *Conn) Write(b []byte) (int, error) {
+func (c *Conn) Write(b []byte) (bytesWritten int, err error) {
 	if c.session == nil {
 		c.core.router.doAdmin(func() {
 			c.startSearch()
@@ -112,7 +116,11 @@ func (c *Conn) Write(b []byte) (int, error) {
 	}
 	packet := p.encode()
 	c.session.bytesSent += uint64(len(b))
-	c.session.send <- packet
+	select {
+	case c.session.send <- packet:
+	case <-c.session.closed:
+		return len(b), errors.New("session was closed")
+	}
 	c.session.core.router.out(packet)
 	return len(b), nil
 }

src/yggdrasil/session.go

@@ -41,6 +41,7 @@ type sessionInfo struct {
 	init            bool      // Reset if coords change
 	send            chan []byte
 	recv            chan *wire_trafficPacket
+	closed          chan interface{}
 	nonceMask       uint64
 	tstamp          int64     // tstamp from their last session ping, replay attack mitigation
 	tstampMutex     int64     // protects the above
@@ -306,6 +307,7 @@ func (ss *sessions) createSession(theirPermKey *crypto.BoxPubKey) *sessionInfo {
 	sinfo.theirSubnet = *address.SubnetForNodeID(crypto.GetNodeID(&sinfo.theirPermPub))
 	sinfo.send = make(chan []byte, 32)
 	sinfo.recv = make(chan *wire_trafficPacket, 32)
+	sinfo.closed = make(chan interface{})
 	ss.sinfos[sinfo.myHandle] = &sinfo
 	ss.byMySes[sinfo.mySesPub] = &sinfo.myHandle
 	ss.byTheirPerm[sinfo.theirPermPub] = &sinfo.myHandle
@@ -364,6 +366,8 @@ func (ss *sessions) cleanup() {
 
 // Closes a session, removing it from sessions maps and killing the worker goroutine.
 func (sinfo *sessionInfo) close() {
+	sinfo.init = false
+	close(sinfo.closed)
 	delete(sinfo.core.sessions.sinfos, sinfo.myHandle)
 	delete(sinfo.core.sessions.byMySes, sinfo.mySesPub)
 	delete(sinfo.core.sessions.byTheirPerm, sinfo.theirPermPub)