Docker/portainer
Docker/portainer
1
0
Fork 0
mirror of https://github.com/portainer/portainer.git synced 2025-07-19 13:29:41 +02:00
Code Issues Releases Activity

feat(websocket): improve websocket code sharing BE-11340 (#61)
Some checks failed
Label Conflicts / triage (push) Has been cancelled

Browse source
This commit is contained in:
andres-portainer 2024-10-25 11:21:49 -03:00 • committed by GitHub
parent b2d67795b3
commit 1d037f2f1f
6 changed files with 122 additions and 153 deletions
Download patch file Download diff file Expand all files Collapse all files

17
api/http/handler/websocket/attach.go
View file

@ -5,10 +5,8 @@ import (
"net/http"
"time"
"github.com/portainer/portainer/api/http/security"
"github.com/rs/zerolog/log"
portainer "github.com/portainer/portainer/api"
"github.com/portainer/portainer/api/ws"
httperror "github.com/portainer/portainer/pkg/libhttp/error"
"github.com/portainer/portainer/pkg/libhttp/request"
@ -76,14 +74,6 @@ func (handler *Handler) websocketAttach(w http.ResponseWriter, r *http.Request)
}
func (handler *Handler) handleAttachRequest(w http.ResponseWriter, r *http.Request, params *webSocketRequestParams) error {
tokenData, err := security.RetrieveTokenData(r)
if err != nil {
log.Warn().
Err(err).
Msg("unable to retrieve user details from authentication token")
return err
}
r.Header.Del("Origin")
if params.endpoint.Type == portainer.AgentOnDockerEnvironment {
@ -98,14 +88,13 @@ func (handler *Handler) handleAttachRequest(w http.ResponseWriter, r *http.Reque
}
defer websocketConn.Close()
return hijackAttachStartOperation(websocketConn, params.endpoint, params.ID, tokenData.Token)
return hijackAttachStartOperation(websocketConn, params.endpoint, params.ID)
}
func hijackAttachStartOperation(
websocketConn *websocket.Conn,
endpoint *portainer.Endpoint,
attachID string,
token string,
) error {
conn, err := initDial(endpoint)
if err != nil {
@ -127,7 +116,7 @@ func hijackAttachStartOperation(
return err
}
return hijackRequest(websocketConn, conn, attachStartRequest, token)
return ws.HijackRequest(websocketConn, conn, attachStartRequest)
}
func createAttachStartRequest(attachID string) (*http.Request, error) {

16
api/http/handler/websocket/exec.go
View file

@ -5,13 +5,12 @@ import (
"net/http"
portainer "github.com/portainer/portainer/api"
"github.com/portainer/portainer/api/http/security"
"github.com/portainer/portainer/api/ws"
httperror "github.com/portainer/portainer/pkg/libhttp/error"
"github.com/portainer/portainer/pkg/libhttp/request"
"github.com/asaskevich/govalidator"
"github.com/gorilla/websocket"
"github.com/rs/zerolog/log"
"github.com/segmentio/encoding/json"
)
@ -79,14 +78,6 @@ func (handler *Handler) websocketExec(w http.ResponseWriter, r *http.Request) *h
}
func (handler *Handler) handleExecRequest(w http.ResponseWriter, r *http.Request, params *webSocketRequestParams) error {
tokenData, err := security.RetrieveTokenData(r)
if err != nil {
log.Warn().
Err(err).
Msg("unable to retrieve user details from authentication token")
return err
}
r.Header.Del("Origin")
if params.endpoint.Type == portainer.AgentOnDockerEnvironment {
@ -102,14 +93,13 @@ func (handler *Handler) handleExecRequest(w http.ResponseWriter, r *http.Request
defer websocketConn.Close()
return hijackExecStartOperation(websocketConn, params.endpoint, params.ID, tokenData.Token)
return hijackExecStartOperation(websocketConn, params.endpoint, params.ID)
}
func hijackExecStartOperation(
websocketConn *websocket.Conn,
endpoint *portainer.Endpoint,
execID string,
token string,
) error {
conn, err := initDial(endpoint)
if err != nil {
@ -121,7 +111,7 @@ func hijackExecStartOperation(
return err
}
return hijackRequest(websocketConn, conn, execStartRequest, token)
return ws.HijackRequest(websocketConn, conn, execStartRequest)
}
func createExecStartRequest(execID string) (*http.Request, error) {

163
api/http/handler/websocket/hijack.go
View file

@ -1,163 +0,0 @@
package websocket
import (
"bufio"
"errors"
"fmt"
"io"
"net"
"net/http"
"sync"
"time"
"github.com/gorilla/websocket"
"github.com/rs/zerolog/log"
)
const (
// Time allowed to write a message to the peer
writeWait = 10 * time.Second
// Send pings to peer with this period
pingPeriod = 50 * time.Second
)
func hijackRequest(
websocketConn *websocket.Conn,
conn net.Conn,
request *http.Request,
token string,
) error {
resp, err := sendHTTPRequest(conn, request)
if err != nil {
return err
}
defer resp.Body.Close()
// Check if the response status code indicates an upgrade (101 Switching Protocols)
if resp.StatusCode != http.StatusSwitchingProtocols {
return fmt.Errorf("unexpected response status code: %d", resp.StatusCode)
}
errorChan := make(chan error, 1)
go readWebSocketToTCP(websocketConn, conn, errorChan)
go writeTCPToWebSocket(websocketConn, conn, errorChan)
err = <-errorChan
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
log.Debug().Msgf("Unexpected close error: %v\n", err)
return err
}
log.Debug().Msgf("session ended")
return nil
}
// sendHTTPRequest sends an HTTP request over the provided net.Conn and parses the response.
func sendHTTPRequest(conn net.Conn, req *http.Request) (*http.Response, error) {
// Send the HTTP request to the server
if err := req.Write(conn); err != nil {
return nil, fmt.Errorf("error writing request: %w", err)
}
// Read the response from the server
resp, err := http.ReadResponse(bufio.NewReader(conn), req)
if err != nil {
return nil, fmt.Errorf("error reading response: %w", err)
}
return resp, nil
}
func readWebSocketToTCP(websocketConn *websocket.Conn, tcpConn net.Conn, errorChan chan error) {
for {
messageType, p, err := websocketConn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
log.Debug().Msgf("Unexpected close error: %v\n", err)
}
errorChan <- err
return
}
if messageType == websocket.TextMessage || messageType == websocket.BinaryMessage {
_, err := tcpConn.Write(p)
if err != nil {
log.Debug().Msgf("Error writing to TCP connection: %v\n", err)
errorChan <- err
return
}
}
}
}
func writeTCPToWebSocket(websocketConn *websocket.Conn, tcpConn net.Conn, errorChan chan error) {
var mu sync.Mutex
out := make([]byte, readerBufferSize)
input := make(chan string)
pingTicker := time.NewTicker(pingPeriod)
defer pingTicker.Stop()
defer websocketConn.Close()
websocketConn.SetReadLimit(2048)
websocketConn.SetPongHandler(func(string) error {
return nil
})
websocketConn.SetPingHandler(func(data string) error {
websocketConn.SetWriteDeadline(time.Now().Add(writeWait))
return websocketConn.WriteMessage(websocket.PongMessage, []byte(data))
})
reader := bufio.NewReader(tcpConn)
go func() {
for {
n, err := reader.Read(out)
if err != nil {
errorChan <- err
if !errors.Is(err, io.EOF) {
log.Debug().Msgf("error reading from server: %v", err)
}
return
}
processedOutput := validString(string(out[:n]))
input <- processedOutput
}
}()
for {
select {
case msg := <-input:
err := wswrite(websocketConn, &mu, msg)
if err != nil {
log.Debug().Msgf("error writing to websocket: %v", err)
errorChan <- err
return
}
case <-pingTicker.C:
if err := wsping(websocketConn, &mu); err != nil {
log.Debug().Msgf("error writing to websocket during pong response: %v", err)
errorChan <- err
return
}
}
}
}
func wswrite(websocketConn *websocket.Conn, mu *sync.Mutex, msg string) error {
mu.Lock()
defer mu.Unlock()
websocketConn.SetWriteDeadline(time.Now().Add(writeWait))
return websocketConn.WriteMessage(websocket.TextMessage, []byte(msg))
}
func wsping(websocketConn *websocket.Conn, mu *sync.Mutex) error {
mu.Lock()
defer mu.Unlock()
websocketConn.SetWriteDeadline(time.Now().Add(writeWait))
return websocketConn.WriteMessage(websocket.PingMessage, nil)
}

5
api/http/handler/websocket/pod.go
View file

@ -9,6 +9,7 @@ import (
portainer "github.com/portainer/portainer/api"
"github.com/portainer/portainer/api/http/proxy/factory/kubernetes"
"github.com/portainer/portainer/api/http/security"
"github.com/portainer/portainer/api/ws"
httperror "github.com/portainer/portainer/pkg/libhttp/error"
"github.com/portainer/portainer/pkg/libhttp/request"
@ -136,8 +137,8 @@ func (handler *Handler) hijackPodExecStartOperation(
// errorChan is used to propagate errors from the go routines to the caller.
errorChan := make(chan error, 1)
go streamFromWebsocketToWriter(websocketConn, stdinWriter, errorChan)
go streamFromReaderToWebsocket(websocketConn, stdoutReader, errorChan)
go ws.StreamFromWebsocketToWriter(websocketConn, stdinWriter, errorChan)
go ws.StreamFromReaderToWebsocket(websocketConn, stdoutReader, errorChan)
// StartExecProcess is a blocking operation which streams IO to/from pod;
// this must execute in asynchronously, since the websocketConn could return errors (e.g. client disconnects) before

70
api/http/handler/websocket/stream.go
View file

@ -1,70 +0,0 @@
package websocket
import (
"io"
"unicode/utf8"
"github.com/gorilla/websocket"
)
const readerBufferSize = 2048
func streamFromWebsocketToWriter(websocketConn *websocket.Conn, writer io.Writer, errorChan chan error) {
for {
_, in, err := websocketConn.ReadMessage()
if err != nil {
errorChan <- err
break
}
_, err = writer.Write(in)
if err != nil {
errorChan <- err
break
}
}
}
func streamFromReaderToWebsocket(websocketConn *websocket.Conn, reader io.Reader, errorChan chan error) {
out := make([]byte, readerBufferSize)
for {
n, err := reader.Read(out)
if err != nil {
errorChan <- err
break
}
processedOutput := validString(string(out[:n]))
err = websocketConn.WriteMessage(websocket.TextMessage, []byte(processedOutput))
if err != nil {
errorChan <- err
break
}
}
}
func validString(s string) string {
if utf8.ValidString(s) {
return s
}
v := make([]rune, 0, len(s))
for i, r := range s {
if r == utf8.RuneError {
_, size := utf8.DecodeRuneInString(s[i:])
if size == 1 {
continue
}
}
v = append(v, r)
}
return string(v)
}