135 lines
2.8 KiB
Go
135 lines
2.8 KiB
Go
package channel
|
|
|
|
import (
|
|
"encoding/base64"
|
|
"net"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
func Wrap(conn Connection, subprotocol string) Connection {
|
|
switch subprotocol {
|
|
case "channel.k8s.io":
|
|
return &kubeWrapper{base64: false, conn: conn}
|
|
case "base64.channel.k8s.io":
|
|
return &kubeWrapper{base64: true, conn: conn}
|
|
case "terminal.gitlab.com":
|
|
return &gitlabWrapper{base64: false, conn: conn}
|
|
case "base64.terminal.gitlab.com":
|
|
return &gitlabWrapper{base64: true, conn: conn}
|
|
}
|
|
|
|
return conn
|
|
}
|
|
|
|
type kubeWrapper struct {
|
|
base64 bool
|
|
conn Connection
|
|
}
|
|
|
|
type gitlabWrapper struct {
|
|
base64 bool
|
|
conn Connection
|
|
}
|
|
|
|
func (w *gitlabWrapper) ReadMessage() (int, []byte, error) {
|
|
mt, data, err := w.conn.ReadMessage()
|
|
if err != nil {
|
|
return mt, data, err
|
|
}
|
|
|
|
if isData(mt) {
|
|
mt = websocket.BinaryMessage
|
|
if w.base64 {
|
|
data, err = decodeBase64(data)
|
|
}
|
|
}
|
|
|
|
return mt, data, err
|
|
}
|
|
|
|
func (w *gitlabWrapper) WriteMessage(mt int, data []byte) error {
|
|
if isData(mt) {
|
|
if w.base64 {
|
|
mt = websocket.TextMessage
|
|
data = encodeBase64(data)
|
|
} else {
|
|
mt = websocket.BinaryMessage
|
|
}
|
|
}
|
|
|
|
return w.conn.WriteMessage(mt, data)
|
|
}
|
|
|
|
func (w *gitlabWrapper) WriteControl(mt int, data []byte, deadline time.Time) error {
|
|
return w.conn.WriteControl(mt, data, deadline)
|
|
}
|
|
|
|
func (w *gitlabWrapper) UnderlyingConn() net.Conn {
|
|
return w.conn.UnderlyingConn()
|
|
}
|
|
|
|
// Coalesces all wsstreams into a single stream. In practice, we should only
|
|
// receive data on stream 1.
|
|
func (w *kubeWrapper) ReadMessage() (int, []byte, error) {
|
|
mt, data, err := w.conn.ReadMessage()
|
|
if err != nil {
|
|
return mt, data, err
|
|
}
|
|
|
|
if isData(mt) {
|
|
mt = websocket.BinaryMessage
|
|
|
|
// Remove the WSStream channel number, decode to raw
|
|
if len(data) > 0 {
|
|
data = data[1:]
|
|
if w.base64 {
|
|
data, err = decodeBase64(data)
|
|
}
|
|
}
|
|
}
|
|
|
|
return mt, data, err
|
|
}
|
|
|
|
// Always sends to wsstream 0
|
|
func (w *kubeWrapper) WriteMessage(mt int, data []byte) error {
|
|
if isData(mt) {
|
|
if w.base64 {
|
|
mt = websocket.TextMessage
|
|
data = append([]byte{'0'}, encodeBase64(data)...)
|
|
} else {
|
|
mt = websocket.BinaryMessage
|
|
data = append([]byte{0}, data...)
|
|
}
|
|
}
|
|
|
|
return w.conn.WriteMessage(mt, data)
|
|
}
|
|
|
|
func (w *kubeWrapper) WriteControl(mt int, data []byte, deadline time.Time) error {
|
|
return w.conn.WriteControl(mt, data, deadline)
|
|
}
|
|
|
|
func (w *kubeWrapper) UnderlyingConn() net.Conn {
|
|
return w.conn.UnderlyingConn()
|
|
}
|
|
|
|
func isData(mt int) bool {
|
|
return mt == websocket.BinaryMessage || mt == websocket.TextMessage
|
|
}
|
|
|
|
func encodeBase64(data []byte) []byte {
|
|
buf := make([]byte, base64.StdEncoding.EncodedLen(len(data)))
|
|
base64.StdEncoding.Encode(buf, data)
|
|
|
|
return buf
|
|
}
|
|
|
|
func decodeBase64(data []byte) ([]byte, error) {
|
|
buf := make([]byte, base64.StdEncoding.DecodedLen(len(data)))
|
|
n, err := base64.StdEncoding.Decode(buf, data)
|
|
return buf[:n], err
|
|
}
|