mirror of
https://github.com/Gouryella/drip.git
synced 2026-02-28 07:11:40 +00:00
feat: Add Bearer Token authentication support and optimize code structure
- Add Bearer Token authentication, supporting tunnel access control via the --auth-bearer parameter - Refactor large modules into smaller, more focused components to improve code maintainability - Update dependency versions, including golang.org/x/crypto, golang.org/x/net, etc. - Add SilenceUsage and SilenceErrors configuration for all CLI commands - Modify connector configuration structure to support the new authentication method - Update recent change log in README with new feature descriptions BREAKING CHANGE: Authentication via Bearer Token is now supported, requiring the new --auth-bearer parameter
This commit is contained in:
113
internal/server/proxy/websocket_handler.go
Normal file
113
internal/server/proxy/websocket_handler.go
Normal file
@@ -0,0 +1,113 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"drip/internal/server/tunnel"
|
||||
"drip/internal/shared/httputil"
|
||||
"drip/internal/shared/netutil"
|
||||
"drip/internal/shared/protocol"
|
||||
"drip/internal/shared/wsutil"
|
||||
)
|
||||
|
||||
type bufferedReadWriteCloser struct {
|
||||
*bufio.Reader
|
||||
net.Conn
|
||||
}
|
||||
|
||||
func (b *bufferedReadWriteCloser) Read(p []byte) (int, error) {
|
||||
return b.Reader.Read(p)
|
||||
}
|
||||
|
||||
func (h *Handler) handleWebSocket(w http.ResponseWriter, r *http.Request, tconn *tunnel.Connection) {
|
||||
stream, err := h.openStreamWithTimeout(tconn)
|
||||
if err != nil {
|
||||
http.Error(w, "Tunnel unavailable", http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
|
||||
tconn.IncActiveConnections()
|
||||
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
stream.Close()
|
||||
tconn.DecActiveConnections()
|
||||
http.Error(w, "WebSocket not supported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
clientConn, clientBuf, err := hj.Hijack()
|
||||
if err != nil {
|
||||
stream.Close()
|
||||
tconn.DecActiveConnections()
|
||||
http.Error(w, "Failed to hijack connection", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
if err := r.Write(stream); err != nil {
|
||||
stream.Close()
|
||||
clientConn.Close()
|
||||
tconn.DecActiveConnections()
|
||||
return
|
||||
}
|
||||
|
||||
go func() {
|
||||
defer stream.Close()
|
||||
defer clientConn.Close()
|
||||
defer tconn.DecActiveConnections()
|
||||
|
||||
var clientRW io.ReadWriteCloser = clientConn
|
||||
if clientBuf != nil && clientBuf.Reader.Buffered() > 0 {
|
||||
clientRW = &bufferedReadWriteCloser{
|
||||
Reader: clientBuf.Reader,
|
||||
Conn: clientConn,
|
||||
}
|
||||
}
|
||||
|
||||
_ = netutil.PipeWithCallbacks(context.Background(), stream, clientRW,
|
||||
func(n int64) { tconn.AddBytesOut(n) },
|
||||
func(n int64) { tconn.AddBytesIn(n) },
|
||||
)
|
||||
}()
|
||||
}
|
||||
|
||||
func (h *Handler) handleTunnelWebSocket(w http.ResponseWriter, r *http.Request) {
|
||||
if !h.IsTransportAllowed("wss") {
|
||||
http.Error(w, "WebSocket transport not allowed on this server", http.StatusForbidden)
|
||||
return
|
||||
}
|
||||
|
||||
if h.wsConnHandler == nil {
|
||||
http.Error(w, "WebSocket tunnel not configured", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
ws, err := h.wsUpgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
h.logger.Error("WebSocket upgrade failed", zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
ws.SetReadLimit(protocol.MaxFrameSize + protocol.FrameHeaderSize + 1024)
|
||||
|
||||
remoteAddr := netutil.ExtractClientIP(r)
|
||||
|
||||
h.logger.Info("WebSocket tunnel connection established",
|
||||
zap.String("remote_addr", remoteAddr),
|
||||
)
|
||||
|
||||
conn := wsutil.NewConnWithPing(ws, 30*time.Second)
|
||||
|
||||
h.wsConnHandler.HandleWSConnection(conn, remoteAddr)
|
||||
}
|
||||
|
||||
func (h *Handler) isWebSocketUpgrade(r *http.Request) bool {
|
||||
return httputil.IsWebSocketUpgrade(r)
|
||||
}
|
||||
Reference in New Issue
Block a user