Simplify subnet preference code considerably

This commit is contained in:
Justin Li 2014-07-17 14:15:05 -04:00
parent 794657274a
commit 6b089342d1

View file

@ -212,13 +212,13 @@ func appendPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, announcer *
break break
} }
if peer.ID == announcer.ID || peer.UserID != 0 && peer.UserID == announcer.UserID { if peersEquivalent(&peer, announcer) {
continue continue
} }
if ip := peer.IP.To4(); ip != nil { if peer.IP.To4() != nil {
ipv4s = append(ipv4s, peer) ipv4s = append(ipv4s, peer)
} else if ip := peer.IP.To16(); ip != nil { } else if peer.IP.To16() != nil {
ipv6s = append(ipv6s, peer) ipv6s = append(ipv6s, peer)
} }
@ -229,52 +229,41 @@ func appendPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, announcer *
} }
func appendSubnetPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, announcer *models.Peer, peers models.PeerMap, wanted int) (models.PeerList, models.PeerList) { func appendSubnetPeers(ipv4s, ipv6s models.PeerList, ann *models.Announce, announcer *models.Peer, peers models.PeerMap, wanted int) (models.PeerList, models.PeerList) {
var ( var subnet net.IPNet
subnet net.IPNet
ipv4Subnet bool
ipv6Subnet bool
)
if aip := announcer.IP.To4(); len(aip) == 4 { if aip := announcer.IP.To4(); aip != nil {
subnet = net.IPNet{aip, net.CIDRMask(ann.Config.PreferredIPv4Subnet, 32)} subnet = net.IPNet{aip, net.CIDRMask(ann.Config.PreferredIPv4Subnet, 32)}
ipv4Subnet = true } else if aip := announcer.IP.To16(); aip != nil {
} else if aip := announcer.IP.To16(); len(aip) == 16 {
subnet = net.IPNet{aip, net.CIDRMask(ann.Config.PreferredIPv6Subnet, 128)} subnet = net.IPNet{aip, net.CIDRMask(ann.Config.PreferredIPv6Subnet, 128)}
ipv6Subnet = true } else {
panic("impossible: missing IP")
} }
// Iterate over the peers twice: first add only peers in the same subnet and // Iterate over the peers twice: first add only peers in the same subnet and
// if we still need more peers grab any that haven't already been added. // if we still need more peers grab any that haven't already been added.
count := 0 count := 0
for _, peersLeftInSubnet := range [2]bool{true, false} { for _, checkInSubnet := range [2]bool{true, false} {
for _, peer := range peers { for _, peer := range peers {
if count >= wanted { if count >= wanted {
break break
} }
if peer.ID == announcer.ID || peer.UserID != 0 && peer.UserID == announcer.UserID { if peersEquivalent(&peer, announcer) || checkInSubnet != subnet.Contains(peer.IP) {
continue continue
} }
if ip := peer.IP.To4(); len(ip) == 4 { if peer.IP.To4() != nil {
if peersLeftInSubnet && ipv4Subnet && subnet.Contains(ip) {
ipv4s = append(ipv4s, peer) ipv4s = append(ipv4s, peer)
count++ } else if peer.IP.To16() != nil {
} else if !peersLeftInSubnet && !subnet.Contains(ip) {
ipv4s = append(ipv4s, peer)
count++
}
} else if ip := peer.IP.To16(); len(ip) == 16 {
if peersLeftInSubnet && ipv6Subnet && subnet.Contains(ip) {
ipv6s = append(ipv6s, peer) ipv6s = append(ipv6s, peer)
count++
} else if !peersLeftInSubnet && !subnet.Contains(ip) {
ipv6s = append(ipv6s, peer)
count++
}
} }
count++
} }
} }
return ipv4s, ipv6s return ipv4s, ipv6s
} }
func peersEquivalent(a, b *models.Peer) bool {
return a.ID == b.ID || a.UserID != 0 && a.UserID == b.UserID
}