Proxy is working but unstable

This commit is contained in:
9seconds
2018-07-06 10:21:04 +03:00
parent 5413a5b090
commit 88faeb7195
10 changed files with 83 additions and 65 deletions
+1 -1
View File
@@ -9,4 +9,4 @@ import (
)
// Init has to initialize client connection based on given config.
type Init func(net.Conn, *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error)
type Init func(net.Conn, *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error)
+2 -2
View File
@@ -15,7 +15,7 @@ import (
const handshakeTimeout = 10 * time.Second
// DirectInit initializes client to access Telegram bypassing middleproxies.
func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error) {
func DirectInit(conn net.Conn, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
if err := config.SetSocketOptions(conn); err != nil {
return nil, nil, errors.Annotate(err, "Cannot set socket options")
}
@@ -37,5 +37,5 @@ func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wr
socket := wrappers.NewTimeoutRWC(conn, conf.PublicIPv4, conf.PublicIPv6)
socket = wrappers.NewStreamCipherRWC(socket, obfs2.Encryptor, obfs2.Decryptor)
return connOpts, socket, nil
return socket, connOpts, nil
}
+3 -3
View File
@@ -9,8 +9,8 @@ import (
"github.com/9seconds/mtg/wrappers"
)
func MiddleInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error) {
opts, newConn, err := DirectInit(conn, conf)
func MiddleInit(conn net.Conn, conf *config.Config) (wrappers.ReadWriteCloserWithAddr, *mtproto.ConnectionOpts, error) {
newConn, opts, err := DirectInit(conn, conf)
if err != nil {
return nil, nil, err
}
@@ -26,5 +26,5 @@ func MiddleInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wr
opts.ConnectionProto = mtproto.ConnectionProtocolIPv6
}
return opts, newConn, nil
return newConn, opts, nil
}
+7 -2
View File
@@ -13,14 +13,19 @@ type ConnectionType uint8
type ConnectionProtocol uint8
type Hacks struct {
SimpleAck bool
QuickAck bool
}
// ConnectionOpts presents an options, metadata on connection requested
// by the user on handshake.
type ConnectionOpts struct {
DC int16
ConnectionType ConnectionType
ConnectionProto ConnectionProtocol
QuickAck bool
SimpleAck bool
ReadHacks Hacks
WriteHacks Hacks
ClientAddr *net.TCPAddr
}
+1 -1
View File
@@ -24,7 +24,7 @@ func (r *ProxyRequest) Bytes(message []byte) []byte {
buf := &bytes.Buffer{}
flags := r.Flags
if r.Options.QuickAck {
if r.Options.ReadHacks.QuickAck {
flags |= proxyRequestFlagsQuickAck
}
+7 -6
View File
@@ -28,9 +28,6 @@ type AbridgedReadWriteCloserWithAddr struct {
func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
return a.BufferedRead(p, func() error {
a.opts.QuickAck = false
a.opts.SimpleAck = false
buf := &bytes.Buffer{}
buf.Grow(3)
@@ -41,7 +38,7 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
buf.Reset()
if msgLength >= abridgedQuickAckLength {
a.opts.QuickAck = true
a.opts.ReadHacks.QuickAck = true
msgLength -= 0x80
}
@@ -56,9 +53,13 @@ func (a *AbridgedReadWriteCloserWithAddr) Read(p []byte) (int, error) {
}
msgLength32 *= 4
if _, err := io.CopyN(a.Buffer, a.conn, int64(msgLength32)); err != nil {
buf.Reset()
buf.Grow(int(msgLength32))
if _, err := io.CopyN(buf, a.conn, int64(msgLength32)); err != nil {
return errors.Annotate(err, "Cannot read message")
}
a.Buffer.Write(buf.Bytes())
return nil
})
@@ -68,7 +69,7 @@ func (a *AbridgedReadWriteCloserWithAddr) Write(p []byte) (int, error) {
if len(p)%4 != 0 {
return 0, errors.Errorf("Incorrect packet length %d", len(p))
}
if a.opts.SimpleAck {
if a.opts.WriteHacks.SimpleAck {
return a.conn.Write(reverseBytes(p))
}
+2 -5
View File
@@ -23,9 +23,6 @@ type IntermediateReadWriteCloserWithAddr struct {
func (i *IntermediateReadWriteCloserWithAddr) Read(p []byte) (int, error) {
return i.BufferedRead(p, func() error {
i.opts.QuickAck = false
i.opts.SimpleAck = false
buf := &bytes.Buffer{}
buf.Grow(4)
@@ -37,7 +34,7 @@ func (i *IntermediateReadWriteCloserWithAddr) Read(p []byte) (int, error) {
buf.Grow(int(length))
if length > intermediateQuickAckLength {
i.opts.QuickAck = true
i.opts.ReadHacks.QuickAck = true
length -= intermediateQuickAckLength
}
@@ -58,7 +55,7 @@ func (i *IntermediateReadWriteCloserWithAddr) Read(p []byte) (int, error) {
}
func (i *IntermediateReadWriteCloserWithAddr) Write(p []byte) (int, error) {
if i.opts.SimpleAck {
if i.opts.WriteHacks.SimpleAck {
return i.conn.Write(p)
}
+19 -19
View File
@@ -22,23 +22,21 @@ type ProxyRequestReadWriteCloserWithAddr struct {
func (p *ProxyRequestReadWriteCloserWithAddr) Read(buf []byte) (int, error) {
return p.BufferedRead(buf, func() error {
ansBuf := &bytes.Buffer{}
ansBuf.Grow(4)
if _, err := io.CopyN(ansBuf, p.conn, 4); err != nil {
ans := make([]byte, 4)
if _, err := io.ReadFull(p.conn, ans); err != nil {
return errors.Annotate(err, "Cannot read RPC tag")
}
switch {
case bytes.Equal(ansBuf.Bytes(), rpc.TagCloseExt):
return p.readCloseExt()
case bytes.Equal(ansBuf.Bytes(), rpc.TagProxyAns):
case bytes.Equal(ans, rpc.TagProxyAns):
return p.readProxyAns(buf)
case bytes.Equal(ansBuf.Bytes(), rpc.TagSimpleAck):
case bytes.Equal(ans, rpc.TagSimpleAck):
return p.readSimpleAck()
case bytes.Equal(ans, rpc.TagCloseExt):
return p.readCloseExt()
}
return errors.Errorf("Unknown RPC answer %s", ansBuf.Bytes())
return errors.Errorf("Unknown RPC answer %v", ans)
})
}
@@ -46,21 +44,21 @@ func (p *ProxyRequestReadWriteCloserWithAddr) readCloseExt() error {
return errors.New("Connection has been closed remotely")
}
func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns(buf []byte) error {
if _, err := io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns(buf []byte) (err error) {
if _, err = io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
return errors.Annotate(err, "Cannot skip flags and connid")
}
for {
n, err := p.conn.Read(buf)
n := len(buf)
preBuffer := &bytes.Buffer{}
for n == len(buf) {
n, err = p.conn.Read(buf)
if err != nil {
return errors.Annotate(err, "Cannot read proxy answer")
}
if n == 0 {
break
}
p.Buffer.Write(buf[:n])
preBuffer.Write(buf[:n])
}
p.Buffer.Write(preBuffer.Bytes())
return nil
}
@@ -69,10 +67,12 @@ func (p *ProxyRequestReadWriteCloserWithAddr) readSimpleAck() error {
if _, err := io.CopyN(ioutil.Discard, p.conn, 8); err != nil {
return errors.Annotate(err, "Cannot skip connid")
}
if _, err := io.CopyN(p.Buffer, p.conn, 4); err != nil {
ackData := make([]byte, 4)
if _, err := io.ReadFull(p.conn, ackData); err != nil {
return errors.Annotate(err, "Cannot read simple ack")
}
p.req.Options.SimpleAck = true
p.Buffer.Write(ackData)
return nil
}
-18
View File
@@ -1,18 +0,0 @@
package proxy
import (
"sync"
"github.com/9seconds/mtg/config"
)
var copyPool sync.Pool
func init() {
copyPool = sync.Pool{
New: func() interface{} {
data := make([]byte, config.BufferSizeCopy)
return &data
},
}
}
+41 -8
View File
@@ -85,8 +85,32 @@ func (s *Server) accept(conn net.Conn) {
wait := &sync.WaitGroup{}
wait.Add(2)
go s.pipe(clientConn, tgConn, wait)
go s.pipe(tgConn, clientConn, wait)
go func() {
defer wait.Done()
for {
connOpts.ReadHacks.QuickAck = false
connOpts.ReadHacks.SimpleAck = false
if err := s.pump(clientConn, tgConn, socketID, "client"); err != nil {
s.logger.Infow("Client stream is aborted",
"socketid", socketID, "error", err)
return
}
}
}()
go func() {
defer wait.Done()
for {
connOpts.WriteHacks.QuickAck = false
connOpts.WriteHacks.SimpleAck = false
if err := s.pump(tgConn, clientConn, socketID, "telegram"); err != nil {
s.logger.Infow("Telegram stream is aborted",
"socketid", socketID, "error", err)
return
}
}
}()
<-ctx.Done()
wait.Wait()
@@ -98,7 +122,7 @@ func (s *Server) accept(conn net.Conn) {
}
func (s *Server) getClientStream(ctx context.Context, cancel context.CancelFunc, conn net.Conn, socketID string) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
connOpts, socket, err := s.clientInit(conn, s.conf)
socket, connOpts, err := s.clientInit(conn, s.conf)
if err != nil {
return nil, nil, errors.Annotate(err, "Cannot init client connection")
}
@@ -128,13 +152,22 @@ func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFun
return conn, nil
}
func (s *Server) pipe(dst io.Writer, src io.Reader, wait *sync.WaitGroup) {
defer wait.Done()
func (s *Server) pump(src io.Reader, dst io.Writer, socketID, name string) (err error) {
copyBuf := make([]byte, 1024*1024*2)
buf := copyPool.Get().(*[]byte)
defer copyPool.Put(buf)
n := config.BufferSizeCopy
for n == config.BufferSizeCopy {
n, err = src.Read(copyBuf)
if err != nil {
break
}
_, err = dst.Write(copyBuf[:n])
if err != nil {
break
}
}
io.CopyBuffer(dst, src, *buf) // nolint: errcheck
return
}
// NewServer creates new instance of MTPROTO proxy.