Compare commits
6 Commits
0.0.1
...
de23b3e815
| Author | SHA1 | Date | |
|---|---|---|---|
| de23b3e815 | |||
| 2e73689762 | |||
| d5de31eda7 | |||
| b8714e52de | |||
| f5741ef60b | |||
| 31add1984b |
1
.gitignore
vendored
1
.gitignore
vendored
@@ -1 +1,2 @@
|
||||
cover.html
|
||||
.DS_Store
|
||||
|
||||
120
client/client.go
120
client/client.go
@@ -2,17 +2,15 @@ package client
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"sync"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"reichard.io/conduit/config"
|
||||
"reichard.io/conduit/types"
|
||||
"reichard.io/conduit/tunnel"
|
||||
)
|
||||
|
||||
func NewTunnel(cfg *config.ClientConfig) (*Tunnel, error) {
|
||||
func NewTunnel(cfg *config.ClientConfig) (*tunnel.Tunnel, error) {
|
||||
// Parse Server URL
|
||||
serverURL, err := url.Parse(cfg.ServerAddress)
|
||||
if err != nil {
|
||||
@@ -43,117 +41,5 @@ func NewTunnel(cfg *config.ClientConfig) (*Tunnel, error) {
|
||||
return nil, fmt.Errorf("failed to connect: %v", err)
|
||||
}
|
||||
|
||||
return &Tunnel{
|
||||
name: cfg.TunnelName,
|
||||
target: cfg.TunnelTarget,
|
||||
serverURL: serverURL,
|
||||
serverConn: serverConn,
|
||||
localConns: make(map[string]net.Conn),
|
||||
}, nil
|
||||
|
||||
}
|
||||
|
||||
type Tunnel struct {
|
||||
name string
|
||||
target string
|
||||
serverURL *url.URL
|
||||
|
||||
serverConn *websocket.Conn
|
||||
localConns map[string]net.Conn
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func (t *Tunnel) Start() error {
|
||||
log.Infof("starting tunnel: %s.%s -> %s\n", t.name, t.serverURL.Hostname(), t.target)
|
||||
defer t.serverConn.Close()
|
||||
|
||||
// Handle Messages
|
||||
for {
|
||||
// Read Message
|
||||
var msg types.Message
|
||||
err := t.serverConn.ReadJSON(&msg)
|
||||
if err != nil {
|
||||
log.Errorf("error reading from tunnel: %v", err)
|
||||
break
|
||||
}
|
||||
|
||||
switch msg.Type {
|
||||
case types.MessageTypeData:
|
||||
localConn, err := t.getLocalConn(msg.StreamID)
|
||||
if err != nil {
|
||||
log.Errorf("failed to get local connection: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Write data to local connection
|
||||
if _, err := localConn.Write(msg.Data); err != nil {
|
||||
log.Errorf("error writing to local connection: %v", err)
|
||||
localConn.Close()
|
||||
t.mu.Lock()
|
||||
delete(t.localConns, msg.StreamID)
|
||||
t.mu.Unlock()
|
||||
}
|
||||
|
||||
case types.MessageTypeClose:
|
||||
t.mu.Lock()
|
||||
if localConn, exists := t.localConns[msg.StreamID]; exists {
|
||||
localConn.Close()
|
||||
delete(t.localConns, msg.StreamID)
|
||||
}
|
||||
t.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Tunnel) getLocalConn(streamID string) (net.Conn, error) {
|
||||
// Get Cached Connection
|
||||
t.mu.RLock()
|
||||
localConn, exists := t.localConns[streamID]
|
||||
t.mu.RUnlock()
|
||||
if exists {
|
||||
return localConn, nil
|
||||
}
|
||||
|
||||
// Initiate Connection & Cache
|
||||
localConn, err := net.Dial("tcp", t.target)
|
||||
if err != nil {
|
||||
log.Errorf("failed to connect to %s: %v", t.target, err)
|
||||
return nil, err
|
||||
}
|
||||
t.mu.Lock()
|
||||
t.localConns[streamID] = localConn
|
||||
t.mu.Unlock()
|
||||
|
||||
// Start Response Relay & Return Connection
|
||||
go t.startResponseRelay(streamID, localConn)
|
||||
return localConn, nil
|
||||
}
|
||||
|
||||
func (t *Tunnel) startResponseRelay(streamID string, localConn net.Conn) {
|
||||
defer func() {
|
||||
t.mu.Lock()
|
||||
delete(t.localConns, streamID)
|
||||
t.mu.Unlock()
|
||||
localConn.Close()
|
||||
}()
|
||||
|
||||
buffer := make([]byte, 4096)
|
||||
for {
|
||||
n, err := localConn.Read(buffer)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
|
||||
response := types.Message{
|
||||
Type: types.MessageTypeData,
|
||||
StreamID: streamID,
|
||||
Data: buffer[:n],
|
||||
}
|
||||
|
||||
if err := t.serverConn.WriteJSON(response); err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
return tunnel.NewClientTunnel(cfg.TunnelName, cfg.TunnelTarget, serverConn)
|
||||
}
|
||||
|
||||
@@ -22,12 +22,7 @@ var tunnelCmd = &cobra.Command{
|
||||
if err != nil {
|
||||
log.Fatal("failed to create tunnel:", err)
|
||||
}
|
||||
|
||||
// Start Tunnel
|
||||
log.Infof("creating TCP tunnel: %s -> %s", cfg.TunnelName, cfg.TunnelTarget)
|
||||
if err := tunnel.Start(); err != nil {
|
||||
log.Fatal("failed to start tunnel:", err)
|
||||
}
|
||||
tunnel.Start()
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -111,7 +111,7 @@ func GetVersion() string {
|
||||
func getConfigValue(cmdFlags *pflag.FlagSet, def ConfigDef) string {
|
||||
// 1. Get Flags First
|
||||
if cmdFlags != nil {
|
||||
if val, err := cmdFlags.GetString(def.Key); err == nil && val != "" {
|
||||
if val, err := cmdFlags.GetString(def.Key); err == nil && val != "" && val != def.Default {
|
||||
return val
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,25 +7,25 @@ import (
|
||||
"net/http"
|
||||
)
|
||||
|
||||
var _ http.ResponseWriter = (*connResponseWriter)(nil)
|
||||
var _ http.ResponseWriter = (*rawHTTPResponseWriter)(nil)
|
||||
|
||||
type connResponseWriter struct {
|
||||
type rawHTTPResponseWriter struct {
|
||||
conn net.Conn
|
||||
header http.Header
|
||||
}
|
||||
|
||||
func (f *connResponseWriter) Header() http.Header {
|
||||
func (f *rawHTTPResponseWriter) Header() http.Header {
|
||||
if f.header == nil {
|
||||
f.header = make(http.Header)
|
||||
}
|
||||
return f.header
|
||||
}
|
||||
|
||||
func (f *connResponseWriter) Write(data []byte) (int, error) {
|
||||
func (f *rawHTTPResponseWriter) Write(data []byte) (int, error) {
|
||||
return f.conn.Write(data)
|
||||
}
|
||||
|
||||
func (f *connResponseWriter) WriteHeader(statusCode int) {
|
||||
func (f *rawHTTPResponseWriter) WriteHeader(statusCode int) {
|
||||
// Write Status
|
||||
status := fmt.Sprintf("HTTP/1.1 %d %s\r\n", statusCode, http.StatusText(statusCode))
|
||||
_, _ = f.conn.Write([]byte(status))
|
||||
@@ -41,7 +41,7 @@ func (f *connResponseWriter) WriteHeader(statusCode int) {
|
||||
_, _ = f.conn.Write([]byte("\r\n"))
|
||||
}
|
||||
|
||||
func (f *connResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||||
func (f *rawHTTPResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||||
// Return Raw Connection & ReadWriter
|
||||
rw := bufio.NewReadWriter(bufio.NewReader(f.conn), bufio.NewWriter(f.conn))
|
||||
return f.conn, rw, nil
|
||||
30
server/reconstructed_conn.go
Normal file
30
server/reconstructed_conn.go
Normal file
@@ -0,0 +1,30 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net"
|
||||
)
|
||||
|
||||
var _ io.ReadWriteCloser = (*reconstructedConn)(nil)
|
||||
|
||||
// reconstructedConn wraps a net.Conn and overrides Read to handle captured data.
|
||||
type reconstructedConn struct {
|
||||
net.Conn
|
||||
reader io.Reader
|
||||
}
|
||||
|
||||
// Read reads from the reconstructed reader (captured data + original conn).
|
||||
func (rc *reconstructedConn) Read(p []byte) (n int, err error) {
|
||||
return rc.reader.Read(p)
|
||||
}
|
||||
|
||||
// newReconstructedConn creates a reconstructed connection that replays captured data
|
||||
// before reading from the original connection.
|
||||
func newReconstructedConn(conn net.Conn, capturedData *bytes.Buffer) net.Conn {
|
||||
allReader := io.MultiReader(capturedData, conn)
|
||||
return &reconstructedConn{
|
||||
Conn: conn,
|
||||
reader: allReader,
|
||||
}
|
||||
}
|
||||
163
server/server.go
163
server/server.go
@@ -17,7 +17,7 @@ import (
|
||||
"github.com/gorilla/websocket"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"reichard.io/conduit/config"
|
||||
"reichard.io/conduit/types"
|
||||
"reichard.io/conduit/tunnel"
|
||||
)
|
||||
|
||||
type InfoResponse struct {
|
||||
@@ -30,19 +30,13 @@ type TunnelInfo struct {
|
||||
Target string `json:"target"`
|
||||
}
|
||||
|
||||
type TunnelConnection struct {
|
||||
*websocket.Conn
|
||||
name string
|
||||
streams map[string]chan []byte
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
host string
|
||||
cfg *config.ServerConfig
|
||||
mu sync.RWMutex
|
||||
|
||||
upgrader websocket.Upgrader
|
||||
tunnels map[string]*TunnelConnection
|
||||
tunnels map[string]*tunnel.Tunnel
|
||||
}
|
||||
|
||||
func NewServer(cfg *config.ServerConfig) (*Server, error) {
|
||||
@@ -56,7 +50,7 @@ func NewServer(cfg *config.ServerConfig) (*Server, error) {
|
||||
return &Server{
|
||||
cfg: cfg,
|
||||
host: serverURL.Host,
|
||||
tunnels: make(map[string]*TunnelConnection),
|
||||
tunnels: make(map[string]*tunnel.Tunnel),
|
||||
upgrader: websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true
|
||||
@@ -79,7 +73,7 @@ func (s *Server) Start() error {
|
||||
for {
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
log.Printf("error accepting connection: %v", err)
|
||||
log.WithError(err).Error("error accepting connection")
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -94,7 +88,7 @@ func (s *Server) getInfo(w http.ResponseWriter, _ *http.Request) {
|
||||
for t, c := range s.tunnels {
|
||||
allTunnels = append(allTunnels, TunnelInfo{
|
||||
Name: t,
|
||||
Target: c.RemoteAddr().String(),
|
||||
Target: c.Source(),
|
||||
})
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
@@ -105,72 +99,17 @@ func (s *Server) getInfo(w http.ResponseWriter, _ *http.Request) {
|
||||
Version: config.GetVersion(),
|
||||
}, "", " ")
|
||||
if err != nil {
|
||||
log.WithError(err).Error("failed to marshal info")
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Send Response
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(d)
|
||||
}
|
||||
|
||||
func (s *Server) proxyRawConnection(clientConn net.Conn, tunnelConn *TunnelConnection, dataReader io.Reader) {
|
||||
defer clientConn.Close()
|
||||
|
||||
// Create Identifiers
|
||||
streamID := fmt.Sprintf("stream_%d", time.Now().UnixNano())
|
||||
responseChan := make(chan []byte, 100)
|
||||
|
||||
// Register Stream
|
||||
s.mu.Lock()
|
||||
if tunnelConn.streams == nil {
|
||||
tunnelConn.streams = make(map[string]chan []byte)
|
||||
}
|
||||
tunnelConn.streams[streamID] = responseChan
|
||||
s.mu.Unlock()
|
||||
|
||||
// Clean Up
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
delete(tunnelConn.streams, streamID)
|
||||
close(responseChan)
|
||||
s.mu.Unlock()
|
||||
|
||||
// Send Close
|
||||
closeMsg := types.Message{
|
||||
Type: types.MessageTypeClose,
|
||||
StreamID: streamID,
|
||||
}
|
||||
_ = tunnelConn.WriteJSON(closeMsg)
|
||||
}()
|
||||
|
||||
// Read & Send Chunks
|
||||
go func() {
|
||||
buffer := make([]byte, 4096)
|
||||
for {
|
||||
n, err := dataReader.Read(buffer)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
if err := tunnelConn.WriteJSON(types.Message{
|
||||
Type: types.MessageTypeData,
|
||||
StreamID: streamID,
|
||||
Data: buffer[:n],
|
||||
}); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Return Response Data
|
||||
for data := range responseChan {
|
||||
if _, err := clientConn.Write(data); err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleRawConnection(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
|
||||
@@ -183,7 +122,7 @@ func (s *Server) handleRawConnection(conn net.Conn) {
|
||||
bufReader := bufio.NewReader(teeReader)
|
||||
|
||||
// Create HTTP Request & Writer
|
||||
w := &connResponseWriter{conn: conn}
|
||||
w := &rawHTTPResponseWriter{conn: conn}
|
||||
r, err := http.ReadRequest(bufReader)
|
||||
if err != nil {
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
@@ -214,15 +153,25 @@ func (s *Server) handleRawConnection(conn net.Conn) {
|
||||
|
||||
// Handle Tunnels
|
||||
s.mu.RLock()
|
||||
tunnelConn, exists := s.tunnels[subdomain]
|
||||
conduitTunnel, exists := s.tunnels[subdomain]
|
||||
s.mu.RUnlock()
|
||||
if exists {
|
||||
log.Infof("relaying %s to tunnel", subdomain)
|
||||
|
||||
// Reconstruct Data & Proxy Connection
|
||||
allReader := io.MultiReader(&capturedData, r.Body)
|
||||
s.proxyRawConnection(conn, tunnelConn, allReader)
|
||||
if !exists {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
_, _ = fmt.Fprintf(w, "unknown tunnel: %s", subdomain)
|
||||
return
|
||||
}
|
||||
|
||||
// Add & Start Stream
|
||||
reconstructedConn := newReconstructedConn(conn, &capturedData)
|
||||
streamID := fmt.Sprintf("stream_%d", time.Now().UnixNano())
|
||||
if err := conduitTunnel.AddStream(streamID, reconstructedConn); err != nil {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = fmt.Fprintf(w, "failed to add stream: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
log.Infof("relaying %s to tunnel", subdomain)
|
||||
_ = conduitTunnel.StartStream(streamID)
|
||||
}
|
||||
|
||||
func (s *Server) handleAsHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -245,40 +194,6 @@ func (s *Server) handleAsHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleTunnelMessages(tunnel *TunnelConnection) {
|
||||
for {
|
||||
var msg types.Message
|
||||
err := tunnel.ReadJSON(&msg)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
if msg.StreamID == "" {
|
||||
log.Infof("tunnel %s missing streamID", tunnel.name)
|
||||
continue
|
||||
}
|
||||
|
||||
switch msg.Type {
|
||||
case types.MessageTypeClose:
|
||||
return
|
||||
case types.MessageTypeData:
|
||||
s.mu.RLock()
|
||||
streamChan, exists := tunnel.streams[msg.StreamID]
|
||||
if !exists {
|
||||
log.Infof("stream %s does not exist", msg.StreamID)
|
||||
s.mu.RUnlock()
|
||||
continue
|
||||
}
|
||||
|
||||
select {
|
||||
case streamChan <- msg.Data:
|
||||
case <-time.After(time.Second):
|
||||
log.Warnf("stream %s channel full, dropping data", msg.StreamID)
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
func (s *Server) createTunnel(w http.ResponseWriter, r *http.Request) {
|
||||
// Get Tunnel Name
|
||||
tunnelName := r.URL.Query().Get("tunnelName")
|
||||
@@ -302,26 +217,20 @@ func (s *Server) createTunnel(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
// Create & Cache TunnelConnection
|
||||
tunnel := &TunnelConnection{
|
||||
Conn: wsConn,
|
||||
name: tunnelName,
|
||||
streams: make(map[string]chan []byte),
|
||||
}
|
||||
// Create Tunnel
|
||||
conduitTunnel := tunnel.NewServerTunnel(tunnelName, wsConn)
|
||||
s.mu.Lock()
|
||||
s.tunnels[tunnelName] = tunnel
|
||||
s.tunnels[tunnelName] = conduitTunnel
|
||||
s.mu.Unlock()
|
||||
log.Infof("tunnel established: %s", tunnelName)
|
||||
|
||||
// Keep connection alive and handle cleanup
|
||||
defer func() {
|
||||
s.mu.Lock()
|
||||
delete(s.tunnels, tunnelName)
|
||||
s.mu.Unlock()
|
||||
_ = wsConn.Close()
|
||||
log.Infof("tunnel closed: %s", tunnelName)
|
||||
}()
|
||||
// Start Tunnel - This is blocking
|
||||
conduitTunnel.Start()
|
||||
|
||||
// Handle tunnel messages
|
||||
s.handleTunnelMessages(tunnel)
|
||||
// Cleanup Tunnel
|
||||
s.mu.Lock()
|
||||
delete(s.tunnels, tunnelName)
|
||||
s.mu.Unlock()
|
||||
_ = wsConn.Close()
|
||||
log.Infof("tunnel closed: %s", tunnelName)
|
||||
}
|
||||
|
||||
104
tunnel/http.go
Normal file
104
tunnel/http.go
Normal file
@@ -0,0 +1,104 @@
|
||||
package tunnel
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/url"
|
||||
"sync"
|
||||
)
|
||||
|
||||
func HTTPConnectionBuilder(targetURL *url.URL) (ConnBuilder, error) {
|
||||
multiConnListener := newMultiConnListener()
|
||||
|
||||
// Create Reverse Proxy
|
||||
proxy := &httputil.ReverseProxy{
|
||||
Director: func(req *http.Request) {
|
||||
req.Host = targetURL.Host
|
||||
req.URL.Host = targetURL.Host
|
||||
req.URL.Scheme = targetURL.Scheme
|
||||
},
|
||||
ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
|
||||
http.Error(w, fmt.Sprintf("Proxy error: %v", err), http.StatusBadGateway)
|
||||
},
|
||||
}
|
||||
|
||||
// Start HTTP Proxy
|
||||
go func() {
|
||||
defer multiConnListener.Close()
|
||||
_ = http.Serve(multiConnListener, proxy)
|
||||
}()
|
||||
|
||||
// Return Connection Builder
|
||||
return func() (conn io.ReadWriteCloser, err error) {
|
||||
clientConn, serverConn := net.Pipe()
|
||||
|
||||
if err := multiConnListener.addConn(serverConn); err != nil {
|
||||
_ = clientConn.Close()
|
||||
_ = serverConn.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return clientConn, nil
|
||||
}, nil
|
||||
}
|
||||
|
||||
type multiConnListener struct {
|
||||
connCh chan net.Conn
|
||||
closed chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func newMultiConnListener() *multiConnListener {
|
||||
return &multiConnListener{
|
||||
connCh: make(chan net.Conn, 100),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *multiConnListener) Accept() (net.Conn, error) {
|
||||
select {
|
||||
case conn := <-l.connCh:
|
||||
if conn == nil {
|
||||
return nil, fmt.Errorf("listener closed")
|
||||
}
|
||||
return conn, nil
|
||||
case <-l.closed:
|
||||
return nil, fmt.Errorf("listener closed")
|
||||
}
|
||||
}
|
||||
|
||||
func (l *multiConnListener) Close() error {
|
||||
l.once.Do(func() {
|
||||
close(l.closed)
|
||||
// Drain any remaining connections
|
||||
go func() {
|
||||
for conn := range l.connCh {
|
||||
if conn != nil {
|
||||
conn.Close()
|
||||
}
|
||||
}
|
||||
}()
|
||||
close(l.connCh)
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *multiConnListener) Addr() net.Addr {
|
||||
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
|
||||
}
|
||||
|
||||
func (l *multiConnListener) addConn(conn net.Conn) error {
|
||||
select {
|
||||
case l.connCh <- conn:
|
||||
return nil
|
||||
case <-l.closed:
|
||||
conn.Close()
|
||||
return fmt.Errorf("listener is closed")
|
||||
default:
|
||||
conn.Close()
|
||||
return fmt.Errorf("connection queue full")
|
||||
}
|
||||
}
|
||||
200
tunnel/tunnel.go
Normal file
200
tunnel/tunnel.go
Normal file
@@ -0,0 +1,200 @@
|
||||
package tunnel
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/url"
|
||||
"sync"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"reichard.io/conduit/types"
|
||||
)
|
||||
|
||||
type ConnBuilder func() (conn io.ReadWriteCloser, err error)
|
||||
|
||||
func NewServerTunnel(name string, wsConn *websocket.Conn) *Tunnel {
|
||||
return &Tunnel{
|
||||
name: name,
|
||||
wsConn: wsConn,
|
||||
streams: make(map[string]io.ReadWriteCloser),
|
||||
}
|
||||
}
|
||||
|
||||
func NewClientTunnel(name, target string, wsConn *websocket.Conn) (*Tunnel, error) {
|
||||
targetURL, err := url.Parse(target)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var connBuilder ConnBuilder
|
||||
switch targetURL.Scheme {
|
||||
case "http", "https":
|
||||
log.Infof("creating HTTP tunnel: %s -> %s", name, target)
|
||||
connBuilder, err = HTTPConnectionBuilder(targetURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
default:
|
||||
log.Infof("creating TCP tunnel: %s -> %s", name, target)
|
||||
connBuilder = func() (conn io.ReadWriteCloser, err error) {
|
||||
return net.Dial("tcp", target)
|
||||
}
|
||||
}
|
||||
|
||||
return &Tunnel{
|
||||
name: name,
|
||||
wsConn: wsConn,
|
||||
streams: make(map[string]io.ReadWriteCloser),
|
||||
connBuilder: connBuilder,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type Tunnel struct {
|
||||
name string
|
||||
wsConn *websocket.Conn
|
||||
streams map[string]io.ReadWriteCloser
|
||||
connBuilder ConnBuilder
|
||||
|
||||
wsMu, streamsMu sync.Mutex
|
||||
}
|
||||
|
||||
func (t *Tunnel) Start() {
|
||||
for {
|
||||
var msg types.Message
|
||||
err := t.wsConn.ReadJSON(&msg)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Validate Stream
|
||||
if msg.StreamID == "" {
|
||||
log.Warnf("tunnel %s missing streamID", t.name)
|
||||
continue
|
||||
}
|
||||
|
||||
// Ensure Stream
|
||||
if err := t.initStreamConnection(msg.StreamID); err != nil {
|
||||
log.WithError(err).Errorf("failed to initialize stream %s connection", t.name)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle Messages
|
||||
switch msg.Type {
|
||||
case types.MessageTypeClose:
|
||||
_ = t.CloseStream(msg.StreamID)
|
||||
case types.MessageTypeData:
|
||||
_ = t.WriteStream(msg.StreamID, msg.Data)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Tunnel) initStreamConnection(streamID string) error {
|
||||
if t.connBuilder == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if _, found := t.getStream(streamID); found {
|
||||
return nil
|
||||
}
|
||||
|
||||
conn, err := t.connBuilder()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := t.AddStream(streamID, conn); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
go t.StartStream(streamID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Tunnel) AddStream(streamID string, conn io.ReadWriteCloser) error {
|
||||
t.streamsMu.Lock()
|
||||
defer t.streamsMu.Unlock()
|
||||
|
||||
if _, found := t.streams[streamID]; found {
|
||||
return fmt.Errorf("stream %s already exists", streamID)
|
||||
}
|
||||
t.streams[streamID] = conn
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Tunnel) StartStream(streamID string) error {
|
||||
// Get Stream
|
||||
conn, found := t.getStream(streamID)
|
||||
if !found {
|
||||
return fmt.Errorf("stream %s does not exist", streamID)
|
||||
}
|
||||
|
||||
// Close Stream
|
||||
defer func() {
|
||||
_ = t.sendWS(&types.Message{
|
||||
Type: types.MessageTypeClose,
|
||||
StreamID: streamID,
|
||||
})
|
||||
|
||||
t.CloseStream(streamID)
|
||||
}()
|
||||
|
||||
// Start Stream
|
||||
buffer := make([]byte, 4096)
|
||||
for {
|
||||
n, err := conn.Read(buffer)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := t.sendWS(&types.Message{
|
||||
Type: types.MessageTypeData,
|
||||
Data: buffer[:n],
|
||||
StreamID: streamID,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Tunnel) WriteStream(streamID string, data []byte) error {
|
||||
// Get Stream
|
||||
conn, found := t.getStream(streamID)
|
||||
if !found {
|
||||
return fmt.Errorf("stream %s does not exist", streamID)
|
||||
}
|
||||
|
||||
_, err := conn.Write(data)
|
||||
return err
|
||||
}
|
||||
|
||||
func (t *Tunnel) CloseStream(streamID string) error {
|
||||
t.streamsMu.Lock()
|
||||
defer t.streamsMu.Unlock()
|
||||
if conn, ok := t.streams[streamID]; ok {
|
||||
delete(t.streams, streamID)
|
||||
return conn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Tunnel) Source() string {
|
||||
return t.wsConn.RemoteAddr().String()
|
||||
}
|
||||
|
||||
func (t *Tunnel) sendWS(msg *types.Message) error {
|
||||
t.wsMu.Lock()
|
||||
defer t.wsMu.Unlock()
|
||||
return t.wsConn.WriteJSON(msg)
|
||||
}
|
||||
|
||||
func (t *Tunnel) getStream(streamID string) (io.ReadWriteCloser, bool) {
|
||||
t.streamsMu.Lock()
|
||||
defer t.streamsMu.Unlock()
|
||||
|
||||
if conn, ok := t.streams[streamID]; ok {
|
||||
return conn, true
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
Reference in New Issue
Block a user