diff --git a/server/federation.go b/server/federation.go index 9f70c56..c85ff3f 100644 --- a/server/federation.go +++ b/server/federation.go @@ -73,7 +73,7 @@ func (s *Server) getNumSubs() int64 { // pings it, so we can determine our own external IP address. func (s *Server) getAndSetExternalIp(msg *pb.ServerMessage) error { log.Println(msg) - myIp, err := UDPPing(msg.Address, msg.Port) + myIp, _, err := UDPPing(msg.Address, msg.Port) if err != nil { return err } diff --git a/server/udp.go b/server/udp.go index 65eb18a..734025f 100644 --- a/server/udp.go +++ b/server/udp.go @@ -149,47 +149,52 @@ func (pong *SPVPong) DecodeAddress() net.IP { ) } +func (pong *SPVPong) DecodeCountry() string { + return pb.Location_Country_name[int32(pong.country)] +} + // UDPPing sends a ping over udp to another hub and returns the ip address of // this hub. -func UDPPing(ip, port string) (net.IP, error) { +func UDPPing(ip, port string) (net.IP, string, error) { address := ip + ":" + port addr, err := net.ResolveUDPAddr("udp", address) if err != nil { - return net.IP{}, err + return net.IP{}, "", err } conn, err := net.DialUDP("udp", nil, addr) if err != nil { - return net.IP{}, err + return net.IP{}, "", err } defer conn.Close() _, err = conn.Write(encodeSPVPing()) if err != nil { - return net.IP{}, err + return net.IP{}, "", err } buffer := make([]byte, maxBufferSize) deadline := time.Now().Add(time.Second) err = conn.SetReadDeadline(deadline) if err != nil { - return net.IP{}, err + return net.IP{}, "", err } n, _, err := conn.ReadFromUDP(buffer) if err != nil { - return net.IP{}, err + return net.IP{}, "", err } pong := decodeSPVPong(buffer[:n]) if pong == nil { - return net.IP{}, errors.Base("Pong decoding failed") + return net.IP{}, "", errors.Base("Pong decoding failed") } myAddr := pong.DecodeAddress() + country := pong.DecodeCountry() - return myAddr, nil + return myAddr, country, nil } // UDPServer is a goroutine that starts an udp server that implements the hubs diff --git a/server/udp_test.go b/server/udp_test.go index 9cc3d9e..5f1071f 100644 --- a/server/udp_test.go +++ b/server/udp_test.go @@ -13,12 +13,14 @@ func TestUDPPing(t *testing.T) { args.StartUDP = false tests := []struct { - name string - want string + name string + wantIP string + wantCountry string } { { - name: "Get the right ip from production server.", - want: "SETME", + name: "Get the right ip from production server.", + wantIP: "SETME", + wantCountry: "US", }, } @@ -28,7 +30,8 @@ func TestUDPPing(t *testing.T) { toAddr := "spv16.lbry.com" toPort := "50001" - ip, err := UDPPing(toAddr, toPort) + ip, country, err := UDPPing(toAddr, toPort) + gotCountry := country if err != nil { log.Println(err) } @@ -41,11 +44,14 @@ func TestUDPPing(t *testing.T) { digIP := strings.TrimSpace(string(res)) udpIP := ip.String() - tt.want = digIP + tt.wantIP = digIP - got1 := udpIP - if got1 != tt.want { - t.Errorf("got: '%s', want: '%s'\n", got1, tt.want) + gotIP := udpIP + if gotIP != tt.wantIP { + t.Errorf("got: '%s', want: '%s'\n", gotIP, tt.wantIP) + } + if gotCountry != tt.wantCountry { + t.Errorf("got: '%s', want: '%s'\n", gotCountry, tt.wantCountry) } }) }