From 4f76cc3ac35a2381e792c09f7b1d800a28058d6f Mon Sep 17 00:00:00 2001
From: Abhishek Kondur <abhi281342@gmail.com>
Date: Tue, 15 Nov 2022 09:31:31 +0530
Subject: [PATCH] forward unknown pkts to all interfaces

---
 logic/peers.go              |  1 +
 nm-proxy/common/common.go   |  2 +-
 nm-proxy/manager/manager.go | 12 +++++++++++-
 nm-proxy/peer/peer.go       |  3 +--
 nm-proxy/server/server.go   | 13 +++++++++++++
 5 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/logic/peers.go b/logic/peers.go
index ed625dd1..69b180f9 100644
--- a/logic/peers.go
+++ b/logic/peers.go
@@ -104,6 +104,7 @@ func GetPeersForProxy(node *models.Node, onlyPeers bool) (manager.ManagerPayload
 					peerConfMap[peer.PublicKey] = manager.PeerConf{
 						IsRelayed: true,
 						RelayedTo: relayTo,
+						Address:   peer.PrimaryAddress(),
 					}
 				}
 
diff --git a/nm-proxy/common/common.go b/nm-proxy/common/common.go
index 2cd2ced1..abde20b5 100644
--- a/nm-proxy/common/common.go
+++ b/nm-proxy/common/common.go
@@ -74,7 +74,7 @@ var WgIFaceMap = make(map[string]map[string]*Conn)
 
 var PeerKeyHashMap = make(map[string]RemotePeer)
 
-var WgIfaceKeyMap = make(map[string]struct{})
+var WgIfaceKeyMap = make(map[string]RemotePeer)
 
 var RelayPeerMap = make(map[string]map[string]RemotePeer)
 
diff --git a/nm-proxy/manager/manager.go b/nm-proxy/manager/manager.go
index 3b92fabb..cfcc6db3 100644
--- a/nm-proxy/manager/manager.go
+++ b/nm-proxy/manager/manager.go
@@ -13,6 +13,7 @@ import (
 	"github.com/gravitl/netmaker/nm-proxy/common"
 	"github.com/gravitl/netmaker/nm-proxy/packet"
 	peerpkg "github.com/gravitl/netmaker/nm-proxy/peer"
+	"github.com/gravitl/netmaker/nm-proxy/proxy"
 	"github.com/gravitl/netmaker/nm-proxy/wg"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 )
@@ -225,7 +226,16 @@ func (m *ManagerAction) AddInterfaceToProxy() error {
 		return err
 	}
 	log.Printf("wg: %+v\n", wgInterface)
-
+	wgListenAddr, err := proxy.GetInterfaceListenAddr(wgInterface.Port)
+	if err != nil {
+		log.Println("failed to get wg listen addr: ", err)
+		return err
+	}
+	common.WgIfaceKeyMap[fmt.Sprintf("%x", md5.Sum([]byte(wgInterface.Device.PublicKey.String())))] = common.RemotePeer{
+		PeerKey:   wgInterface.Device.PublicKey.String(),
+		Interface: wgInterface.Name,
+		Endpoint:  wgListenAddr,
+	}
 	for _, peerI := range m.Payload.Peers {
 		peerConf := m.Payload.PeerMap[peerI.PublicKey.String()]
 		if peerI.Endpoint == nil && !(peerConf.IsAttachedExtClient || peerConf.IsExtClient) {
diff --git a/nm-proxy/peer/peer.go b/nm-proxy/peer/peer.go
index f8072abc..c7c2e6bd 100644
--- a/nm-proxy/peer/peer.go
+++ b/nm-proxy/peer/peer.go
@@ -1,7 +1,6 @@
 package peer
 
 import (
-	"crypto/md5"
 	"errors"
 	"fmt"
 	"log"
@@ -105,6 +104,6 @@ func AddNewPeer(wgInterface *wg.WGIface, peer *wgtypes.PeerConfig,
 		common.WgIFaceMap[wgInterface.Name] = make(map[string]*common.Conn)
 		common.WgIFaceMap[wgInterface.Name][peer.PublicKey.String()] = &peerConn
 	}
-	common.WgIfaceKeyMap[fmt.Sprintf("%x", md5.Sum([]byte(wgInterface.Device.PublicKey.String())))] = struct{}{}
+
 	return nil
 }
diff --git a/nm-proxy/server/server.go b/nm-proxy/server/server.go
index 96a2407f..5454e3e2 100644
--- a/nm-proxy/server/server.go
+++ b/nm-proxy/server/server.go
@@ -125,9 +125,22 @@ func (p *ProxyServer) Listen(ctx context.Context) {
 
 					}
 				}
+				continue
 
 			}
+			// forward to all interfaces
+			for _, ifaceCfg := range common.WgIfaceKeyMap {
+				log.Println("###--------> Forwarding Unknown PKT to ", ifaceCfg.Interface)
+				conn, err := net.DialUDP("udp", source, ifaceCfg.Endpoint)
+				if err == nil {
+					_, err := conn.Write(buffer[:n])
+					if err != nil {
+						log.Println("Failed to forward the unknown pkt to ifcace: ", ifaceCfg.Interface, err)
+					}
+					conn.Close()
+				}
 
+			}
 		}
 
 	}