From c511a3c8959342a3e82e03f1620f9d47071f2694 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Sat, 21 Mar 2020 10:24:38 +0100
Subject: [PATCH] Add network.upstream helper

---
 downstream.go |  7 +------
 user.go       | 11 +++++++----
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/downstream.go b/downstream.go
index f864661..7844b49 100644
--- a/downstream.go
+++ b/downstream.go
@@ -143,12 +143,7 @@ func (dc *downstreamConn) upstream() *upstreamConn {
 	if dc.network == nil {
 		return nil
 	}
-
-	var upstream *upstreamConn
-	dc.forEachUpstream(func(uc *upstreamConn) {
-		upstream = uc
-	})
-	return upstream
+	return dc.network.upstream()
 }
 
 func (dc *downstreamConn) marshalEntity(uc *upstreamConn, entity string) string {
diff --git a/user.go b/user.go
index b3c5873..629e0a1 100644
--- a/user.go
+++ b/user.go
@@ -67,6 +67,12 @@ func (net *network) run() {
 	}
 }
 
+func (net *network) upstream() *upstreamConn {
+	net.lock.Lock()
+	defer net.lock.Unlock()
+	return net.conn
+}
+
 type user struct {
 	User
 	srv *Server
@@ -99,10 +105,7 @@ func (u *user) forEachNetwork(f func(*network)) {
 func (u *user) forEachUpstream(f func(uc *upstreamConn)) {
 	u.lock.Lock()
 	for _, network := range u.networks {
-		network.lock.Lock()
-		uc := network.conn
-		network.lock.Unlock()
-
+		uc := network.upstream()
 		if uc == nil || !uc.registered || uc.closed {
 			continue
 		}