Skip to content
This repository has been archived by the owner on May 26, 2022. It is now read-only.

Commit

Permalink
Merge pull request #41 from libp2p/fix/modernize
Browse files Browse the repository at this point in the history
modernize request handling
  • Loading branch information
Stebalien committed Apr 4, 2019
2 parents b9aa343 + 1161ec4 commit 93daaa6
Show file tree
Hide file tree
Showing 5 changed files with 26 additions and 43 deletions.
1 change: 0 additions & 1 deletion .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ env:
global:
- GOTFLAGS="-race"
matrix:
- BUILD_DEPTYPE=gx
- BUILD_DEPTYPE=gomod


Expand Down
18 changes: 8 additions & 10 deletions conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ var _ net.Conn = (*Conn)(nil)
type Conn struct {
*ws.Conn
DefaultMessageType int
done func()
reader io.Reader
closeOnce sync.Once
}
Expand Down Expand Up @@ -85,14 +84,14 @@ func (c *Conn) Write(b []byte) (n int, err error) {
func (c *Conn) Close() error {
var err error
c.closeOnce.Do(func() {
if c.done != nil {
c.done()
// Be nice to GC
c.done = nil
err1 := c.Conn.WriteControl(ws.CloseMessage, nil, time.Now().Add(GracefulCloseTimeout))
err2 := c.Conn.Close()
switch {
case err1 != nil:
err = err1
case err2 != nil:
err = err2
}

c.Conn.WriteControl(ws.CloseMessage, nil, time.Now().Add(GracefulCloseTimeout))
err = c.Conn.Close()
})
return err
}
Expand Down Expand Up @@ -122,10 +121,9 @@ func (c *Conn) SetWriteDeadline(t time.Time) error {
}

// NewConn creates a Conn given a regular gorilla/websocket Conn.
func NewConn(raw *ws.Conn, done func()) *Conn {
func NewConn(raw *ws.Conn) *Conn {
return &Conn{
Conn: raw,
DefaultMessageType: ws.BinaryMessage,
done: done,
}
}
30 changes: 3 additions & 27 deletions listener.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package websocket

import (
"context"
"fmt"
"net"
"net/http"
Expand All @@ -21,7 +20,7 @@ type listener struct {

func (l *listener) serve() {
defer close(l.closed)
http.Serve(l.Listener, l)
_ = http.Serve(l.Listener, l)
}

func (l *listener) ServeHTTP(w http.ResponseWriter, r *http.Request) {
Expand All @@ -31,35 +30,12 @@ func (l *listener) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}

ctx, cancel := context.WithCancel(context.Background())

var cnCh <-chan bool
if cn, ok := w.(http.CloseNotifier); ok {
cnCh = cn.CloseNotify()
}

wscon := NewConn(c, cancel)
// Just to make sure.
defer wscon.Close()

select {
case l.incoming <- wscon:
case l.incoming <- NewConn(c):
case <-l.closed:
c.Close()
return
case <-cnCh:
return
}

// wait until conn gets closed, otherwise the handler closes it early
select {
case <-ctx.Done():
case <-l.closed:
c.Close()
return
case <-cnCh:
return
}
// The connection has been hijacked, it's safe to return.
}

func (l *listener) Accept() (manet.Conn, error) {
Expand Down
2 changes: 1 addition & 1 deletion websocket.go
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ func (t *WebsocketTransport) maDial(ctx context.Context, raddr ma.Multiaddr) (ma
return nil, err
}

mnc, err := manet.WrapNetConn(NewConn(wscon, nil))
mnc, err := manet.WrapNetConn(NewConn(wscon))
if err != nil {
wscon.Close()
return nil, err
Expand Down
18 changes: 14 additions & 4 deletions websocket_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,14 @@ func TestWebsocketListen(t *testing.T) {
return
}

c.Write(msg)
c.Close()
_, err = c.Write(msg)
if err != nil {
t.Error(err)
}
err = c.Close()
if err != nil {
t.Error(err)
}
}()

c, err := l.Accept()
Expand Down Expand Up @@ -120,8 +126,12 @@ func TestConcurrentClose(t *testing.T) {
return
}

go c.Write(msg)
go c.Close()
go func() {
_, _ = c.Write(msg)
}()
go func() {
_ = c.Close()
}()
}
}()

Expand Down

0 comments on commit 93daaa6

Please sign in to comment.