docs/beam/ws/ws.go

73 lines
1.8 KiB
Go

package ws
import (
"errors"
"github.com/docker/libswarm/beam"
"github.com/docker/libswarm/beam/http2"
"github.com/docker/spdystream/ws"
"github.com/gorilla/websocket"
"net/http"
)
// Connect to a Beam server over a Websocket connection as a client
func NewSender(wsConn *websocket.Conn) (beam.Sender, error) {
session, err := http2.NewStreamSession(ws.NewConnection(wsConn))
if err != nil {
return nil, err
}
return session, nil
}
// Upgrade an HTTP connection to a Beam over HTTP2 over
// Websockets connection.
type Upgrader struct {
Upgrader websocket.Upgrader
}
func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*http2.Server, error) {
wsConn, err := u.Upgrader.Upgrade(w, r, responseHeader)
if err != nil {
return nil, err
}
netConn := ws.NewConnection(wsConn)
server, err := http2.NewServer(netConn)
if err != nil {
netConn.Close()
return nil, err
}
return server, nil
}
// Returns true if a handshake error occured in websockets, which means
// a response has already been written to the stream.
func IsHandshakeError(err error) bool {
_, ok := err.(websocket.HandshakeError)
return ok
}
type BeamFunc func(beam.Receiver)
// Handler function for serving Beam over HTTP. Will invoke f and
// then close the server's Beam endpoint after f returns.
func Serve(u *Upgrader, f BeamFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
if r.Method != "GET" {
u.Upgrader.Error(w, r, http.StatusMethodNotAllowed, errors.New("Method not allowed"))
return
}
server, err := u.Upgrade(w, r, nil)
if err != nil {
if !IsHandshakeError(err) {
u.Upgrader.Error(w, r, http.StatusInternalServerError, errors.New("Unable to open an HTTP2 connection over Websockets"))
}
return
}
defer server.Close()
f(server)
}
}