feat: better proxy handling

control/control.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"net"
+	"sync/atomic"
 	"time"
 
 	"github.com/pkg/errors"
@@ -12,10 +13,11 @@ import (
 )
 
 type Control struct {
-	encoder *json.Encoder
-	decoder *json.Decoder
-	stream  *smux.Stream
-	sess    *smux.Session
+	encoder    *json.Encoder
+	decoder    *json.Decoder
+	stream     *smux.Stream
+	sess       *smux.Session
+	proxyClock int64
 }
 
 func (c *Control) AuthRequest(credentials interface{}) (bool, error) {
@@ -37,17 +39,18 @@ func (c *Control) AuthRequest(credentials interface{}) (bool, error) {
 	return authResPayload.Success, nil
 }
 
-type CloseStream func()
-
 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)
+
 	req := NewMessage(TypeProxyRequest, &ProxyRequestPayload{
-		Network: network,
-		Address: address,
+		RequestID: requestID,
+		Network:   network,
+		Address:   address,
 	})
 
 	ctx = logger.With(ctx, logger.F("network", network), logger.F("address", address))
@@ -65,6 +68,21 @@ func (c *Control) Proxy(ctx context.Context, network, address string) (net.Conn,
 		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
 }
 

control/message.go

@@ -10,6 +10,7 @@ const (
 	TypeAuthRequest  MessageType = "auth-req"
 	TypeAuthResponse MessageType = "auth-res"
 	TypeProxyRequest MessageType = "proxy-req"
+	TypeCloseProxy   MessageType = "close-proxy"
 )
 
 type MessageType string
@@ -61,6 +62,8 @@ func unmarshalPayload(mType MessageType, data []byte) (interface{}, error) {
 		payload = &AuthResponsePayload{}
 	case TypeProxyRequest:
 		payload = &ProxyRequestPayload{}
+	case TypeCloseProxy:
+		payload = &CloseProxyPayload{}
 	default:
 		return nil, errors.Wrapf(ErrUnexpectedMessage, "unexpected message type '%s'", mType)
 	}

control/proxy.go

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