fix: added support for slave side listener

compose/network.go

@@ -17,15 +17,16 @@ const (
 )
 
 type Network struct {
-	Type    string   `yaml:"type" json:"type"`
-	Network *string  `yaml:"network,omitempty" json:"network,omitempty"`
-	Address *string  `yaml:"address,omitempty" json:"address,omitempty"`
-	Data    *Network `yaml:"data,omitempty" json:"data,omitempty"`
-	Control *Network `yaml:"control,omitempty" json:"control,omitempty"`
-	Auth    *Auth    `yaml:"auth,omitempty" json:"auth,omitempty"`
-	Crypto  *Crypto  `yaml:"crypto,omitempty" json:"crypto,omitempty"`
-	master  *reverse.Master
-	sshConn *ssh.Client
+	Type          string   `yaml:"type" json:"type"`
+	Network       *string  `yaml:"network,omitempty" json:"network,omitempty"`
+	Address       *string  `yaml:"address,omitempty" json:"address,omitempty"`
+	Data          *Network `yaml:"data,omitempty" json:"data,omitempty"`
+	Control       *Network `yaml:"control,omitempty" json:"control,omitempty"`
+	Auth          *Auth    `yaml:"auth,omitempty" json:"auth,omitempty"`
+	Crypto        *Crypto  `yaml:"crypto,omitempty" json:"crypto,omitempty"`
+	SlaveListener *bool    `yaml:"slaveListener,omitempty" json:"slaveListener,omitempty"`
+	master        *reverse.Master
+	sshConn       *ssh.Client
 }
 
 func (n *Network) setupBasicListener(listen network.ListenFunc) (net.Listener, error) {
@@ -63,8 +64,24 @@ func (n *Network) getMaster() (*reverse.Master, error) {
 	return n.master, nil
 }
 
+type slaveListenerWrapper struct {
+	*reverse.Master
+}
+
+func (slw *slaveListenerWrapper) Accept() (net.Conn, error) {
+	return slw.Master.SlaveAccept()
+}
+
 func (n *Network) setupMasterListener() (ll net.Listener, err error) {
-	return n.getMaster()
+	master, err := n.getMaster()
+	if err != nil {
+		return nil, err
+	}
+	ll = master
+	if n.SlaveListener != nil && *n.SlaveListener {
+		return &slaveListenerWrapper{master}, nil
+	}
+	return ll, nil
 }
 
 type sshWrapper struct {

compose/network_test.go

@@ -4,6 +4,8 @@ import (
 	"net"
 	"testing"
 
+	"github.com/shoriwe/fullproxy/v3/reverse"
+	"github.com/shoriwe/fullproxy/v3/utils/network"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -121,6 +123,62 @@ func TestNetwork_setupMasterListener(t *testing.T) {
 		_, err := l.setupMasterListener()
 		assert.NotNil(tt, err)
 	})
+	t.Run("SlaveListener", func(tt *testing.T) {
+		l := Network{
+			SlaveListener: new(bool),
+			Data: &Network{
+				Type:    NetworkBasic,
+				Network: new(string),
+				Address: new(string),
+			},
+			Control: &Network{
+				Type:    NetworkBasic,
+				Network: new(string),
+				Address: new(string),
+			},
+		}
+		*l.SlaveListener = true
+		*l.Data.Network = "tcp"
+		*l.Data.Address = "localhost:0"
+		*l.Control.Network = "tcp"
+		*l.Control.Address = "localhost:0"
+		// Request the slave listener
+		ll, err := l.setupMasterListener()
+		assert.Nil(tt, err)
+		defer ll.Close()
+		//
+		checkCh := make(chan struct{}, 1)
+		ssL := network.ListenAny()
+		defer ssL.Close()
+		// Start slave
+		go func() {
+			masterConn := network.Dial(l.master.Control.Addr().String())
+			s := &reverse.Slave{
+				Listener: ssL,
+				Dial:     net.Dial,
+				Master:   masterConn,
+			}
+			go s.Serve()
+			<-checkCh
+		}()
+		testMsg := []byte("TEST")
+		// Start Dial to slave listener
+		go func() {
+			conn := network.Dial(ssL.Addr().String())
+			defer conn.Close()
+			_, err := conn.Write(testMsg)
+			assert.Nil(tt, err)
+		}()
+		// Accept connection
+		conn, err := ll.Accept()
+		assert.Nil(tt, err)
+		defer conn.Close()
+		buffer := make([]byte, len(testMsg))
+		_, err = conn.Read(buffer)
+		assert.Nil(tt, err)
+		assert.Equal(tt, testMsg, buffer)
+		checkCh <- struct{}{}
+	})
 }
 
 func TestNetwork_setupSSHListener(t *testing.T) {

reverse/master.go

@@ -28,11 +28,11 @@ func (m *Master) init() (err error) {
 }
 
 func (m *Master) handle(req *Request) (conn net.Conn, err error) {
-	defer network.CloseOnError(&err, conn)
 	err = m.init()
 	if err == nil {
 		conn, err = m.cSession.Open()
 		if err == nil {
+			defer network.CloseOnError(&err, conn)
 			err = gob.NewEncoder(conn).Encode(req)
 			if err == nil {
 				var response Response