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

Simplify relay

tags/v2.0.0-rc1
9seconds 5 лет назад
Родитель
Сommit
4c3f42e264
3 измененных файлов: 23 добавлений и 12 удалений
  1. 1
    1
      mtglib/internal/faketls/conn.go
  2. 10
    6
      mtglib/internal/relay/conn.go
  3. 12
    5
      mtglib/internal/relay/relay.go

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

29
 		}
29
 		}
30
 
30
 
31
 		switch rec.Type { // nolint: exhaustive
31
 		switch rec.Type { // nolint: exhaustive
32
-		case record.TypeChangeCipherSpec:
33
 		case record.TypeApplicationData:
32
 		case record.TypeApplicationData:
34
 			rec.Payload.WriteTo(&c.readBuffer) // nolint: errcheck
33
 			rec.Payload.WriteTo(&c.readBuffer) // nolint: errcheck
35
 
34
 
36
 			return c.readBuffer.Read(p)
35
 			return c.readBuffer.Read(p)
36
+		case record.TypeChangeCipherSpec:
37
 		default:
37
 		default:
38
 			return 0, fmt.Errorf("unsupported record type %v", rec.Type)
38
 			return 0, fmt.Errorf("unsupported record type %v", rec.Type)
39
 		}
39
 		}

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

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

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

21
 func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
21
 func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
22
 	eastConn = conn{
22
 	eastConn = conn{
23
 		ReadWriteCloser: eastConn,
23
 		ReadWriteCloser: eastConn,
24
-		relay:           r,
24
+		ctx:             r.ctx,
25
+		tickChannel:     r.tickChannel,
25
 	}
26
 	}
26
 	westConn = conn{
27
 	westConn = conn{
27
 		ReadWriteCloser: westConn,
28
 		ReadWriteCloser: westConn,
28
-		relay:           r,
29
+		ctx:             r.ctx,
30
+		tickChannel:     r.tickChannel,
29
 	}
31
 	}
30
 
32
 
31
 	defer func() {
33
 	defer func() {
37
 	wg := &sync.WaitGroup{}
39
 	wg := &sync.WaitGroup{}
38
 	wg.Add(3) // nolint: gomnd
40
 	wg.Add(3) // nolint: gomnd
39
 
41
 
40
-	go r.runObserver(r.ctx, wg)
42
+	go r.runObserver(wg)
41
 
43
 
42
 	go r.transmit(eastConn, westConn, r.westBuffer, "west", wg)
44
 	go r.transmit(eastConn, westConn, r.westBuffer, "west", wg)
43
 
45
 
66
 
68
 
67
 		select {
69
 		select {
68
 		case <-r.ctx.Done():
70
 		case <-r.ctx.Done():
71
+			err = r.ctx.Err()
72
+		default:
73
+		}
74
+
75
+		select {
69
 		case r.errorChannel <- err:
76
 		case r.errorChannel <- err:
70
 		default:
77
 		default:
71
 		}
78
 		}
72
 	}
79
 	}
73
 }
80
 }
74
 
81
 
75
-func (r *Relay) runObserver(ctx context.Context, wg *sync.WaitGroup) {
82
+func (r *Relay) runObserver(wg *sync.WaitGroup) {
76
 	ticker := time.NewTicker(time.Second)
83
 	ticker := time.NewTicker(time.Second)
77
 
84
 
78
 	defer func() {
85
 	defer func() {
90
 
97
 
91
 	for {
98
 	for {
92
 		select {
99
 		select {
93
-		case <-ctx.Done():
100
+		case <-r.ctx.Done():
94
 			return
101
 			return
95
 		case <-r.tickChannel:
102
 		case <-r.tickChannel:
96
 			lastTickAt = time.Now()
103
 			lastTickAt = time.Now()

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