Debug direct connection

This commit is contained in:
9seconds
2018-07-07 20:45:10 +03:00
parent f82ff1f6fe
commit cd63483503
6 changed files with 32 additions and 17 deletions
+4 -6
View File
@@ -21,25 +21,23 @@ func NewProxyDirect(conf *config.Config) *Proxy {
return &Proxy{
conf: conf,
acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
connID string, wait *sync.WaitGroup, conf *config.Config) error {
connID string, wait *sync.WaitGroup, conf *config.Config) (io.Closer, io.Closer, error) {
client, opts, err := client.DirectInit(ctx, cancel, clientSocket, connID, conf)
if err != nil {
return errors.Annotate(err, "Cannot initialize client connection")
return nil, nil, errors.Annotate(err, "Cannot initialize client connection")
}
defer client.Close()
server, err := directTelegramStream(ctx, cancel, opts, connID, tg)
if err != nil {
return errors.Annotate(err, "Cannot initialize telegram connection")
return client, nil, errors.Annotate(err, "Cannot initialize telegram connection")
}
defer server.Close()
wait.Add(2)
go directPipe(client, server, wait)
go directPipe(server, client, wait)
return nil
return client, server, nil
},
}
}
+5 -6
View File
@@ -2,6 +2,7 @@ package proxy
import (
"context"
"io"
"net"
"sync"
@@ -20,25 +21,23 @@ func NewProxyMiddle(conf *config.Config) *Proxy {
return &Proxy{
conf: conf,
acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
connID string, wait *sync.WaitGroup, conf *config.Config) error {
connID string, wait *sync.WaitGroup, conf *config.Config) (io.Closer, io.Closer, error) {
client, opts, err := client.MiddleInit(ctx, cancel, clientSocket, connID, conf)
if err != nil {
return errors.Annotate(err, "Cannot initialize client connection")
return nil, nil, errors.Annotate(err, "Cannot initialize client connection")
}
defer client.Close()
server, err := middleTelegramStream(ctx, cancel, opts, connID, tg)
if err != nil {
return errors.Annotate(err, "Cannot initialize telegram connection")
return client, nil, errors.Annotate(err, "Cannot initialize telegram connection")
}
defer server.Close()
wait.Add(2)
go middlePipe(client, server, wait, &opts.ReadHacks)
go middlePipe(server, client, wait, &opts.WriteHacks)
return nil
return client, server, nil
},
}
}
+12 -2
View File
@@ -2,6 +2,7 @@ package proxy
import (
"context"
"io"
"net"
"sync"
@@ -12,7 +13,7 @@ import (
"github.com/9seconds/mtg/config"
)
type proxyAcceptCallback func(context.Context, context.CancelFunc, net.Conn, string, *sync.WaitGroup, *config.Config) error
type proxyAcceptCallback func(context.Context, context.CancelFunc, net.Conn, string, *sync.WaitGroup, *config.Config) (io.Closer, io.Closer, error)
type Proxy struct {
conf *config.Config
@@ -51,7 +52,16 @@ func (p *Proxy) accept(conn net.Conn) {
ctx, cancel := context.WithCancel(context.Background())
wait := &sync.WaitGroup{}
if err := p.acceptCallback(ctx, cancel, conn, connID, wait, p.conf); err != nil {
client, server, err := p.acceptCallback(ctx, cancel, conn, connID, wait, p.conf)
defer func() {
if client != nil {
client.Close()
}
if server != nil {
server.Close()
}
}()
if err != nil {
log.Errorw("Cannot initialize connection", "error", err)
cancel()
}