Add SocketId() and PeerSocketId() to ConnRequest (#104)

Author: ioppermann

conn_request.go

@@ -26,6 +26,12 @@ type ConnRequest interface {
 	// to decide what to do with the connection.
 	StreamId() string
 
+	// SocketId return the socketid of the connection.
+	SocketId() uint32
+
+	// PeerSocketId returns the socketid of the peer of the connection.
+	PeerSocketId() uint32
+
 	// IsEncrypted returns whether the connection is encrypted. If it is
 	// encrypted, use SetPassphrase to set the passphrase for decrypting.
 	IsEncrypted() bool
@@ -54,6 +60,7 @@ type connRequest struct {
 	addr            net.Addr
 	start           time.Time
 	socketId        uint32
+	peerSocketId    uint32
 	timestamp       uint32
 	config          Config
 	handshake       *packet.CIFHandshake
@@ -231,13 +238,13 @@ func newConnRequest(ln *listener, p packet.Packet) *connRequest {
 		}
 
 		req := &connRequest{
-			ln:        ln,
-			addr:      p.Header().Addr,
-			start:     time.Now(),
-			socketId:  cif.SRTSocketId,
-			timestamp: p.Header().Timestamp,
-			config:    config,
-			handshake: cif,
+			ln:           ln,
+			addr:         p.Header().Addr,
+			start:        time.Now(),
+			peerSocketId: cif.SRTSocketId,
+			timestamp:    p.Header().Timestamp,
+			config:       config,
+			handshake:    cif,
 		}
 
 		if cif.SRTKM != nil {
@@ -263,11 +270,25 @@ func newConnRequest(ln *listener, p packet.Packet) *connRequest {
 		}
 		ln.lock.Unlock()
 
-		// we received a duplicate request: reject silently
+		// We received a duplicate request: reject silently
 		if exists {
 			return nil
 		}
 
+		// Already reserve a socketId for this connection
+		ln.lock.Lock()
+		socketId, err := req.generateSocketId()
+		if err == nil {
+			ln.conns[socketId] = nil
+			req.socketId = socketId
+		}
+		ln.lock.Unlock()
+
+		// We couldn't create a socketId: reject silently
+		if err != nil {
+			return nil
+		}
+
 		return req
 	} else {
 		if cif.HandshakeType.IsRejection() {
@@ -293,6 +314,14 @@ func (req *connRequest) StreamId() string {
 	return req.handshake.StreamId
 }
 
+func (req *connRequest) SocketId() uint32 {
+	return req.socketId
+}
+
+func (req *connRequest) PeerSocketId() uint32 {
+	return req.peerSocketId
+}
+
 func (req *connRequest) IsEncrypted() bool {
 	return req.crypto != nil
 }
@@ -321,7 +350,7 @@ func (req *connRequest) Reject(reason RejectionReason) {
 	req.ln.lock.Lock()
 	defer req.ln.lock.Unlock()
 
-	if _, hasReq := req.ln.connReqs[req.socketId]; !hasReq {
+	if _, hasReq := req.ln.connReqs[req.peerSocketId]; !hasReq {
 		return
 	}
 
@@ -331,14 +360,15 @@ func (req *connRequest) Reject(reason RejectionReason) {
 	p.Header().SubType = 0
 	p.Header().TypeSpecific = 0
 	p.Header().Timestamp = uint32(time.Since(req.ln.start).Microseconds())
-	p.Header().DestinationSocketId = req.socketId
+	p.Header().DestinationSocketId = req.peerSocketId
 	req.handshake.HandshakeType = packet.HandshakeType(reason)
 	p.MarshalCIF(req.handshake)
 	req.ln.log("handshake:send:dump", func() string { return p.Dump() })
 	req.ln.log("handshake:send:cif", func() string { return req.handshake.String() })
 	req.ln.send(p)
 
-	delete(req.ln.connReqs, req.socketId)
+	delete(req.ln.connReqs, req.peerSocketId)
+	delete(req.ln.conns, req.socketId)
 }
 
 // generateSocketId generates an SRT SocketID that can be used for this connection
@@ -367,16 +397,10 @@ func (req *connRequest) Accept() (Conn, error) {
 	req.ln.lock.Lock()
 	defer req.ln.lock.Unlock()
 
-	if _, hasReq := req.ln.connReqs[req.socketId]; !hasReq {
+	if _, hasReq := req.ln.connReqs[req.peerSocketId]; !hasReq {
 		return nil, fmt.Errorf("connection already accepted")
 	}
 
-	// Create a new socket ID
-	socketId, err := req.generateSocketId()
-	if err != nil {
-		return nil, fmt.Errorf("could not generate socket id: %w", err)
-	}
-
 	// Select the largest TSBPD delay advertised by the caller, but at least 120ms
 	recvTsbpdDelay := uint16(req.config.ReceiverLatency.Milliseconds())
 	sendTsbpdDelay := uint16(req.config.PeerLatency.Milliseconds())
@@ -402,8 +426,8 @@ func (req *connRequest) Accept() (Conn, error) {
 		remoteAddr:                  req.addr,
 		config:                      req.config,
 		start:                       req.start,
-		socketId:                    socketId,
-		peerSocketId:                req.handshake.SRTSocketId,
+		socketId:                    req.socketId,
+		peerSocketId:                req.peerSocketId,
 		tsbpdTimeBase:               uint64(req.timestamp),
 		tsbpdDelay:                  uint64(recvTsbpdDelay) * 1000,
 		peerTsbpdDelay:              uint64(sendTsbpdDelay) * 1000,
@@ -417,7 +441,7 @@ func (req *connRequest) Accept() (Conn, error) {
 
 	req.ln.log("connection:new", func() string { return fmt.Sprintf("%#08x (%s)", conn.SocketId(), conn.StreamId()) })
 
-	req.handshake.SRTSocketId = socketId
+	req.handshake.SRTSocketId = req.socketId
 	req.handshake.SynCookie = 0
 
 	if req.handshake.Version == 5 {
@@ -441,14 +465,14 @@ func (req *connRequest) Accept() (Conn, error) {
 	p.Header().SubType = 0
 	p.Header().TypeSpecific = 0
 	p.Header().Timestamp = uint32(time.Since(req.start).Microseconds())
-	p.Header().DestinationSocketId = req.socketId
+	p.Header().DestinationSocketId = req.peerSocketId
 	p.MarshalCIF(req.handshake)
 	req.ln.log("handshake:send:dump", func() string { return p.Dump() })
 	req.ln.log("handshake:send:cif", func() string { return req.handshake.String() })
 	req.ln.send(p)
 
-	req.ln.conns[socketId] = conn
-	delete(req.ln.connReqs, req.socketId)
+	req.ln.conns[req.socketId] = conn
+	delete(req.ln.connReqs, req.peerSocketId)
 
 	return conn, nil
 }

listen.go

@@ -347,6 +347,9 @@ func (ln *listener) Close() {
 
 		ln.lock.RLock()
 		for _, conn := range ln.conns {
+			if conn == nil {
+				continue
+			}
 			conn.close()
 		}
 		ln.lock.RUnlock()
@@ -402,7 +405,7 @@ func (ln *listener) reader(ctx context.Context) {
 			conn, ok := ln.conns[p.Header().DestinationSocketId]
 			ln.lock.RUnlock()
 
-			if !ok {
+			if !ok || conn == nil {
 				// ignore the packet, we don't know the destination
 				break
 			}