From 4a962d826477ace329ae23adb66700ca0da74a6b Mon Sep 17 00:00:00 2001 From: samosvalishe Date: Wed, 20 May 2026 11:55:35 +0300 Subject: [PATCH] =?UTF-8?q?refactor(bondframe):=20=D0=B2=D1=8B=D0=BD=D0=B5?= =?UTF-8?q?=D1=81=D1=82=D0=B8=20Reorder;=20=D1=80=D0=B0=D0=B7=D0=B4=D0=B5?= =?UTF-8?q?=D0=BB=D0=B8=D1=82=D1=8C=20=D0=BC=D0=B5=D0=B6=D0=B4=D1=83=20bon?= =?UTF-8?q?dclient=20=D0=B8=20bondserver?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/proxy/bondclient/bondclient.go | 58 +++---------------- internal/proxy/bondserver/bondserver.go | 58 +++---------------- internal/wire/bondframe/bondframe.go | 76 +++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 102 deletions(-) diff --git a/internal/proxy/bondclient/bondclient.go b/internal/proxy/bondclient/bondclient.go index 8be3fd8..51a16d9 100644 --- a/internal/proxy/bondclient/bondclient.go +++ b/internal/proxy/bondclient/bondclient.go @@ -216,55 +216,11 @@ func writeBondFrameToNextLane(ctx context.Context, lanes []*lane, typ byte, seq } func (h *Handler) copyBondToTCP(ctx context.Context, connID uint64, tcpConn net.Conn, recvCh <-chan bondframe.Frame) { - pending := make(map[uint64][]byte) - var expect uint64 - var finSeq *uint64 - - for { - if finSeq != nil && expect == *finSeq { - bondframe.CloseWrite(tcpConn, h.Deps.log().Debugf) - h.Deps.log().Debugf("[bond %d] download finished chunks=%d", connID, expect) - return - } - - select { - case <-ctx.Done(): - return - case f, ok := <-recvCh: - if !ok { - return - } - switch f.Type { - case bondframe.FrameData: - if len(pending) >= 1024 { - h.Deps.log().Errorf("[bond %d] pending map overflow (>1024), closing", connID) - return - } - pending[f.Seq] = f.Data - case bondframe.FrameFIN: - v := f.Seq - if finSeq == nil || v < *finSeq { - finSeq = &v - } - default: - h.Deps.log().Errorf("[bond %d] unknown frame type %d", connID, f.Type) - return - } - - for { - data, ok := pending[expect] - if !ok { - break - } - delete(pending, expect) - if len(data) > 0 { - if _, err := tcpConn.Write(data); err != nil { - h.Deps.log().Errorf("[bond %d] local TCP write error: %v", connID, err) - return - } - } - expect++ - } - } - } + chunks := bondframe.Reorder(ctx, tcpConn, recvCh, bondframe.ReorderHooks{ + OnOverflow: func(have int) { h.Deps.log().Errorf("[bond %d] pending map overflow (>%d), closing", connID, bondframe.PendingCap) }, + OnUnknownType: func(typ byte) { h.Deps.log().Errorf("[bond %d] unknown frame type %d", connID, typ) }, + OnWriteError: func(err error) { h.Deps.log().Errorf("[bond %d] local TCP write error: %v", connID, err) }, + OnCloseWrite: h.Deps.log().Debugf, + }) + h.Deps.log().Debugf("[bond %d] download finished chunks=%d", connID, chunks) } diff --git a/internal/proxy/bondserver/bondserver.go b/internal/proxy/bondserver/bondserver.go index 409f90d..5a80722 100644 --- a/internal/proxy/bondserver/bondserver.go +++ b/internal/proxy/bondserver/bondserver.go @@ -265,57 +265,13 @@ func (c *serverConn) run() { } func (c *serverConn) copyBondToBackend(backendConn net.Conn) { - pending := make(map[uint64][]byte) - var expect uint64 - var finSeq *uint64 - - for { - if finSeq != nil && expect == *finSeq { - bondframe.CloseWrite(backendConn, c.deps.log().Debugf) - c.deps.log().Debugf("[bond %d] upload to backend finished chunks=%d", c.id, expect) - return - } - - select { - case <-c.ctx.Done(): - return - case f := <-c.recvCh: - switch f.Type { - case bondframe.FrameData: - // Bound the reorder buffer so a malicious or buggy peer - // that emits seq with permanent gaps cannot exhaust memory. - // Matches bondclient.copyBondToTCP. - if len(pending) >= 1024 { - c.deps.log().Errorf("[bond %d] pending map overflow (>1024), closing", c.id) - return - } - pending[f.Seq] = f.Data - case bondframe.FrameFIN: - v := f.Seq - if finSeq == nil || v < *finSeq { - finSeq = &v - } - default: - c.deps.log().Errorf("[bond %d] unknown frame type %d", c.id, f.Type) - return - } - - for { - data, ok := pending[expect] - if !ok { - break - } - delete(pending, expect) - if len(data) > 0 { - if _, err := backendConn.Write(data); err != nil { - c.deps.log().Errorf("[bond %d] backend write error: %v", c.id, err) - return - } - } - expect++ - } - } - } + chunks := bondframe.Reorder(c.ctx, backendConn, c.recvCh, bondframe.ReorderHooks{ + OnOverflow: func(have int) { c.deps.log().Errorf("[bond %d] pending map overflow (>%d), closing", c.id, bondframe.PendingCap) }, + OnUnknownType: func(typ byte) { c.deps.log().Errorf("[bond %d] unknown frame type %d", c.id, typ) }, + OnWriteError: func(err error) { c.deps.log().Errorf("[bond %d] backend write error: %v", c.id, err) }, + OnCloseWrite: c.deps.log().Debugf, + }) + c.deps.log().Debugf("[bond %d] upload to backend finished chunks=%d", c.id, chunks) } func (c *serverConn) copyBackendToBond(backendConn net.Conn) { diff --git a/internal/wire/bondframe/bondframe.go b/internal/wire/bondframe/bondframe.go index 3f6f631..e8adec7 100644 --- a/internal/wire/bondframe/bondframe.go +++ b/internal/wire/bondframe/bondframe.go @@ -4,6 +4,7 @@ package bondframe import ( + "context" "encoding/binary" "fmt" "io" @@ -117,6 +118,81 @@ func ReadFrame(r io.Reader) (Frame, error) { return f, nil } +// PendingCap bounds the per-bond reorder buffer. A peer that emits seq with +// permanent gaps cannot grow it past this many frames. +const PendingCap = 1024 + +// ReorderHooks plugs caller-specific logging into Reorder. All fields may be nil. +type ReorderHooks struct { + OnOverflow func(have int) + OnUnknownType func(typ byte) + OnWriteError func(err error) + OnCloseWrite func(format string, v ...any) +} + +// Reorder consumes Frames from recv, writes payloads to dst in Seq order, and +// returns the number of fully delivered chunks. It returns when: +// - the FIN seq has been reached (CloseWrite is called on dst); +// - recv is closed; +// - ctx is cancelled; +// - an unknown frame type, pending overflow, or write error fires. +func Reorder(ctx context.Context, dst net.Conn, recv <-chan Frame, h ReorderHooks) uint64 { + pending := make(map[uint64][]byte) + var expect uint64 + var finSeq *uint64 + + for { + if finSeq != nil && expect == *finSeq { + CloseWrite(dst, h.OnCloseWrite) + return expect + } + select { + case <-ctx.Done(): + return expect + case f, ok := <-recv: + if !ok { + return expect + } + switch f.Type { + case FrameData: + if len(pending) >= PendingCap { + if h.OnOverflow != nil { + h.OnOverflow(len(pending)) + } + return expect + } + pending[f.Seq] = f.Data + case FrameFIN: + v := f.Seq + if finSeq == nil || v < *finSeq { + finSeq = &v + } + default: + if h.OnUnknownType != nil { + h.OnUnknownType(f.Type) + } + return expect + } + for { + data, ok := pending[expect] + if !ok { + break + } + delete(pending, expect) + if len(data) > 0 { + if _, err := dst.Write(data); err != nil { + if h.OnWriteError != nil { + h.OnWriteError(err) + } + return expect + } + } + expect++ + } + } + } +} + // CloseWrite half-closes the write side of conn if the underlying type // supports it (TCPConn, smux.Stream, …); otherwise it is a no-op. errf is // invoked with the error if CloseWrite fails; callers typically pass a