Replace string with concret address struct

bzEq · Apr 26, 2024 · 3a44eb8 · 3a44eb8
1 parent deabfa0
commit 3a44eb8
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/proxy/intrinsic/client.go b/proxy/intrinsic/client.go
@@ -78,12 +78,16 @@ func (self *ClientContext) Dial(network string, addr string) (net.Conn, error) {
 }
 
 func (self *ClientContext) dialUDP(network, addr string) (net.Conn, error) {
+	raddr, err := net.ResolveUDPAddr(network, addr)
+	if err != nil {
+		return nil, err
+	}
 	local := core.MakePipe()
 	c := local[1]
 	go func() {
 		defer c.Close()
 		disp := self.router.C.(*UDPDispatcher)
-		id := disp.NewEntry(addr)
+		id := disp.NewEntry(raddr)
 		defer disp.DeleteEntry(id)
 		cp := core.NewSyncPortWithTimeout(c, nil, core.DEFAULT_UDP_TIMEOUT)
 		r, err := self.router.NewRoute(id, cp)
@@ -97,6 +101,10 @@ func (self *ClientContext) dialUDP(network, addr string) (net.Conn, error) {
 }
 
 func (self *ClientContext) dialTCP(network, addr string) (net.Conn, error) {
+	raddr, err := net.ResolveTCPAddr(network, addr)
+	if err != nil {
+		return nil, err
+	}
 	local := core.MakePipe()
 	go func() {
 		defer local[1].Close()
@@ -109,7 +117,7 @@ func (self *ClientContext) dialTCP(network, addr string) (net.Conn, error) {
 		i := Intrinsic{Func: RELAY_TCP}
 		{
 			data := &bytes.Buffer{}
-			req := TCPRequest{Addr: addr}
+			req := TCPRequest{Addr: *raddr}
 			enc := gob.NewEncoder(data)
 			if err := enc.Encode(&req); err != nil {
 				log.Println(err)
@@ -132,11 +140,11 @@ func (self *ClientContext) dialTCP(network, addr string) (net.Conn, error) {
 }
 
 type UDPDispatcher struct {
-	t core.Map[core.RouteId, string]
+	t core.Map[core.RouteId, *net.UDPAddr]
 	c uint64
 }
 
-func (self *UDPDispatcher) NewEntry(addr string) core.RouteId {
+func (self *UDPDispatcher) NewEntry(addr *net.UDPAddr) core.RouteId {
 	id := core.RouteId(atomic.AddUint64(&self.c, 1))
 	self.t.Store(id, addr)
 	return id
@@ -151,7 +159,7 @@ func (self *UDPDispatcher) Encode(id core.RouteId, data *iovec.IoVec) error {
 	if !in {
 		return fmt.Errorf("Remote address of RouteId #%d doesn't exist", id)
 	}
-	msg := UDPMessage{Id: id, Addr: raddr, Data: data.Consume()}
+	msg := UDPMessage{Id: id, Addr: *raddr, Data: data.Consume()}
 	var buf bytes.Buffer
 	enc := gob.NewEncoder(&buf)
 	if err := enc.Encode(&msg); err != nil {

diff --git a/proxy/intrinsic/server.go b/proxy/intrinsic/server.go
@@ -25,12 +25,12 @@ const (
 )
 
 type TCPRequest struct {
-	Addr string
+	Addr net.TCPAddr
 }
 
 type UDPMessage struct {
 	Id   core.RouteId
-	Addr string
+	Addr net.UDPAddr
 	Data []byte
 }
 
@@ -64,7 +64,7 @@ func (self *Server) relayUDP() error {
 				log.Println(err)
 				return
 			}
-			c, err := net.Dial("udp", msg.Addr)
+			c, err := net.Dial("udp", msg.Addr.String())
 			if err != nil {
 				log.Println(err)
 				return
@@ -128,7 +128,7 @@ func (self *Server) Run() {
 			log.Println(err)
 			return
 		}
-		if err := self.relayTCP(req.Addr); err != nil {
+		if err := self.relayTCP(req.Addr.String()); err != nil {
 			log.Println(err)
 			return
 		}