fix race condition

listener.go

@@ -83,8 +83,13 @@ func (l *listener) Accept() (tpt.CapableConn, error) {
  if ok {
  select {
  case holePunch.connCh <- conn:
+ // We need to delete the entry from the map here,
+ // in case we accept two connections from the same address.
+ l.transport.holePunchingMx.Lock()
+ delete(l.transport.holePunching, key)
+ l.transport.holePunchingMx.Unlock()
  continue
- case <-holePunch.ctx.Done():
+ default:
  }
  }
 

transport.go

@@ -241,12 +241,6 @@ func (t *transport) holePunch(ctx context.Context, network string, addr *net.UDP
  t.holePunching[key] = activeHolePunch{connCh: connCh, ctx: ctx}
  t.holePunchingMx.Unlock()
 
- defer func() {
- t.holePunchingMx.Lock()
- delete(t.holePunching, key)
- t.holePunchingMx.Unlock()
- }()
-
  payload := make([]byte, 64)
  for i := 0; ; i++ {
  if _, err := rand.Read(payload); err != nil {
@@ -260,12 +254,14 @@ func (t *transport) holePunch(ctx context.Context, network string, addr *net.UDP
  if maxSleep > 200 {
  maxSleep = 200
  }
- sleep := 10*time.Millisecond + time.Duration(rand.Intn(maxSleep))*time.Millisecond
  select {
  case c := <-connCh:
  return c, nil
- case <-time.After(sleep):
+ case <-time.After(10*time.Millisecond + time.Duration(rand.Intn(maxSleep))*time.Millisecond):
  case <-ctx.Done():
+ t.holePunchingMx.Lock()
+ delete(t.holePunching, key)
+ t.holePunchingMx.Unlock()
  return nil, ErrHolePunching
  }
  }