Просмотр исходного кода

Proxy is working in a simple mode now

tags/v2.0.0-rc1
9seconds 5 лет назад
Родитель
Сommit
a3362c7ea4

+ 2
- 2
mtglib/conns.go Просмотреть файл

26
 		})
26
 		})
27
 	}
27
 	}
28
 
28
 
29
-	return n, err
29
+	return n, err // nolint: wrapcheck
30
 }
30
 }
31
 
31
 
32
 func (c connTelegramTraffic) Write(b []byte) (int, error) {
32
 func (c connTelegramTraffic) Write(b []byte) (int, error) {
41
 		})
41
 		})
42
 	}
42
 	}
43
 
43
 
44
-	return n, err
44
+	return n, err // nolint: wrapcheck
45
 }
45
 }

+ 4
- 5
mtglib/internal/obfuscated2/client_handshake.go Просмотреть файл

23
 }
23
 }
24
 
24
 
25
 func (c *clientHandhakeFrame) encryptor(secret []byte) cipher.Stream {
25
 func (c *clientHandhakeFrame) encryptor(secret []byte) cipher.Stream {
26
-	arr := clientHandhakeFrame{}
27
-	invertByteSlices(arr.data[:], c.data[:])
26
+	invertedHandshake := c.invert()
28
 
27
 
29
 	hasher := acquireSha256Hasher()
28
 	hasher := acquireSha256Hasher()
30
 	defer releaseSha256Hasher(hasher)
29
 	defer releaseSha256Hasher(hasher)
31
 
30
 
32
-	hasher.Write(arr.key()) // nolint: errcheck
33
-	hasher.Write(secret)    // nolint: errcheck
31
+	hasher.Write(invertedHandshake.key()) // nolint: errcheck
32
+	hasher.Write(secret)                  // nolint: errcheck
34
 
33
 
35
-	return makeAesCtr(hasher.Sum(nil), arr.iv())
34
+	return makeAesCtr(hasher.Sum(nil), invertedHandshake.iv())
36
 }
35
 }
37
 
36
 
38
 func ClientHandshake(secret []byte, reader io.Reader) (int, cipher.Stream, cipher.Stream, error) {
37
 func ClientHandshake(secret []byte, reader io.Reader) (int, cipher.Stream, cipher.Stream, error) {

+ 10
- 0
mtglib/internal/obfuscated2/handshake_frame.go Просмотреть файл

60
 func (h *handshakeFrame) connectionType() []byte {
60
 func (h *handshakeFrame) connectionType() []byte {
61
 	return h.data[handshakeFrameOffsetConnectionType:handshakeFrameOffsetDC]
61
 	return h.data[handshakeFrameOffsetConnectionType:handshakeFrameOffsetDC]
62
 }
62
 }
63
+
64
+func (h *handshakeFrame) invert() handshakeFrame {
65
+	copyFrame := *h
66
+
67
+	for i := 0; i < handshakeFrameLenKey+handshakeFrameLenIV; i++ {
68
+		copyFrame.data[handshakeFrameOffsetKey+i] = h.data[handshakeFrameOffsetConnectionType-1-i]
69
+	}
70
+
71
+	return copyFrame
72
+}

+ 10
- 11
mtglib/internal/obfuscated2/server_handshake.go Просмотреть файл

5
 	"crypto/rand"
5
 	"crypto/rand"
6
 	"encoding/binary"
6
 	"encoding/binary"
7
 	"fmt"
7
 	"fmt"
8
-	"net"
8
+	"io"
9
 )
9
 )
10
 
10
 
11
 type serverHandshakeFrame struct {
11
 type serverHandshakeFrame struct {
13
 }
13
 }
14
 
14
 
15
 func (s *serverHandshakeFrame) decryptor() cipher.Stream {
15
 func (s *serverHandshakeFrame) decryptor() cipher.Stream {
16
-	return makeAesCtr(s.key(), s.iv())
16
+	invertedHandshake := s.invert()
17
+
18
+	return makeAesCtr(invertedHandshake.key(), invertedHandshake.iv())
17
 }
19
 }
18
 
20
 
19
 func (s *serverHandshakeFrame) encryptor() cipher.Stream {
21
 func (s *serverHandshakeFrame) encryptor() cipher.Stream {
20
-	arr := serverHandshakeFrame{}
21
-	invertByteSlices(arr.data[:], s.data[:])
22
-
23
-	return makeAesCtr(arr.key(), arr.iv())
22
+	return makeAesCtr(s.key(), s.iv())
24
 }
23
 }
25
 
24
 
26
-func ServerHandshake(conn net.Conn) (cipher.Stream, cipher.Stream, error) {
25
+func ServerHandshake(writer io.Writer) (cipher.Stream, cipher.Stream, error) {
27
 	handshake := generateServerHanshakeFrame()
26
 	handshake := generateServerHanshakeFrame()
28
 	copyHandshake := handshake
27
 	copyHandshake := handshake
29
 	encryptor := handshake.encryptor()
28
 	encryptor := handshake.encryptor()
33
 	copy(handshake.key(), copyHandshake.key())
32
 	copy(handshake.key(), copyHandshake.key())
34
 	copy(handshake.iv(), copyHandshake.iv())
33
 	copy(handshake.iv(), copyHandshake.iv())
35
 
34
 
36
-	if _, err := conn.Write(handshake.data[:]); err != nil {
35
+	if _, err := writer.Write(handshake.data[:]); err != nil {
37
 		return nil, nil, fmt.Errorf("cannot send a handshake frame to telegram: %w", err)
36
 		return nil, nil, fmt.Errorf("cannot send a handshake frame to telegram: %w", err)
38
 	}
37
 	}
39
 
38
 
48
 			panic(err)
47
 			panic(err)
49
 		}
48
 		}
50
 
49
 
51
-		if frame.data[0] == 0xef {
50
+		if frame.data[0] == 0xef { // nolint: gomnd // taken from tg sources
52
 			continue
51
 			continue
53
 		}
52
 		}
54
 
53
 
55
 		switch binary.LittleEndian.Uint32(frame.data[:4]) {
54
 		switch binary.LittleEndian.Uint32(frame.data[:4]) {
56
-		case 0x44414548, 0x54534f50, 0x20544547, 0x4954504f, 0xeeeeeeee:
55
+		case 0x44414548, 0x54534f50, 0x20544547, 0x4954504f, 0xeeeeeeee: // nolint: gomnd // taken from tg sources
57
 			continue
56
 			continue
58
 		}
57
 		}
59
 
58
 
60
-		if (frame.data[4] | frame.data[5] | frame.data[6] | frame.data[7]) == 0 {
59
+		if frame.data[4]|frame.data[5]|frame.data[6]|frame.data[7] == 0 {
61
 			continue
60
 			continue
62
 		}
61
 		}
63
 
62
 

+ 0
- 8
mtglib/internal/obfuscated2/utils.go Просмотреть файл

13
 
13
 
14
 	return cipher.NewCTR(block, iv)
14
 	return cipher.NewCTR(block, iv)
15
 }
15
 }
16
-
17
-func invertByteSlices(dst, src []byte) {
18
-	lenDst := len(dst) - 1
19
-
20
-	for i, v := range src {
21
-		dst[lenDst-i] = v
22
-	}
23
-}

+ 33
- 0
mtglib/internal/relay/conn.go Просмотреть файл

1
+package relay
2
+
3
+import "io"
4
+
5
+type conn struct {
6
+	io.ReadWriteCloser
7
+
8
+	relay *Relay
9
+}
10
+
11
+func (c conn) Read(p []byte) (int, error) {
12
+	ctx := c.relay.ctx
13
+	n, err := c.ReadWriteCloser.Read(p)
14
+
15
+	select {
16
+	case <-ctx.Done():
17
+	case c.relay.tickChannel <- struct{}{}:
18
+	}
19
+
20
+	return n, err // nolint: wrapcheck
21
+}
22
+
23
+func (c conn) Write(p []byte) (int, error) {
24
+	ctx := c.relay.ctx
25
+	n, err := c.ReadWriteCloser.Write(p)
26
+
27
+	select {
28
+	case <-ctx.Done():
29
+	case c.relay.tickChannel <- struct{}{}:
30
+	}
31
+
32
+	return n, err // nolint: wrapcheck
33
+}

+ 46
- 0
mtglib/internal/relay/pools.go Просмотреть файл

1
+package relay
2
+
3
+import (
4
+	"context"
5
+	"sync"
6
+	"time"
7
+)
8
+
9
+var relayPool = sync.Pool{
10
+	New: func() interface{} {
11
+		return &Relay{
12
+			tickChannel:  make(chan struct{}),
13
+			errorChannel: make(chan error, 1),
14
+		}
15
+	},
16
+}
17
+
18
+func AcquireRelay(ctx context.Context, logger Logger, bufferSize int, idleTimeout time.Duration) *Relay {
19
+	ctx, cancel := context.WithCancel(ctx)
20
+
21
+	r := relayPool.Get().(*Relay)
22
+	r.ctx = ctx
23
+	r.ctxCancel = cancel
24
+	r.logger = logger
25
+	r.tickTimeout = idleTimeout
26
+
27
+	if len(r.eastBuffer) != bufferSize {
28
+		r.eastBuffer = make([]byte, bufferSize)
29
+	}
30
+
31
+	if len(r.westBuffer) != bufferSize {
32
+		r.westBuffer = make([]byte, bufferSize)
33
+	}
34
+
35
+	return r
36
+}
37
+
38
+func ReleaseRelay(r *Relay) {
39
+	r.ctxCancel()
40
+
41
+	r.ctx = nil
42
+	r.ctxCancel = nil
43
+	r.logger = nil
44
+
45
+	relayPool.Put(r)
46
+}

+ 100
- 0
mtglib/internal/relay/relay.go Просмотреть файл

1
+package relay
2
+
3
+import (
4
+	"context"
5
+	"io"
6
+	"sync"
7
+	"time"
8
+)
9
+
10
+type Relay struct {
11
+	ctx          context.Context
12
+	ctxCancel    context.CancelFunc
13
+	logger       Logger
14
+	eastBuffer   []byte
15
+	westBuffer   []byte
16
+	tickChannel  chan struct{}
17
+	errorChannel chan error
18
+	tickTimeout  time.Duration
19
+}
20
+
21
+func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
22
+	eastConn = conn{
23
+		ReadWriteCloser: eastConn,
24
+		relay:           r,
25
+	}
26
+	westConn = conn{
27
+		ReadWriteCloser: westConn,
28
+		relay:           r,
29
+	}
30
+
31
+	defer func() {
32
+		r.ctxCancel()
33
+		eastConn.Close()
34
+		westConn.Close()
35
+	}()
36
+
37
+	go r.runObserver()
38
+
39
+	wg := &sync.WaitGroup{}
40
+	wg.Add(2) // nolint: gomnd
41
+
42
+	go r.transmit(eastConn, westConn, r.westBuffer, "west", wg)
43
+
44
+	r.transmit(westConn, eastConn, r.eastBuffer, "east", wg)
45
+
46
+	wg.Wait()
47
+
48
+	return <-r.errorChannel
49
+}
50
+
51
+func (r *Relay) transmit(src io.ReadCloser, dst io.WriteCloser,
52
+	buffer []byte, direction string, wg *sync.WaitGroup) {
53
+	defer func() {
54
+		wg.Done()
55
+		src.Close()
56
+		dst.Close()
57
+	}()
58
+
59
+	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {
60
+		r.logger.Printf("error '%v' happened on direction %s", err, direction)
61
+
62
+		select {
63
+		case <-r.ctx.Done():
64
+		case r.errorChannel <- err:
65
+		default:
66
+		}
67
+	}
68
+}
69
+
70
+func (r *Relay) runObserver() {
71
+	ticker := time.NewTicker(time.Second)
72
+
73
+	defer func() {
74
+		ticker.Stop()
75
+
76
+		select {
77
+		case <-ticker.C:
78
+		default:
79
+		}
80
+	}()
81
+
82
+	lastTickAt := time.Now()
83
+	ctx := r.ctx
84
+
85
+	for {
86
+		select {
87
+		case <-ctx.Done():
88
+			return
89
+		case <-r.tickChannel:
90
+			lastTickAt = time.Now()
91
+		case <-ticker.C:
92
+			if time.Since(lastTickAt) > r.tickTimeout {
93
+				r.logger.Printf("exit due to a timeout")
94
+				r.ctxCancel()
95
+
96
+				return
97
+			}
98
+		}
99
+	}
100
+}

+ 15
- 0
mtglib/proxy.go Просмотреть файл

9
 	"time"
9
 	"time"
10
 
10
 
11
 	"github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2"
11
 	"github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2"
12
+	"github.com/9seconds/mtg/v2/mtglib/internal/relay"
12
 	"github.com/9seconds/mtg/v2/mtglib/internal/telegram"
13
 	"github.com/9seconds/mtg/v2/mtglib/internal/telegram"
13
 	"github.com/panjf2000/ants/v2"
14
 	"github.com/panjf2000/ants/v2"
14
 )
15
 )
19
 	streamWaitGroup sync.WaitGroup
20
 	streamWaitGroup sync.WaitGroup
20
 
21
 
21
 	idleTimeout time.Duration
22
 	idleTimeout time.Duration
23
+	bufferSize  int
22
 	workerPool  *ants.PoolWithFunc
24
 	workerPool  *ants.PoolWithFunc
23
 	telegram    *telegram.Telegram
25
 	telegram    *telegram.Telegram
24
 
26
 
64
 
66
 
65
 		return
67
 		return
66
 	}
68
 	}
69
+
70
+	rel := relay.AcquireRelay(ctx, p.logger.Named("relay"), p.bufferSize, p.idleTimeout)
71
+	defer relay.ReleaseRelay(rel)
72
+
73
+	if err := rel.Process(ctx.clientConn, ctx.telegramConn); err != nil {
74
+		p.logger.DebugError("relay has been finished", err)
75
+	}
67
 }
76
 }
68
 
77
 
69
 func (p *Proxy) Serve(listener net.Listener) error {
78
 func (p *Proxy) Serve(listener net.Listener) error {
185
 		idleTimeout = DefaultIdleTimeout
194
 		idleTimeout = DefaultIdleTimeout
186
 	}
195
 	}
187
 
196
 
197
+	bufferSize := opts.BufferSize
198
+	if bufferSize < 1 {
199
+		bufferSize = DefaultBufferSize
200
+	}
201
+
188
 	ctx, cancel := context.WithCancel(context.Background())
202
 	ctx, cancel := context.WithCancel(context.Background())
189
 	proxy := &Proxy{
203
 	proxy := &Proxy{
190
 		ctx:             ctx,
204
 		ctx:             ctx,
195
 		eventStream:     opts.EventStream,
209
 		eventStream:     opts.EventStream,
196
 		logger:          opts.Logger.Named("proxy"),
210
 		logger:          opts.Logger.Named("proxy"),
197
 		idleTimeout:     idleTimeout,
211
 		idleTimeout:     idleTimeout,
212
+		bufferSize:      int(bufferSize),
198
 		telegram:        tg,
213
 		telegram:        tg,
199
 	}
214
 	}
200
 
215
 

Загрузка…
Отмена
Сохранить