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 switch rec.Type { // nolint: exhaustive
case record.TypeChangeCipherSpec:
case record.TypeApplicationData: case record.TypeApplicationData:
rec.Payload.WriteTo(&c.readBuffer) // nolint: errcheck rec.Payload.WriteTo(&c.readBuffer) // nolint: errcheck
return c.readBuffer.Read(p) return c.readBuffer.Read(p)
case record.TypeChangeCipherSpec:
default: default:
return 0, fmt.Errorf("unsupported record type %v", rec.Type) return 0, fmt.Errorf("unsupported record type %v", rec.Type)
} }
+10 -6
View File
@@ -1,19 +1,23 @@
package relay package relay
import "io" import (
"context"
"io"
)
type conn struct { type conn struct {
io.ReadWriteCloser io.ReadWriteCloser
relay *Relay ctx context.Context
tickChannel chan struct{}
} }
func (c conn) Read(p []byte) (int, error) { func (c conn) Read(p []byte) (int, error) {
n, err := c.ReadWriteCloser.Read(p) n, err := c.ReadWriteCloser.Read(p)
select { select {
case <-c.relay.ctx.Done(): case <-c.ctx.Done():
case c.relay.tickChannel <- struct{}{}: case c.tickChannel <- struct{}{}:
} }
return n, err // nolint: wrapcheck return n, err // nolint: wrapcheck
@@ -23,8 +27,8 @@ func (c conn) Write(p []byte) (int, error) {
n, err := c.ReadWriteCloser.Write(p) n, err := c.ReadWriteCloser.Write(p)
select { select {
case <-c.relay.ctx.Done(): case <-c.ctx.Done():
case c.relay.tickChannel <- struct{}{}: case c.tickChannel <- struct{}{}:
} }
return n, err // nolint: wrapcheck 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 { func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
eastConn = conn{ eastConn = conn{
ReadWriteCloser: eastConn, ReadWriteCloser: eastConn,
relay: r, ctx: r.ctx,
tickChannel: r.tickChannel,
} }
westConn = conn{ westConn = conn{
ReadWriteCloser: westConn, ReadWriteCloser: westConn,
relay: r, ctx: r.ctx,
tickChannel: r.tickChannel,
} }
defer func() { defer func() {
@@ -37,7 +39,7 @@ func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
wg := &sync.WaitGroup{} wg := &sync.WaitGroup{}
wg.Add(3) // nolint: gomnd wg.Add(3) // nolint: gomnd
go r.runObserver(r.ctx, wg) go r.runObserver(wg)
go r.transmit(eastConn, westConn, r.westBuffer, "west", 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 { select {
case <-r.ctx.Done(): case <-r.ctx.Done():
err = r.ctx.Err()
default:
}
select {
case r.errorChannel <- err: case r.errorChannel <- err:
default: default:
} }
} }
} }
func (r *Relay) runObserver(ctx context.Context, wg *sync.WaitGroup) { func (r *Relay) runObserver(wg *sync.WaitGroup) {
ticker := time.NewTicker(time.Second) ticker := time.NewTicker(time.Second)
defer func() { defer func() {
@@ -90,7 +97,7 @@ func (r *Relay) runObserver(ctx context.Context, wg *sync.WaitGroup) {
for { for {
select { select {
case <-ctx.Done(): case <-r.ctx.Done():
return return
case <-r.tickChannel: case <-r.tickChannel:
lastTickAt = time.Now() lastTickAt = time.Now()