telebit/mplexer/websockettunnel.go

148 lines
4.1 KiB
Go

package telebit
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)
}