fix to include v6

refraction-networking · Jul 11, 2024 · 8926272 · 8926272
1 parent 2c8b400
commit 8926272
Showing 1 changed file with 18 additions and 10 deletions.
diff --git a/pkg/station/lib/registration_ingest.go b/pkg/station/lib/registration_ingest.go
@@ -420,8 +420,8 @@ func (rm *RegistrationManager) NewRegistrationC2SWrapper(c2sw *pb.C2SWrapper, in
 	// signed by a registration server and has overrides that should be applied
 	var dstPort = -1
 
-	// Used to apply IPv4 overrides from the registration response
-        var ipv4Override net.IP
+	// Used to apply phantom IP overrides from the registration response
+        var ipOverride net.IP
 
 	if rr := c2sw.GetRegistrationResponse(); rr != nil {
 		if rr.DstPort != nil {
@@ -433,12 +433,19 @@ func (rm *RegistrationManager) NewRegistrationC2SWrapper(c2sw *pb.C2SWrapper, in
 		if rr.GetTransportParams() != nil && !c2s.GetDisableRegistrarOverrides() {
 			c2s.TransportParams = rr.GetTransportParams()
 		}
+		if !includeV6 {
+			// apply the ipv4 address from the registration response, if rr.Ipv4Addr is not empty
+			if rr.Ipv4Addr != nil && *rr.Ipv4Addr != 0 {
+				ipv4Bytes := make([]byte, 4)
+				binary.BigEndian.PutUint32(ipv4Bytes, *rr.Ipv4Addr)
+			        ipOverride = net.IP(ipv4Bytes)
+			}
+		} else {
+			// apply the ipv6 address from the registration response, if rr.Ipv6Addr is not empty
+			if rr.Ipv6Addr != nil {
+                                ipOverride = net.IP(rr.Ipv6Addr)
+                        }
 
-		// apply the ip addresses from the registration response, if the Ipv4Addr is not empty
-		if rr.Ipv4Addr != nil && *rr.Ipv4Addr != 0 {
-			ipv4Bytes := make([]byte, 4)
-	                binary.BigEndian.PutUint32(ipv4Bytes, *rr.Ipv4Addr)
-		        ipv4Override = net.IP(ipv4Bytes)
 		}
 
 	}
@@ -448,9 +455,10 @@ func (rm *RegistrationManager) NewRegistrationC2SWrapper(c2sw *pb.C2SWrapper, in
 		return nil, fmt.Errorf("failed to build registration: %w", err)
 	}
 
-	if ipv4Override != nil {
-		// If the Ipv4Addr from the registration response is not empty, use it to override the IPv4 that the station derived
-		reg.PhantomIp = ipv4Override
+	if ipOverride != nil {
+		// If the ipOverride (which is populated by Ipv4Addr or Ipv6Addr from the registration response) 
+		// is not empty, use it to override the phantom IP that the station derived
+		reg.PhantomIp = ipOverride
         }
 
 	clientAddr := net.IP(c2sw.GetRegistrationAddress())