added initial typing of forwarders, now only tcp

app/operator/proxy.go

@@ -14,6 +14,7 @@ func (oper *Operator) ListForwarders(ctx context.Context, params *mlbctl.ListFor
 	}
 	for _, forw := range oper.proxy.Forwarders {
 		oForw := &mlbctl.Forwarder{
+			Type:         forw.Type,
 			Lport:        forw.Lport,
 			Dport:        forw.Dport,
 			Destinations: make([]*mlbctl.Destination, 0),
@@ -32,7 +33,7 @@ func (oper *Operator) ListForwarders(ctx context.Context, params *mlbctl.ListFor
 func (oper *Operator) CreateForwarder(ctx context.Context, params *mlbctl.CreateForwarderParams) (*mlbctl.CreateForwarderResult, error) {
 	var err error
 	res := &mlbctl.CreateForwarderResult{}
-	err = oper.proxy.AddForwarder(ctx, params.Lport, params.Dport, params.Destinations...)
+	err = oper.proxy.AddForwarder(ctx, params.Type, params.Lport, params.Dport, params.Destinations...)
 	if err != err {
 		return res, err
 	}

app/rproxy/forwarder.go

@@ -11,7 +11,13 @@ import (
 	"helmet/app/logger"
 )
 
+const (
+	TCP = "tcp"
+	UDP = "udp"
+)
+
 type Forwarder struct {
+	Type   string             `json:"type"   yaml:"type"`
 	listen net.Listener       `json:"-"      yaml:"-"`
 	ctx    context.Context    `json:"-"      yaml:"-"`
 	cancel context.CancelFunc `json:"-"      yaml:"-"`
@@ -21,7 +27,7 @@ type Forwarder struct {
 	log    *logger.Logger
 }
 
-func NewForwarder(ctx context.Context, lport, dport uint32, addrs ...string) (*Forwarder, error) {
+func NewForwarder(ctx context.Context, typ string, lport, dport uint32, addrs ...string) (*Forwarder, error) {
 	ctx, cancel := context.WithCancel(context.Background())
 	forw := &Forwarder{
 		Dests:  make([]*Destination, 0),
@@ -29,6 +35,7 @@ func NewForwarder(ctx context.Context, lport, dport uint32, addrs ...string) (*F
 		Dport:  dport,
 		ctx:    ctx,
 		cancel: cancel,
+		Type:   typ,
 	}
 	id := strconv.FormatUint(uint64(lport), 10)
 	forw.log = logger.NewLogger("forwarder:" + id)
@@ -50,7 +57,7 @@ func NewForwarder(ctx context.Context, lport, dport uint32, addrs ...string) (*F
 	return forw, err
 }
 
-func (forw *Forwarder) Listen(wg *sync.WaitGroup) {
+func (forw *Forwarder) ListenTCP(wg *sync.WaitGroup) {
 	forw.log.Debugf("Start listening on %d", forw.Lport)
 	defer wg.Done()
 	for {

app/rproxy/rproxy.go

@@ -27,7 +27,7 @@ func NewProxy() *Proxy {
 	}
 }
 
-func (bal *Proxy) AddForwarder(ctx context.Context, lport, dport uint32, addrs ...string) error {
+func (bal *Proxy) AddForwarder(ctx context.Context, typ string, lport, dport uint32, addrs ...string) error {
 	var err error
 	if lport == 0 {
 		return errors.New("Zero lport")
@@ -35,13 +35,13 @@ func (bal *Proxy) AddForwarder(ctx context.Context, lport, dport uint32, addrs .
 	if dport == 0 {
 		return errors.New("Zero dport")
 	}
-	forw, err := NewForwarder(ctx, lport, dport, addrs...)
+	forw, err := NewForwarder(ctx, typ, lport, dport, addrs...)
 	if err != nil {
 		return err
 	}
 	bal.Forwarders = append(bal.Forwarders, forw)
 	bal.wg.Add(1)
-	go forw.Listen(&bal.wg)
+	go forw.ListenTCP(&bal.wg)
 	return err
 }
 
@@ -78,7 +78,7 @@ func (bal *Proxy) Start() error {
 	var err error
 	for _, forw := range bal.Forwarders {
 		bal.wg.Add(1)
-		go forw.Listen(&bal.wg)
+		go forw.ListenTCP(&bal.wg)
 	}
 	bal.wg.Wait()
 	return err