Răsfoiți Sursa

Fix SplitHTTP race condition when creating new sessions (#3533)

Co-authored-by: nobody <[email protected]>
Co-authored-by: mmmray <[email protected]>
yuhan6665 1 an în urmă
părinte
comite
02cd3b8c74

+ 12 - 0
transport/internet/splithttp/hub.go

@@ -27,6 +27,7 @@ type requestHandler struct {
 	host      string
 	path      string
 	ln        *Listener
+	sessionMu *sync.Mutex
 	sessions  sync.Map
 	localAddr gonet.TCPAddr
 }
@@ -56,11 +57,21 @@ func (h *requestHandler) maybeReapSession(isFullyConnected *done.Instance, sessi
 }
 
 func (h *requestHandler) upsertSession(sessionId string) *httpSession {
+	// fast path
 	currentSessionAny, ok := h.sessions.Load(sessionId)
 	if ok {
 		return currentSessionAny.(*httpSession)
 	}
 
+	// slow path
+	h.sessionMu.Lock()
+	defer h.sessionMu.Unlock()
+
+	currentSessionAny, ok = h.sessions.Load(sessionId)
+	if ok {
+		return currentSessionAny.(*httpSession)
+	}
+
 	s := &httpSession{
 		uploadQueue:      NewUploadQueue(int(2 * h.ln.config.GetNormalizedMaxConcurrentUploads())),
 		isFullyConnected: done.New(),
@@ -277,6 +288,7 @@ func ListenSH(ctx context.Context, address net.Address, port net.Port, streamSet
 		host:      shSettings.Host,
 		path:      shSettings.GetNormalizedPath(),
 		ln:        l,
+		sessionMu: &sync.Mutex{},
 		sessions:  sync.Map{},
 		localAddr: localAddr,
 	}

+ 1 - 1
transport/internet/splithttp/splithttp_test.go

@@ -63,8 +63,8 @@ func Test_listenSHAndDial(t *testing.T) {
 	}
 
 	common.Must(conn.Close())
-	<-time.After(time.Second * 5)
 	conn, err = Dial(ctx, net.TCPDestination(net.DomainAddress("localhost"), listenPort), streamSettings)
+
 	common.Must(err)
 	_, err = conn.Write([]byte("Test connection 2"))
 	common.Must(err)