From 31cf571a6aac3dcfe77cc35c328af967b799d44a Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Wed, 3 Jun 2026 23:33:27 +0800 Subject: [PATCH] refactor(handler/router): extract conn.go (packetConn, lockWriter) from associate.go --- handler/router/associate.go | 66 -------------------------------- handler/router/conn.go | 75 +++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 66 deletions(-) create mode 100644 handler/router/conn.go diff --git a/handler/router/associate.go b/handler/router/associate.go index 864645e6..9bfff18c 100644 --- a/handler/router/associate.go +++ b/handler/router/associate.go @@ -3,13 +3,9 @@ package router import ( "bytes" "context" - "encoding/binary" - "errors" "fmt" "io" - "math" "net" - "sync" "time" "github.com/go-gost/core/common/bufpool" @@ -294,65 +290,3 @@ func (h *routerHandler) getAddrforRoute(ctx context.Context, routerID, gateway s return raddr } -type packetConn struct { - net.Conn -} - -func (c *packetConn) Read(b []byte) (n int, err error) { - var bb [2]byte - _, err = io.ReadFull(c.Conn, bb[:]) - if err != nil { - return - } - - dlen := int(binary.BigEndian.Uint16(bb[:])) - if len(b) >= dlen { - return io.ReadFull(c.Conn, b[:dlen]) - } - - buf := bufpool.Get(dlen) - defer bufpool.Put(buf) - - n, err = io.ReadFull(c.Conn, buf) - copy(b, buf[:n]) - - return -} - -func (c *packetConn) Write(b []byte) (n int, err error) { - if len(b) > math.MaxUint16 { - err = errors.New("write: data maximum exceeded") - return - } - - buf := bufpool.Get(len(b) + 2) - defer bufpool.Put(buf) - - binary.BigEndian.PutUint16(buf[:2], uint16(len(b))) - n = copy(buf[2:], b) - - return c.Conn.Write(buf) -} - -type lockWriter struct { - w io.Writer - mu sync.Mutex -} - -func LockWriter(w io.Writer) io.Writer { - return &lockWriter{w: w} -} - -func (w *lockWriter) Write(p []byte) (int, error) { - w.mu.Lock() - defer w.mu.Unlock() - - return w.w.Write(p) -} - -func (w *lockWriter) Close() error { - if closer, ok := w.w.(io.Closer); ok { - return closer.Close() - } - return nil -} diff --git a/handler/router/conn.go b/handler/router/conn.go new file mode 100644 index 00000000..214d954b --- /dev/null +++ b/handler/router/conn.go @@ -0,0 +1,75 @@ +package router + +import ( + "encoding/binary" + "errors" + "io" + "math" + "net" + "sync" + + "github.com/go-gost/core/common/bufpool" +) + +type packetConn struct { + net.Conn +} + +func (c *packetConn) Read(b []byte) (n int, err error) { + var bb [2]byte + _, err = io.ReadFull(c.Conn, bb[:]) + if err != nil { + return + } + + dlen := int(binary.BigEndian.Uint16(bb[:])) + if len(b) >= dlen { + return io.ReadFull(c.Conn, b[:dlen]) + } + + buf := bufpool.Get(dlen) + defer bufpool.Put(buf) + + n, err = io.ReadFull(c.Conn, buf) + copy(b, buf[:n]) + + return +} + +func (c *packetConn) Write(b []byte) (n int, err error) { + if len(b) > math.MaxUint16 { + err = errors.New("write: data maximum exceeded") + return + } + + buf := bufpool.Get(len(b) + 2) + defer bufpool.Put(buf) + + binary.BigEndian.PutUint16(buf[:2], uint16(len(b))) + n = copy(buf[2:], b) + + return c.Conn.Write(buf) +} + +type lockWriter struct { + w io.Writer + mu sync.Mutex +} + +func LockWriter(w io.Writer) io.Writer { + return &lockWriter{w: w} +} + +func (w *lockWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + + return w.w.Write(p) +} + +func (w *lockWriter) Close() error { + if closer, ok := w.w.(io.Closer); ok { + return closer.Close() + } + return nil +}