github.com/xmplusdev/xmcore@v1.8.11-0.20240412132628-5518b55526af/common/buf/override.go (about)

     1  package buf
     2  
     3  import (
     4  	"github.com/xmplusdev/xmcore/common/net"
     5  )
     6  
     7  type EndpointOverrideReader struct {
     8  	Reader
     9  	Dest         net.Address
    10  	OriginalDest net.Address
    11  }
    12  
    13  func (r *EndpointOverrideReader) ReadMultiBuffer() (MultiBuffer, error) {
    14  	mb, err := r.Reader.ReadMultiBuffer()
    15  	if err == nil {
    16  		for _, b := range mb {
    17  			if b.UDP != nil && b.UDP.Address == r.OriginalDest {
    18  				b.UDP.Address = r.Dest
    19  			}
    20  		}
    21  	}
    22  	return mb, err
    23  }
    24  
    25  type EndpointOverrideWriter struct {
    26  	Writer
    27  	Dest         net.Address
    28  	OriginalDest net.Address
    29  }
    30  
    31  func (w *EndpointOverrideWriter) WriteMultiBuffer(mb MultiBuffer) error {
    32  	for _, b := range mb {
    33  		if b.UDP != nil && b.UDP.Address == w.Dest {
    34  			b.UDP.Address = w.OriginalDest
    35  		}
    36  	}
    37  	return w.Writer.WriteMultiBuffer(mb)
    38  }