1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-22 01:57:12 -05:00

rewrite connection interface

This commit is contained in:
Shelikhoo 2021-12-30 15:12:32 +00:00 committed by Xiaokang Wang (Shelikhoo)
parent 899b8fefca
commit cdfef7e66b
3 changed files with 52 additions and 30 deletions

View File

@ -69,16 +69,17 @@ func (c *packetConnectionAdaptor) ReadFrom(p []byte) (n int, addr gonet.Addr, er
} }
} }
c.readerBuffer, n = buf.SplitFirstBytes(c.readerBuffer, p) c.readerBuffer, n = buf.SplitFirstBytes(c.readerBuffer, p)
var w []byte var w *buf.Buffer
w, addr = ExtractAddressFromPacket(p[:n]) w, addr, err = ExtractAddressFromPacket(buf.FromBytes(p[:n]))
n = copy(p, w) n = copy(p, w.Bytes())
w.Release()
return return
} }
func (c *packetConnectionAdaptor) WriteTo(p []byte, addr gonet.Addr) (n int, err error) { func (c *packetConnectionAdaptor) WriteTo(p []byte, addr gonet.Addr) (n int, err error) {
payloadLen := len(p) payloadLen := len(p)
p = AttachAddressToPacket(p, addr) var buffer *buf.Buffer
buffer := buf.FromBytes(p) buffer, err = AttachAddressToPacket(buf.FromBytes(p), addr)
mb := buf.MultiBuffer{buffer} mb := buf.MultiBuffer{buffer}
err = c.link.Writer.WriteMultiBuffer(mb) err = c.link.Writer.WriteMultiBuffer(mb)
if err != nil { if err != nil {
@ -134,18 +135,26 @@ func (pc *packetConnWrapper) Read(p []byte) (n int, err error) {
if err != nil { if err != nil {
return 0, err return 0, err
} }
result := AttachAddressToPacket(recbuf.Bytes()[0:n], addr) recbuf.Resize(0, int32(n))
n = copy(p, result) result, err := AttachAddressToPacket(&recbuf, addr)
recbuf.Release() if err != nil {
return 0, err
}
n = copy(p, result.Bytes())
result.Release()
return n, nil return n, nil
} }
func (pc *packetConnWrapper) Write(p []byte) (n int, err error) { func (pc *packetConnWrapper) Write(p []byte) (n int, err error) {
data, addr := ExtractAddressFromPacket(p) data, addr, err := ExtractAddressFromPacket(buf.FromBytes(p))
_, err = pc.PacketConn.WriteTo(data, addr)
if err != nil { if err != nil {
return 0, err return 0, err
} }
_, err = pc.PacketConn.WriteTo(data.Bytes(), addr)
if err != nil {
return 0, err
}
data.Release()
return len(p), nil return len(p), nil
} }

View File

@ -13,35 +13,43 @@ var addrParser = protocol.NewAddressParser(
protocol.AddressFamilyByte(0x02, net.AddressFamilyIPv6), protocol.AddressFamilyByte(0x02, net.AddressFamilyIPv6),
) )
func AttachAddressToPacket(data []byte, address gonet.Addr) []byte { // AttachAddressToPacket
packetBuf := buf.StackNew() // relinquish ownership of data
// gain ownership of the returning value
func AttachAddressToPacket(data *buf.Buffer, address gonet.Addr) (*buf.Buffer, error) {
packetBuf := buf.New()
udpaddr := address.(*gonet.UDPAddr) udpaddr := address.(*gonet.UDPAddr)
port, err := net.PortFromInt(uint32(udpaddr.Port)) port, err := net.PortFromInt(uint32(udpaddr.Port))
if err != nil { if err != nil {
panic(err) return nil, err
} }
err = addrParser.WriteAddressPort(&packetBuf, net.IPAddress(udpaddr.IP), port) err = addrParser.WriteAddressPort(packetBuf, net.IPAddress(udpaddr.IP), port)
if err != nil { if err != nil {
panic(err) return nil, err
} }
//Incorrect buffer reuse _, err = packetBuf.Write(data.Bytes())
data = append(packetBuf.Bytes(), data...) if err != nil {
//packetBuf.Release() return nil, err
return data }
data.Release()
return packetBuf, nil
} }
func ExtractAddressFromPacket(data []byte) ([]byte, gonet.Addr) { // ExtractAddressFromPacket
// relinquish ownership of data
// gain ownership of the returning value
func ExtractAddressFromPacket(data *buf.Buffer) (*buf.Buffer, gonet.Addr, error) {
packetBuf := buf.StackNew() packetBuf := buf.StackNew()
address, port, err := addrParser.ReadAddressPort(&packetBuf, bytes.NewReader(data)) address, port, err := addrParser.ReadAddressPort(&packetBuf, bytes.NewReader(data.Bytes()))
if err != nil { if err != nil {
panic(err) return nil, nil, err
} }
var addr = &gonet.UDPAddr{ var addr = &gonet.UDPAddr{
IP: address.IP(), IP: address.IP(),
Port: int(port.Value()), Port: int(port.Value()),
Zone: "", Zone: "",
} }
payload := data[int(packetBuf.Len()):] data.Advance(packetBuf.Len())
packetBuf.Release() packetBuf.Release()
return payload, addr return data, addr, nil
} }

View File

@ -2,6 +2,7 @@ package packetaddr
import ( import (
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/v2fly/v2ray-core/v4/common/buf"
sysnet "net" sysnet "net"
"testing" "testing"
) )
@ -12,11 +13,13 @@ func TestPacketEncodingIPv4(t *testing.T) {
Port: 1234, Port: 1234,
} }
var packetData [256]byte var packetData [256]byte
wrapped := AttachAddressToPacket(packetData[:], packetAddress) wrapped, err := AttachAddressToPacket(buf.FromBytes(packetData[:]), packetAddress)
assert.NoError(t, err)
packetPayload, decodedAddress := ExtractAddressFromPacket(wrapped) packetPayload, decodedAddress, err := ExtractAddressFromPacket(wrapped)
assert.NoError(t, err)
assert.Equal(t, packetPayload, packetData[:]) assert.Equal(t, packetPayload.Bytes(), packetData[:])
assert.Equal(t, packetAddress, decodedAddress) assert.Equal(t, packetAddress, decodedAddress)
} }
@ -26,10 +29,12 @@ func TestPacketEncodingIPv6(t *testing.T) {
Port: 1234, Port: 1234,
} }
var packetData [256]byte var packetData [256]byte
wrapped := AttachAddressToPacket(packetData[:], packetAddress) wrapped, err := AttachAddressToPacket(buf.FromBytes(packetData[:]), packetAddress)
assert.NoError(t, err)
packetPayload, decodedAddress := ExtractAddressFromPacket(wrapped) packetPayload, decodedAddress, err := ExtractAddressFromPacket(wrapped)
assert.NoError(t, err)
assert.Equal(t, packetPayload, packetData[:]) assert.Equal(t, packetPayload.Bytes(), packetData[:])
assert.Equal(t, packetAddress, decodedAddress) assert.Equal(t, packetAddress, decodedAddress)
} }