Simplify relay

This commit is contained in:
9seconds
2021-03-29 10:37:09 +03:00
parent 3992054560
commit 4c3f42e264
3 changed files with 23 additions and 12 deletions
+1 -1
View File
@@ -29,11 +29,11 @@ func (c *Conn) Read(p []byte) (int, error) {
}
switch rec.Type { // nolint: exhaustive
case record.TypeChangeCipherSpec:
case record.TypeApplicationData:
rec.Payload.WriteTo(&c.readBuffer) // nolint: errcheck
return c.readBuffer.Read(p)
case record.TypeChangeCipherSpec:
default:
return 0, fmt.Errorf("unsupported record type %v", rec.Type)
}
+10 -6
View File
@@ -1,19 +1,23 @@
package relay
import "io"
import (
"context"
"io"
)
type conn struct {
io.ReadWriteCloser
relay *Relay
ctx context.Context
tickChannel chan struct{}
}
func (c conn) Read(p []byte) (int, error) {
n, err := c.ReadWriteCloser.Read(p)
select {
case <-c.relay.ctx.Done():
case c.relay.tickChannel <- struct{}{}:
case <-c.ctx.Done():
case c.tickChannel <- struct{}{}:
}
return n, err // nolint: wrapcheck
@@ -23,8 +27,8 @@ func (c conn) Write(p []byte) (int, error) {
n, err := c.ReadWriteCloser.Write(p)
select {
case <-c.relay.ctx.Done():
case c.relay.tickChannel <- struct{}{}:
case <-c.ctx.Done():
case c.tickChannel <- struct{}{}:
}
return n, err // nolint: wrapcheck
+12 -5
View File
@@ -21,11 +21,13 @@ type Relay struct {
func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
eastConn = conn{
ReadWriteCloser: eastConn,
relay: r,
ctx: r.ctx,
tickChannel: r.tickChannel,
}
westConn = conn{
ReadWriteCloser: westConn,
relay: r,
ctx: r.ctx,
tickChannel: r.tickChannel,
}
defer func() {
@@ -37,7 +39,7 @@ func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
wg := &sync.WaitGroup{}
wg.Add(3) // nolint: gomnd
go r.runObserver(r.ctx, wg)
go r.runObserver(wg)
go r.transmit(eastConn, westConn, r.westBuffer, "west", wg)
@@ -66,13 +68,18 @@ func (r *Relay) transmit(src io.ReadCloser, dst io.WriteCloser,
select {
case <-r.ctx.Done():
err = r.ctx.Err()
default:
}
select {
case r.errorChannel <- err:
default:
}
}
}
func (r *Relay) runObserver(ctx context.Context, wg *sync.WaitGroup) {
func (r *Relay) runObserver(wg *sync.WaitGroup) {
ticker := time.NewTicker(time.Second)
defer func() {
@@ -90,7 +97,7 @@ func (r *Relay) runObserver(ctx context.Context, wg *sync.WaitGroup) {
for {
select {
case <-ctx.Done():
case <-r.ctx.Done():
return
case <-r.tickChannel:
lastTickAt = time.Now()