package packer import ( "context" "fmt" "io" "net" "net/http" "os" "strings" "time" "github.com/gorilla/websocket" ) // WebsocketTunnel wraps a websocket.Conn instance to behave like net.Conn. // TODO make conform. type WebsocketTunnel struct { wsconn WSConn tmpr io.Reader //w io.WriteCloser //pingCh chan struct{} } // WSConn defines a interface for gorilla websockets for the purpose of testing type WSConn interface { NextReader() (messageType int, r io.Reader, err error) NextWriter(messageType int) (io.WriteCloser, error) WriteControl(messageType int, data []byte, deadline time.Time) error WriteMessage(messageType int, data []byte) error SetReadDeadline(t time.Time) error SetWriteDeadline(t time.Time) error Close() error RemoteAddr() net.Addr // LocalAddr() net.Addr } // NewWebsocketTunnel allocates a new websocket connection wrapper func NewWebsocketTunnel(wsconn WSConn) net.Conn { return &WebsocketTunnel{ wsconn: wsconn, tmpr: nil, } } // DialWebsocketTunnel connects to the given websocket relay as wraps it as net.Conn func DialWebsocketTunnel(ctx context.Context, relay, authz string) (net.Conn, error) { wsd := websocket.Dialer{} headers := http.Header{} headers.Set("Authorization", fmt.Sprintf("Bearer %s", authz)) // *http.Response sep := "?" if strings.Contains(relay, sep) { sep = "&" } wsconn, _, err := wsd.DialContext(ctx, relay+sep+"access_token="+authz+"&versions=v1", headers) return NewWebsocketTunnel(wsconn), err } func (wsw *WebsocketTunnel) Read(b []byte) (int, error) { if nil == wsw.tmpr { _, msgr, err := wsw.wsconn.NextReader() if nil != err { //fmt.Println("debug wsw NextReader err:", err) return 0, err } wsw.tmpr = msgr } n, err := wsw.tmpr.Read(b) if nil != err { //fmt.Println("debug wsw Read err:", err) if io.EOF == err { wsw.tmpr = nil // ignore the message EOF because it's not the websocket EOF err = nil } } return n, err } func (wsw *WebsocketTunnel) Write(b []byte) (int, error) { // TODO create or reset ping deadline // TODO document that more complete writes are preferred? msgw, err := wsw.wsconn.NextWriter(websocket.BinaryMessage) if nil != err { //fmt.Println("debug wsw NextWriter err:", err) return 0, err } n, err := msgw.Write(b) if nil != err { //fmt.Println("debug wsw Write err:", err) return n, err } // if the message error fails, we can assume the websocket is damaged return n, msgw.Close() } // Close will close the websocket with a control message func (wsw *WebsocketTunnel) Close() error { //fmt.Println("[debug] closing the websocket.Conn") // TODO handle EOF as websocket.CloseNormal? message := websocket.FormatCloseMessage(websocket.CloseGoingAway, "closing connection") deadline := time.Now().Add(10 * time.Second) err := wsw.wsconn.WriteControl(websocket.CloseMessage, message, deadline) if nil != err { fmt.Fprintf(os.Stderr, "failed to write close message to websocket: %s\n", err) } _ = wsw.wsconn.Close() return err } // LocalAddr is not implemented and will panic func (wsw *WebsocketTunnel) LocalAddr() net.Addr { // TODO do we reverse this since the "local" address is that of the relay? // return wsw.wsconn.RemoteAddr() panic("LocalAddr() not implemented") } // RemoteAddr is not implemented and will panic. Additionally, it wouldn't mean anything useful anyway. func (wsw *WebsocketTunnel) RemoteAddr() net.Addr { // TODO do we reverse this since the "remote" address means nothing / is that of one of the clients? // return wsw.wsconn.LocalAddr() panic("RemoteAddr() not implemented") } // SetDeadline sets the read and write deadlines associated func (wsw *WebsocketTunnel) SetDeadline(t time.Time) error { err := wsw.SetReadDeadline(t) if nil == err { err = wsw.SetWriteDeadline(t) } return err } // SetReadDeadline sets the deadline for future Read calls func (wsw *WebsocketTunnel) SetReadDeadline(t time.Time) error { return wsw.wsconn.SetReadDeadline(t) } // SetWriteDeadline sets the deadline for future Write calls func (wsw *WebsocketTunnel) SetWriteDeadline(t time.Time) error { return wsw.wsconn.SetWriteDeadline(t) }