Move client IP extraction to RemoteAddr() func + some validity checks

2 files changed, 43 insertions, 19 deletions

diff --git a/proxy-go/snowflake.go b/proxy-go/snowflake.go
index 80867e7..69390d2 100644
--- a/proxy-go/snowflake.go
+++ b/proxy-go/snowflake.go
@@ -50,7 +50,6 @@ type webRTCConn struct {
 	dc *webrtc.DataChannel
 	pc *webrtc.PeerConnection
 	pr *io.PipeReader
-	client_ip string
 }
 
 func (c *webRTCConn) Read(b []byte) (int, error) {
@@ -73,6 +72,24 @@ func (c *webRTCConn) LocalAddr() net.Addr {
 }
 
 func (c *webRTCConn) RemoteAddr() net.Addr {
+	//Parse Remote SDP offer and extract client IP
+	remoteSDP := c.pc.RemoteDescription().Sdp
+	splitSDP := strings.Split(remoteSDP, "\r\n")
+	for i := range splitSDP {
+		if strings.HasPrefix(splitSDP[i], "c=")	{
+			candidSplit := strings.Split(splitSDP[i], " ")
+			if len(candidSplit) >= 3 {
+				addr := candidSplit[2]
+
+				clientIP := net.ParseIP(addr)
+				if clientIP != nil {
+					// Return client IP address
+					clientAddr := &net.IPAddr{clientIP, ""}
+					return clientAddr
+				}
+			}
+		}
+	}
 	return nil
 }
 
@@ -183,7 +200,19 @@ func datachannelHandler(conn *webRTCConn) {
 	defer conn.Close()
 	defer retToken()
 
-	wsConn, err := websocket.Dial(opt.relay + "?client_ip=" + conn.client_ip, "", opt.relay)
+	// Retrieve client IP address
+	if conn.RemoteAddr() != nil &&  conn.RemoteAddr().String() != "" {
+		// Encode client IP address in relay URL
+		q := opt.relayURL.Query()
+		clientIP := conn.RemoteAddr().String()
+		q.Set("client_ip", clientIP)
+		opt.relayURL.RawQuery = q.Encode()
+	} else {
+		log.Printf("no remote address given in websocket")
+	}
+
+
+	wsConn, err := websocket.Dial(opt.relayURL.String(), "", opt.relay)
 	if err != nil {
 		log.Printf("error dialing relay: %s", err)
 		return
@@ -236,21 +265,7 @@ func makePeerConnectionFromOffer(sdp *webrtc.SessionDescription, config *webrtc.
 			}
 		}
 
-		//Parse Remote SDP to pass client IP to the server
-		var clientIP string = ""
-		remoteSDP := pc.RemoteDescription().Sdp
-		splitSDP := strings.Split(remoteSDP, "\r\n")
-		for i := range splitSDP {
-			if strings.HasPrefix(splitSDP[i], "c=")	{
-				candidSplit := strings.Split(splitSDP[i], " ")
-				if len(candidSplit) >= 3 {
-					clientIP = candidSplit[2]
-					break
-				}
-			}
-		}
-
-		conn := &webRTCConn{pc: pc, dc: dc, pr: pr, client_ip: clientIP}
+		conn := &webRTCConn{pc: pc, dc: dc, pr: pr}
 		go datachannelHandler(conn)
 	}
 
diff --git a/server/server.go b/server/server.go
index 0928ecf..edf5977 100644
--- a/server/server.go
+++ b/server/server.go
@@ -139,10 +139,19 @@ func webSocketHandler(ws *websocket.WebSocket) {
 		handlerChan <- -1
 	}()
 
-	// Pass the address of client as the remote address of incoming connection
-	addr:= ws.ReqURL.Query().Get("client_ip")
+	// Check if client addr is a valid IP
+	addr:= ws.Request().URL.Query().Get("client_ip")
+	clientIP:= net.ParseIP(addr)
+
+	if clientIP == nil {
+		// Set client addr to empty
+		log.Printf("failed to validate client_ip: %s", addr)
+		addr = ""
+	}
 
+	// Pass the address of client as the remote address of incoming connection
 	or, err := pt.DialOr(&ptInfo, addr, ptMethodName)
+
 	if err != nil {
 		log.Printf("failed to connect to ORPort: %s", err)
 		return