Update listener to use stream session

Listener is currently broken when multiple connections are created, using sessions allows proper separation of each connection.

Signed-off-by: Derek McGowan <derek@mcgstyle.net> (github: dmcgowan)

Author: dmcgowan

http2/listener.go

@@ -1,21 +1,14 @@
 package http2
 
 import (
-	"github.com/docker/libchan"
-	"github.com/docker/spdystream"
 	"net"
-	"net/http"
-	"sync"
 )
 
 // ListenSession is a session manager which accepts new
 // connections and spawns spdy connection managers.
 type ListenSession struct {
-	listener       net.Listener
-	streamChan     chan *spdystream.Stream
-	streamLock     sync.RWMutex
-	subStreamChans map[string]chan *spdystream.Stream
-	auth           Authenticator
+	listener net.Listener
+	auth     Authenticator
 }
 
 // NewListenSession creates a new listen session using
@@ -25,78 +18,32 @@ type ListenSession struct {
 // ListenSession will not perform tls handshakes.
 func NewListenSession(listener net.Listener, auth Authenticator) (*ListenSession, error) {
 	return &ListenSession{
-		listener:       listener,
-		streamChan:     make(chan *spdystream.Stream),
-		subStreamChans: make(map[string]chan *spdystream.Stream),
-		auth:           auth,
+		listener: listener,
+		auth:     auth,
 	}, nil
 }
 
-func (l *ListenSession) streamHandler(stream *spdystream.Stream) {
-	// TODO authorize stream
-	stream.SendReply(http.Header{}, false)
-	streamChan := l.getStreamChan(stream.Parent())
-	streamChan <- stream
-}
-
-func (l *ListenSession) addStreamChan(stream *spdystream.Stream, streamChan chan *spdystream.Stream) {
-	l.streamLock.Lock()
-	l.subStreamChans[stream.String()] = streamChan
-	l.streamLock.Unlock()
-}
-
-func (l *ListenSession) getStreamChan(stream *spdystream.Stream) chan *spdystream.Stream {
-	if stream == nil {
-		return l.streamChan
-	}
-	l.streamLock.RLock()
-	defer l.streamLock.RUnlock()
-	streamChan, ok := l.subStreamChans[stream.String()]
-	if ok {
-		return streamChan
-	}
-	return l.streamChan
+// Close closes the underlying listener
+func (l *ListenSession) Close() error {
+	return l.listener.Close()
 }
 
-// Serve runs the listen accept loop, spawning connection manager
-// for each accepted connection. This function will block until
-// the listener is closed or an error occurs on accept.
-func (l *ListenSession) Serve() {
+// AcceptSessions waits for a new network connections
+// and creates a new stream.  Connections which fail
+// authentication will not be returned.
+func (l *ListenSession) AcceptSession() (*StreamSession, error) {
 	for {
 		conn, err := l.listener.Accept()
 		if err != nil {
+			return nil, err
+		}
+		authErr := l.auth(conn)
+		if authErr != nil {
 			// TODO log
-			break
+			conn.Close()
+			continue
 		}
 
-		go func() {
-			authErr := l.auth(conn)
-			if authErr != nil {
-				// TODO log
-				conn.Close()
-				return
-			}
-
-			spdyConn, spdyErr := spdystream.NewConnection(conn, true)
-			if spdyErr != nil {
-				// TODO log
-				conn.Close()
-				return
-			}
-
-			go spdyConn.Serve(l.streamHandler)
-		}()
+		return newStreamSession(conn, true)
 	}
 }
-
-// Close closes the underlying listener
-func (l *ListenSession) Close() error {
-	return l.listener.Close()
-}
-
-// AcceptReceiver waits for a stream to be created in
-// the session and returns a receiver for that stream.
-func (l *ListenSession) AcceptReceiver() (libchan.Receiver, error) {
-	stream := <-l.streamChan
-	return &StreamReceiver{stream: stream, streamChans: l, ret: &StreamSender{stream: stream, streamChans: l}}, nil
-}

http2/listener_test.go

@@ -14,17 +14,19 @@ func TestListenSession(t *testing.T) {
 		t.Fatalf("Error creating listener: %s", listenErr)
 	}
 
-	session, sessionErr := NewListenSession(listener, NoAuthenticator)
-	if sessionErr != nil {
-		t.Fatalf("Error creating session: %s", sessionErr)
+	listenSession, listenSessionErr := NewListenSession(listener, NoAuthenticator)
+	if listenSessionErr != nil {
+		t.Fatalf("Error creating session: %s", listenSessionErr)
 	}
 
-	go session.Serve()
-
 	end := make(chan bool)
 	go exerciseServer(t, listen, end)
 
-	receiver, receiverErr := session.AcceptReceiver()
+	session, sessionErr := listenSession.AcceptSession()
+	if sessionErr != nil {
+		t.Fatalf("Error accepting session: %s", sessionErr)
+	}
+	receiver, receiverErr := session.ReceiverWait()
 	if receiverErr != nil {
 		t.Fatalf("Error accepting receiver: %s", receiverErr)
 	}

http2/stream.go

@@ -16,24 +16,31 @@ import (
 type StreamSession struct {
 	conn *spdystream.Connection
 
-	streamLock     sync.Mutex
+	streamCond     *sync.Cond
 	streamChan     chan *spdystream.Stream
 	subStreamChans map[string]chan *spdystream.Stream
 }
 
 func (s *StreamSession) addStreamChan(stream *spdystream.Stream, streamChan chan *spdystream.Stream) {
+	s.streamCond.L.Lock()
+	defer s.streamCond.L.Unlock()
 	s.subStreamChans[stream.String()] = streamChan
+	s.streamCond.Broadcast()
 }
 
 func (s *StreamSession) getStreamChan(stream *spdystream.Stream) chan *spdystream.Stream {
 	if stream == nil {
 		return s.streamChan
 	}
+	s.streamCond.L.Lock()
+	defer s.streamCond.L.Unlock()
 	streamChan, ok := s.subStreamChans[stream.String()]
-	if ok {
-		return streamChan
+	for !ok {
+		// TODO Test for stream being closed
+		s.streamCond.Wait()
+		streamChan, ok = s.subStreamChans[stream.String()]
 	}
-	return s.streamChan
+	return streamChan
 }
 
 func (s *StreamSession) newStreamHandler(stream *spdystream.Stream) {
@@ -44,14 +51,21 @@ func (s *StreamSession) newStreamHandler(stream *spdystream.Stream) {
 
 // NewStreamSession creates a new stream session from the
 // provided network connection.  The network connection is
-// expected to already provide a tls session.
+// expected to already provide a tls session.  NewStreamSession
+// should only be called on client connections, server connections
+// should be created through ListenSession.
 func NewStreamSession(conn net.Conn) (*StreamSession, error) {
+	return newStreamSession(conn, false)
+}
+
+func newStreamSession(conn net.Conn, server bool) (*StreamSession, error) {
 	session := &StreamSession{
 		streamChan:     make(chan *spdystream.Stream),
 		subStreamChans: make(map[string]chan *spdystream.Stream),
+		streamCond:     sync.NewCond(new(sync.Mutex)),
 	}
 
-	spdyConn, spdyErr := spdystream.NewConnection(conn, false)
+	spdyConn, spdyErr := spdystream.NewConnection(conn, server)
 	if spdyErr != nil {
 		return nil, spdyErr
 	}
@@ -76,9 +90,21 @@ func (s *StreamSession) NewSender() (libchan.Sender, error) {
 	if waitErr != nil {
 		return nil, waitErr
 	}
+
+	streamChan := make(chan *spdystream.Stream)
+	s.addStreamChan(stream, streamChan)
 	return &StreamSender{stream: stream, streamChans: s}, nil
 }
 
+// ReceiverWaits waits a new stream to be created and returns a
+// receiver with that stream.
+func (s *StreamSession) ReceiverWait() (libchan.Receiver, error) {
+	stream := <-s.streamChan
+	streamChan := make(chan *spdystream.Stream)
+	s.addStreamChan(stream, streamChan)
+	return &StreamReceiver{stream: stream, streamChans: s, ret: &StreamSender{stream: stream, streamChans: s}}, nil
+}
+
 // Close closes the underlying network connection, causing
 // each stream created by this session to closed.
 func (s *StreamSession) Close() error {
@@ -115,6 +141,7 @@ func (s *StreamReceiver) Receive(mode int) (*libchan.Message, error) {
 	} else {
 		streamChan := s.streamChans.getStreamChan(s.stream)
 		stream := <-streamChan
+		s.streamChans.addStreamChan(stream, make(chan *spdystream.Stream))
 
 		data, dataErr := extractDataHeader(stream.Headers())
 		if dataErr != nil {

http2/stream_test.go

@@ -103,17 +103,20 @@ func runServer(listen string, t *testing.T, endChan chan bool) (io.Closer, error
 		return nil, lErr
 	}
 
-	session, sessionErr := NewListenSession(listener, NoAuthenticator)
-	if sessionErr != nil {
-		t.Fatalf("Error creating session: %s", sessionErr)
+	listenSession, listenSessionErr := NewListenSession(listener, NoAuthenticator)
+	if listenSessionErr != nil {
+		t.Fatalf("Error creating session listener: %s", listenSessionErr)
 	}
 
-	go session.Serve()
-
 	go func() {
 		defer close(endChan)
 
-		receiver, receiverErr := session.AcceptReceiver()
+		session, sessionErr := listenSession.AcceptSession()
+		if sessionErr != nil {
+			t.Fatalf("Error accepting session: %s", sessionErr)
+		}
+
+		receiver, receiverErr := session.ReceiverWait()
 		if receiverErr != nil {
 			t.Fatalf("Error accepting receiver: %s", receiverErr)
 		}