fix: enhance proxy stability

control/control.go

@@ -3,9 +3,6 @@ package control
 import (
 	"context"
 	"encoding/json"
-	"net"
-	"sync/atomic"
-	"time"
 
 	"github.com/pkg/errors"
 	"github.com/xtaci/smux"
@@ -13,11 +10,39 @@ import (
 )
 
 type Control struct {
-	encoder    *json.Encoder
-	decoder    *json.Decoder
-	stream     *smux.Stream
-	sess       *smux.Session
-	proxyClock int64
+	encoder *json.Encoder
+	decoder *json.Decoder
+	stream  *smux.Stream
+}
+
+func (c *Control) Init(ctx context.Context, sess *smux.Session, serverMode bool) error {
+	config := smux.DefaultConfig()
+	config.Version = 2
+
+	logger.Debug(ctx, "creating control stream")
+
+	var (
+		stream *smux.Stream
+		err    error
+	)
+
+	if serverMode {
+		stream, err = sess.AcceptStream()
+		if err != nil {
+			return errors.WithStack(err)
+		}
+	} else {
+		stream, err = sess.OpenStream()
+		if err != nil {
+			return errors.WithStack(err)
+		}
+	}
+
+	c.stream = stream
+	c.decoder = json.NewDecoder(stream)
+	c.encoder = json.NewEncoder(stream)
+
+	return nil
 }
 
 func (c *Control) AuthRequest(credentials interface{}) (bool, error) {
@@ -39,89 +64,82 @@ func (c *Control) AuthRequest(credentials interface{}) (bool, error) {
 	return authResPayload.Success, nil
 }
 
-func (c *Control) Proxy(ctx context.Context, network, address string) (net.Conn, error) {
-	var (
-		stream *smux.Stream
-		err    error
-	)
-
-	requestID := atomic.AddInt64(&c.proxyClock, 1)
-
+func (c *Control) ProxyReq(ctx context.Context, network, address string) error {
 	req := NewMessage(TypeProxyRequest, &ProxyRequestPayload{
-		RequestID: requestID,
-		Network:   network,
-		Address:   address,
+		Network: network,
+		Address: address,
 	})
 
-	ctx = logger.With(ctx, logger.F("network", network), logger.F("address", address))
-
-	logger.Debug(ctx, "proxying")
-
 	if err := c.Write(req); err != nil {
-		return nil, errors.WithStack(err)
+		return errors.WithStack(err)
 	}
 
-	logger.Debug(ctx, "opening stream")
-
-	stream, err = c.sess.AcceptStream()
-	if err != nil {
-		return nil, errors.WithStack(err)
-	}
-
-	go func() {
-		<-ctx.Done()
-
-		req := NewMessage(TypeCloseProxy, &CloseProxyPayload{
-			RequestID: requestID,
-		})
-
-		if err := c.Write(req); err != nil {
-			logger.Error(ctx, "error while closing proxy", logger.E(err))
-		}
-
-		logger.Debug(ctx, "closing proxy conn")
-		stream.Close()
-	}()
-
-	return stream, nil
+	return nil
 }
 
 func (c *Control) Listen(ctx context.Context, handlers Handlers) error {
-	for {
-		logger.Debug(ctx, "reading next message")
+	errChan := make(chan error)
+	msgChan := make(chan *Message)
+	dieChan := c.stream.GetDieCh()
 
-		req, err := c.Read()
-		if err != nil {
-			return errors.WithStack(err)
-		}
+	go func(msgChan chan *Message, errChan chan error) {
+		for {
+			logger.Debug(ctx, "reading next message")
 
-		go func() {
-			subCtx := logger.With(ctx, logger.F("messageType", req.Type))
-
-			handler, exists := handlers[req.Type]
-			if !exists {
-				logger.Error(subCtx, "no message handler registered")
-
-				return
-			}
-
-			res, err := handler(subCtx, req)
+			msg, err := c.Read()
 			if err != nil {
-				logger.Error(subCtx, "error while handling message", logger.E(err))
+				errChan <- errors.WithStack(err)
+
+				close(errChan)
+				close(msgChan)
 
 				return
 			}
 
-			if res == nil {
-				return
-			}
+			msgChan <- msg
+		}
+	}(msgChan, errChan)
 
-			if err := c.Write(res); err != nil {
-				logger.Error(subCtx, "error while write message response", logger.E(err))
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
 
-				return
-			}
-		}()
+		case <-dieChan:
+			return errors.WithStack(ErrStreamClosed)
+
+		case err := <-errChan:
+			return errors.WithStack(err)
+
+		case msg := <-msgChan:
+			go func() {
+				subCtx := logger.With(ctx, logger.F("messageType", msg.Type))
+
+				handler, exists := handlers[msg.Type]
+				if !exists {
+					logger.Error(subCtx, "no message handler registered")
+
+					return
+				}
+
+				res, err := handler(subCtx, msg)
+				if err != nil {
+					logger.Error(subCtx, "error while handling message", logger.E(err))
+
+					return
+				}
+
+				if res == nil {
+					return
+				}
+
+				if err := c.Write(res); err != nil {
+					logger.Error(subCtx, "error while write message response", logger.E(err))
+
+					return
+				}
+			}()
+		}
 	}
 }
 
@@ -164,10 +182,6 @@ func (c *Control) read(m *Message) error {
 }
 
 func (c *Control) write(m *Message) error {
-	if err := c.stream.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
-		return errors.WithStack(err)
-	}
-
 	if err := c.encoder.Encode(m); err != nil {
 		return errors.WithStack(err)
 	}
@@ -175,11 +189,6 @@ func (c *Control) write(m *Message) error {
 	return nil
 }
 
-func New(sess *smux.Session, controlStream *smux.Stream) *Control {
-	return &Control{
-		encoder: json.NewEncoder(controlStream),
-		decoder: json.NewDecoder(controlStream),
-		sess:    sess,
-		stream:  controlStream,
-	}
+func New() *Control {
+	return &Control{}
 }

control/error.go

@@ -3,5 +3,6 @@ package control
 import "errors"
 
 var (
+	ErrStreamClosed      = errors.New("stream closed")
 	ErrUnexpectedMessage = errors.New("unexpected message")
 )

control/proxy.go

@@ -1,9 +1,8 @@
 package control
 
 type ProxyRequestPayload struct {
-	RequestID int64  `json:"i"`
-	Network   string `json:"n"`
-	Address   string `json:"a"`
+	Network string `json:"n"`
+	Address string `json:"a"`
 }
 
 type CloseProxyPayload struct {