Fix WriteAddrPort usage

This commit is contained in:
世界 2023-09-07 09:06:46 +08:00
parent 0eec7bbe19
commit 03c21c0a12
No known key found for this signature in database
GPG key ID: CD109927C34A63C4
5 changed files with 39 additions and 14 deletions

View file

@ -30,7 +30,11 @@ func NewLazyConn(conn net.Conn, request Request) *Conn {
func (c *LazyClientConn) Write(p []byte) (n int, err error) {
if !c.requestWritten {
request := EncodeRequest(c.request)
var request *buf.Buffer
request, err = EncodeRequest(c.request)
if err != nil {
return
}
err = c.writer.WriteVectorised([]*buf.Buffer{request, buf.As(p)})
if err != nil {
return
@ -43,8 +47,12 @@ func (c *LazyClientConn) Write(p []byte) (n int, err error) {
func (c *LazyClientConn) WriteVectorised(buffers []*buf.Buffer) error {
if !c.requestWritten {
request := EncodeRequest(c.request)
err := c.writer.WriteVectorised(append([]*buf.Buffer{request}, buffers...))
request, err := EncodeRequest(c.request)
if err != nil {
return err
}
err = c.writer.WriteVectorised(append([]*buf.Buffer{request}, buffers...))
c.requestWritten = true
return err
}

View file

@ -51,20 +51,25 @@ func ReadRequest(reader io.Reader) (*Request, error) {
return &request, nil
}
func EncodeRequest(request Request) *buf.Buffer {
func EncodeRequest(request Request) (*buf.Buffer, error) {
var bufferLen int
bufferLen += 1 // isConnect
bufferLen += M.SocksaddrSerializer.AddrPortLen(request.Destination)
buffer := buf.NewSize(bufferLen)
common.Must(
binary.Write(buffer, binary.BigEndian, request.IsConnect),
M.SocksaddrSerializer.WriteAddrPort(buffer, request.Destination),
)
return buffer
common.Must(binary.Write(buffer, binary.BigEndian, request.IsConnect))
err := M.SocksaddrSerializer.WriteAddrPort(buffer, request.Destination)
if err != nil {
buffer.Release()
return nil, err
}
return buffer, nil
}
func WriteRequest(writer io.Writer, request Request) error {
buffer := EncodeRequest(request)
buffer, err := EncodeRequest(request)
if err != nil {
return err
}
defer buffer.Release()
return common.Error(writer.Write(buffer.Bytes()))
}

View file

@ -104,7 +104,10 @@ func (c *AssociatePacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Sock
func (c *AssociatePacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
header := buf.With(buffer.ExtendHeader(3 + M.SocksaddrSerializer.AddrPortLen(destination)))
common.Must(header.WriteZeroN(3))
common.Must(M.SocksaddrSerializer.WriteAddrPort(header, destination))
err := M.SocksaddrSerializer.WriteAddrPort(header, destination)
if err != nil {
return err
}
return common.Error(bufio.WritePacketBuffer(c.NetPacketConn, buffer, c.remoteAddr))
}

View file

@ -43,7 +43,10 @@ func (v *VectorisedAssociatePacketConn) WriteVectorisedPacket(buffers []*buf.Buf
header := buf.NewSize(3 + M.SocksaddrSerializer.AddrPortLen(destination))
defer header.Release()
common.Must(header.WriteZeroN(3))
common.Must(M.SocksaddrSerializer.WriteAddrPort(header, destination))
err := M.SocksaddrSerializer.WriteAddrPort(header, destination)
if err != nil {
return err
}
return v.VectorisedPacketWriter.WriteVectorisedPacket(append([]*buf.Buffer{header}, buffers...), destination)
}

View file

@ -193,8 +193,11 @@ func WriteRequest(writer io.Writer, request Request) error {
buffer.WriteByte(Version),
buffer.WriteByte(request.Command),
buffer.WriteZero(),
M.SocksaddrSerializer.WriteAddrPort(buffer, request.Destination),
)
err := M.SocksaddrSerializer.WriteAddrPort(buffer, request.Destination)
if err != nil {
return err
}
return rw.WriteBytes(writer, buffer.Bytes())
}
@ -244,8 +247,11 @@ func WriteResponse(writer io.Writer, response Response) error {
buffer.WriteByte(Version),
buffer.WriteByte(response.ReplyCode),
buffer.WriteZero(),
M.SocksaddrSerializer.WriteAddrPort(buffer, bind),
)
err := M.SocksaddrSerializer.WriteAddrPort(buffer, bind)
if err != nil {
return err
}
return rw.WriteBytes(writer, buffer.Bytes())
}