1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-21 09:36:34 -05:00

Apply buffered reader and writer to socks

This commit is contained in:
v2ray 2016-02-27 11:02:42 +01:00
parent cba26dabe8
commit ef51c600fb
4 changed files with 39 additions and 40 deletions

View File

@ -26,7 +26,7 @@ func (this *BufferedWriter) Write(b []byte) (int, error) {
}
nBytes, _ := this.buffer.Write(b)
if this.buffer.IsFull() {
err := this.flush()
err := this.Flush()
if err != nil {
return nBytes, err
}
@ -34,7 +34,7 @@ func (this *BufferedWriter) Write(b []byte) (int, error) {
return nBytes, nil
}
func (this *BufferedWriter) flush() error {
func (this *BufferedWriter) Flush() error {
nBytes, err := this.writer.Write(this.buffer.Value)
this.buffer.SliceFrom(nBytes)
if !this.buffer.IsEmpty() {
@ -54,7 +54,7 @@ func (this *BufferedWriter) Cached() bool {
func (this *BufferedWriter) SetCached(cached bool) {
this.cached = cached
if !cached && !this.buffer.IsEmpty() {
this.flush()
this.Flush()
}
}

View File

@ -295,16 +295,16 @@ func (r *Socks5Response) SetDomain(domain string) {
r.Domain = domain
}
func (r *Socks5Response) Write(buffer *alloc.Buffer) {
buffer.AppendBytes(r.Version, r.Error, 0x00 /* reserved */, r.AddrType)
func (r *Socks5Response) Write(writer io.Writer) {
writer.Write([]byte{r.Version, r.Error, 0x00 /* reserved */, r.AddrType})
switch r.AddrType {
case 0x01:
buffer.Append(r.IPv4[:])
writer.Write(r.IPv4[:])
case 0x03:
buffer.AppendBytes(byte(len(r.Domain)))
buffer.Append([]byte(r.Domain))
writer.Write([]byte{byte(len(r.Domain))})
writer.Write([]byte(r.Domain))
case 0x04:
buffer.Append(r.IPv6[:])
writer.Write(r.IPv6[:])
}
buffer.Append(r.Port.Bytes())
writer.Write(r.Port.Bytes())
}

View File

@ -2,8 +2,8 @@ package protocol
import (
"errors"
"io"
"github.com/v2ray/v2ray-core/common/alloc"
v2net "github.com/v2ray/v2ray-core/common/net"
)
@ -32,8 +32,8 @@ func NewSocks4AuthenticationResponse(result byte, port v2net.Port, ip []byte) *S
}
}
func (r *Socks4AuthenticationResponse) Write(buffer *alloc.Buffer) {
buffer.AppendBytes(
byte(0x00), r.result, byte(r.port>>8), byte(r.port),
r.ip[0], r.ip[1], r.ip[2], r.ip[3])
func (r *Socks4AuthenticationResponse) Write(writer io.Writer) {
writer.Write([]byte{
byte(0x00), r.result, byte(r.port >> 8), byte(r.port),
r.ip[0], r.ip[1], r.ip[2], r.ip[3]})
}

View File

@ -7,7 +7,6 @@ import (
"time"
"github.com/v2ray/v2ray-core/app/dispatcher"
"github.com/v2ray/v2ray-core/common/alloc"
v2io "github.com/v2ray/v2ray-core/common/io"
"github.com/v2ray/v2ray-core/common/log"
v2net "github.com/v2ray/v2ray-core/common/net"
@ -94,7 +93,10 @@ func (this *SocksServer) Listen(port v2net.Port) error {
func (this *SocksServer) handleConnection(connection *hub.TCPConn) {
defer connection.Close()
reader := v2net.NewTimeOutReader(120, connection)
timedReader := v2net.NewTimeOutReader(120, connection)
reader := v2io.NewBufferedReader(timedReader)
writer := v2io.NewBufferedWriter(connection)
auth, auth4, err := protocol.ReadAuthentication(reader)
if err != nil && err != protocol.Socks4Downgrade {
@ -103,13 +105,13 @@ func (this *SocksServer) handleConnection(connection *hub.TCPConn) {
}
if err != nil && err == protocol.Socks4Downgrade {
this.handleSocks4(reader, connection, auth4)
this.handleSocks4(reader, writer, auth4)
} else {
this.handleSocks5(reader, connection, auth)
this.handleSocks5(reader, writer, auth)
}
}
func (this *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Writer, auth protocol.Socks5AuthenticationRequest) error {
func (this *SocksServer) handleSocks5(reader *v2io.BufferedReader, writer *v2io.BufferedWriter, auth protocol.Socks5AuthenticationRequest) error {
expectedAuthMethod := protocol.AuthNotRequired
if this.config.AuthType == AuthTypePassword {
expectedAuthMethod = protocol.AuthUserPass
@ -118,6 +120,7 @@ func (this *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Wri
if !auth.HasAuthMethod(expectedAuthMethod) {
authResponse := protocol.NewAuthenticationResponse(protocol.AuthNoMatchingMethod)
err := protocol.WriteAuthentication(writer, authResponse)
writer.Flush()
if err != nil {
log.Error("Socks: failed to write authentication: ", err)
return err
@ -128,6 +131,7 @@ func (this *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Wri
authResponse := protocol.NewAuthenticationResponse(expectedAuthMethod)
err := protocol.WriteAuthentication(writer, authResponse)
writer.Flush()
if err != nil {
log.Error("Socks: failed to write authentication: ", err)
return err
@ -144,6 +148,7 @@ func (this *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Wri
}
upResponse := protocol.NewSocks5UserPassResponse(status)
err = protocol.WriteUserPassResponse(writer, upResponse)
writer.Flush()
if err != nil {
log.Error("Socks: failed to write user pass response: ", err)
return err
@ -170,10 +175,8 @@ func (this *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Wri
response.Port = v2net.Port(0)
response.SetIPv4([]byte{0, 0, 0, 0})
responseBuffer := alloc.NewSmallBuffer().Clear()
response.Write(responseBuffer)
_, err = writer.Write(responseBuffer.Value)
responseBuffer.Release()
response.Write(writer)
writer.Flush()
if err != nil {
log.Error("Socks: failed to write response: ", err)
return err
@ -189,15 +192,15 @@ func (this *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Wri
response.Port = v2net.Port(1717)
response.SetIPv4([]byte{0, 0, 0, 0})
responseBuffer := alloc.NewSmallBuffer().Clear()
response.Write(responseBuffer)
_, err = writer.Write(responseBuffer.Value)
responseBuffer.Release()
response.Write(writer)
if err != nil {
log.Error("Socks: failed to write response: ", err)
return err
}
reader.SetCached(false)
writer.SetCached(false)
dest := request.Destination()
log.Info("Socks: TCP Connect request to ", dest)
@ -206,7 +209,7 @@ func (this *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.Wri
return nil
}
func (this *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer) error {
func (this *SocksServer) handleUDP(reader io.Reader, writer *v2io.BufferedWriter) error {
response := protocol.NewSocks5Response()
response.Error = protocol.ErrorSuccess
@ -222,18 +225,14 @@ func (this *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer
response.SetDomain(udpAddr.Address().Domain())
}
responseBuffer := alloc.NewSmallBuffer().Clear()
response.Write(responseBuffer)
_, err := writer.Write(responseBuffer.Value)
responseBuffer.Release()
response.Write(writer)
err := writer.Flush()
if err != nil {
log.Error("Socks: failed to write response: ", err)
return err
}
reader.SetTimeOut(300) /* 5 minutes */
v2io.ReadFrom(reader, nil) // Just in case of anything left in the socket
// The TCP connection closes after this method returns. We need to wait until
// the client closes it.
// TODO: get notified from UDP part
@ -242,23 +241,23 @@ func (this *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writer
return nil
}
func (this *SocksServer) handleSocks4(reader io.Reader, writer io.Writer, auth protocol.Socks4AuthenticationRequest) error {
func (this *SocksServer) handleSocks4(reader *v2io.BufferedReader, writer *v2io.BufferedWriter, auth protocol.Socks4AuthenticationRequest) error {
result := protocol.Socks4RequestGranted
if auth.Command == protocol.CmdBind {
result = protocol.Socks4RequestRejected
}
socks4Response := protocol.NewSocks4AuthenticationResponse(result, auth.Port, auth.IP[:])
responseBuffer := alloc.NewSmallBuffer().Clear()
socks4Response.Write(responseBuffer)
writer.Write(responseBuffer.Value)
responseBuffer.Release()
socks4Response.Write(writer)
if result == protocol.Socks4RequestRejected {
log.Warning("Socks: Unsupported socks 4 command ", auth.Command)
return ErrorUnsupportedSocksCommand
}
reader.SetCached(false)
writer.SetCached(false)
dest := v2net.TCPDestination(v2net.IPAddress(auth.IP[:]), auth.Port)
packet := v2net.NewPacket(dest, nil, true)
this.transport(reader, writer, packet)