// Dispatch processes the connection.
func (d *dispatcher) Dispatch(g grim.GrimReaper, conn net.Conn) {
// Create buffer and read connection type
buf := make([]byte, 1)
if _, err := conn.Read(buf); err != nil {
d.logger.Warn("error reading StreamType",
"remote", conn.RemoteAddr().String, "err", err)
conn.Close()
return
}
// Dispatch connection based on first byte
ctype := StreamType(buf[0])
if handler, ok := d.handlers[ctype]; ok {
g.SpawnFunc(func(c context.Context) {
handler.Handle(c, conn)
})
return
}
// Process unknown streams if there is an unknown stream handler
if d.unknown != nil {
g.SpawnFunc(func(c context.Context) {
d.unknown.Handle(c, conn)
})
return
}
// If there is not an unknown stream handler, close the stream.
d.logger.Warn("closing connection: unknown stream type",
"type", ctype, "remote", conn.RemoteAddr().String())
conn.Close()
}
func processStreams(g grim.GrimReaper, conn net.Conn, streamCh chan net.Conn, dispatcher Dispatcher) grim.TaskFunc {
return func(ctx context.Context) {
defer conn.Close()
for {
select {
case <-ctx.Done():
return
case stream, ok := <-streamCh:
if !ok {
return
}
dispatcher.Dispatch(g.New(), stream)
}
}
}
}
func (y *yamuxer) handleConn(g grim.GrimReaper, conn net.Conn) {
defer g.Wait()
conf := yamux.DefaultConfig()
conf.LogOutput = y.logOutput
session, _ := yamux.Server(conn, conf)
streamCh := make(chan net.Conn)
g.SpawnFunc(processStreams(g.New(), conn, streamCh, y.dispatcher))
g.SpawnFunc(acceptStreams(y.logger, session, streamCh))
}