From df4a2d77a28eb73999b39693920e2a0fff5fc0ea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9E=97=E5=BF=97=E5=AE=87?= <evanlinjin@pm.me>
Date: Fri, 14 Jun 2019 00:44:16 +0800
Subject: [PATCH] Fixed possible panics of transport re-establishment.

---
 pkg/transport/manager.go | 37 +++++++++++++++++++------------------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/pkg/transport/manager.go b/pkg/transport/manager.go
index b432f6c64..ebc4d4145 100644
--- a/pkg/transport/manager.go
+++ b/pkg/transport/manager.go
@@ -3,6 +3,7 @@ package transport
 import (
 	"context"
 	"errors"
+	"io"
 	"math/big"
 	"strings"
 	"sync"
@@ -241,21 +242,17 @@ func (tm *Manager) CreateTransport(ctx context.Context, remote cipher.PubKey, tp
 // DeleteTransport disconnects and removes the Transport of Transport ID.
 func (tm *Manager) DeleteTransport(id uuid.UUID) error {
 	tm.mu.Lock()
-	tr := tm.transports[id]
-	delete(tm.transports, id)
+	if tr, ok := tm.transports[id]; ok {
+		delete(tm.transports, id)
+		_ = tr.Close() //nolint:errcheck
+	}
 	tm.mu.Unlock()
 
-	tr.Close()
-
 	if _, err := tm.config.DiscoveryClient.UpdateStatuses(context.Background(), &Status{ID: id, IsUp: false}); err != nil {
 		tm.Logger.Warnf("Failed to change transport status: %s", err)
 	}
 
 	tm.Logger.Infof("Unregistered transport %s", id)
-	if tr != nil {
-		return tr.Close()
-	}
-
 	return nil
 }
 
@@ -330,11 +327,13 @@ func (tm *Manager) createTransport(ctx context.Context, remote cipher.PubKey, tp
 	tm.transports[entry.ID] = mTr
 	tm.mu.Unlock()
 
-	tm.TrChan <- mTr
-
-	go tm.manageTransport(ctx, mTr, factory, remote, public, false)
-
-	return mTr, nil
+	select {
+	case <-tm.doneChan:
+		return nil, io.ErrClosedPipe
+	case tm.TrChan <- mTr:
+		go tm.manageTransport(ctx, mTr, factory, remote, public, false)
+		return mTr, nil
+	}
 }
 
 func (tm *Manager) acceptTransport(ctx context.Context, factory Factory) (*ManagedTransport, error) {
@@ -371,11 +370,13 @@ func (tm *Manager) acceptTransport(ctx context.Context, factory Factory) (*Manag
 	tm.transports[entry.ID] = mTr
 	tm.mu.Unlock()
 
-	tm.TrChan <- mTr
-
-	go tm.manageTransport(ctx, mTr, factory, remote, true, true)
-
-	return mTr, nil
+	select {
+	case <-tm.doneChan:
+		return nil, io.ErrClosedPipe
+	case tm.TrChan <- mTr:
+		go tm.manageTransport(ctx, mTr, factory, remote, true, true)
+		return mTr, nil
+	}
 }
 
 func (tm *Manager) walkEntries(walkFunc func(*Entry) bool) *Entry {