moved websocket wrapper to iowrappers package.

This commit is contained in:
Erik Brakkee 2024-07-18 21:14:23 +02:00
parent a6a0f287dc
commit 6b9dc2bfbe
3 changed files with 49 additions and 60 deletions

View File

@ -1,19 +1,18 @@
package main package main
import ( import (
"fmt" "cidebug/pkg/iowrappers"
"github.com/gorilla/websocket"
"io" "io"
"log" "log"
"net" "net"
"github.com/gorilla/websocket"
) )
func handleConnection(tcpConn net.Conn, wsURL string) { func handleConnection(tcpConn net.Conn, wsURL string) {
defer tcpConn.Close() defer tcpConn.Close()
conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
wsConn := NewWebSocketConn(conn) wsConn := iowrappers.NewWebSocketConn(conn)
if err != nil { if err != nil {
log.Println("WebSocket connection error:", err) log.Println("WebSocket connection error:", err)
return return
@ -26,29 +25,18 @@ func handleConnection(tcpConn net.Conn, wsURL string) {
defer func() { defer func() {
waitChannel <- true waitChannel <- true
}() }()
_, _ = io.Copy(tcpConn, wsConn) _, err = io.Copy(tcpConn, wsConn)
// TODO print error if err != nil {
log.Printf("error %v", err)
}
}() }()
go func() { go func() {
defer func() { defer func() {
waitChannel <- true waitChannel <- true
}() }()
for { _, err = io.Copy(wsConn, tcpConn)
buffer := make([]byte, 1024) log.Printf("Error %v", err)
n, err := tcpConn.Read(buffer)
if err != nil {
if err != io.EOF {
log.Println("TCP read error:", err)
}
return
}
err = conn.WriteMessage(websocket.BinaryMessage, buffer[:n])
if err != nil {
log.Println("WebSocket write error:", err)
return
}
}
}() }()
<-waitChannel <-waitChannel
@ -76,41 +64,3 @@ func main() {
go handleConnection(conn, wsURL) go handleConnection(conn, wsURL)
} }
} }
type WebSocketConn struct {
conn *websocket.Conn
buf []byte
}
func NewWebSocketConn(conn *websocket.Conn) *WebSocketConn {
return &WebSocketConn{conn: conn}
}
func (websocketConn *WebSocketConn) Read(p []byte) (n int, err error) {
if len(websocketConn.buf) == 0 {
_, message, err := websocketConn.conn.ReadMessage()
fmt.Println("Got message ", message)
if err != nil {
return 0, err
}
websocketConn.buf = message
}
n = copy(p, websocketConn.buf)
websocketConn.buf = websocketConn.buf[n:]
log.Println("Read ", n, " bytes")
return n, err
}
func (websocketConn *WebSocketConn) Write(p []byte) (n int, err error) {
err = websocketConn.conn.WriteMessage(websocket.BinaryMessage, p)
if err != nil {
n = len(p)
}
return n, err
}
func (websocketConn *WebSocketConn) Close() error {
return websocketConn.conn.Close()
}

2
go.mod
View File

@ -1,4 +1,4 @@
module ssh module cidebug
go 1.18 go 1.18

39
pkg/iowrappers/io.go Normal file
View File

@ -0,0 +1,39 @@
package iowrappers
import "github.com/gorilla/websocket"
type WebSocketConn struct {
conn *websocket.Conn
buf []byte
}
func NewWebSocketConn(conn *websocket.Conn) *WebSocketConn {
return &WebSocketConn{conn: conn}
}
func (websocketConn *WebSocketConn) Read(p []byte) (n int, err error) {
if len(websocketConn.buf) == 0 {
_, message, err := websocketConn.conn.ReadMessage()
if err != nil {
return 0, err
}
websocketConn.buf = message
}
n = copy(p, websocketConn.buf)
websocketConn.buf = websocketConn.buf[n:]
return n, err
}
func (websocketConn *WebSocketConn) Write(p []byte) (n int, err error) {
err = websocketConn.conn.WriteMessage(websocket.BinaryMessage, p)
if err == nil {
n = len(p)
}
return n, err
}
func (websocketConn *WebSocketConn) Close() error {
return websocketConn.conn.Close()
}