Merge branch 'master' of github.com:schollz/croc

src/croc/croc.go

@@ -542,12 +542,15 @@ func (c *Client) Receive() (err error) {
 				if portToUse == "" {
 					portToUse = "9009"
 				}
-				c.Options.RelayAddress = net.JoinHostPort(discoveries[0].Address, portToUse)
-				c.ExternalIPConnected = c.Options.RelayAddress
-				c.Options.RelayAddress6 = ""
-				usingLocal = true
-
-				break
+				address := net.JoinHostPort(discoveries[0].Address, portToUse)
+				if tcp.PingServer(address) == nil {
+					log.Debugf("succesfully pinged '%s'", address)
+					c.Options.RelayAddress = address
+					c.ExternalIPConnected = c.Options.RelayAddress
+					c.Options.RelayAddress6 = ""
+					usingLocal = true
+					break
+				}
 			}
 		}
 		log.Debugf("discoveries: %+v", discoveries)

src/tcp/tcp.go

@@ -37,6 +37,7 @@ type roomMap struct {
 }
 
 var timeToRoomDeletion = 10 * time.Minute
+var pingRoom = "pinglkasjdlfjsaldjf"
 
 // Run starts a tcp listener, run async
 func Run(debugLevel, port, password string, banner ...string) (err error) {
@@ -100,8 +101,16 @@ func (s *server) run() (err error) {
 		go func(port string, connection net.Conn) {
 			c := comm.New(connection)
 			room, errCommunication := s.clientCommunication(port, c)
+			log.Debugf("room: %+v", room)
+			log.Debugf("err: %+v", errCommunication)
 			if errCommunication != nil {
 				log.Debugf("relay-%s: %s", connection.RemoteAddr().String(), errCommunication.Error())
+				connection.Close()
+				return
+			}
+			if room == pingRoom {
+				log.Debugf("got ping")
+				connection.Close()
 				return
 			}
 			for {
@@ -151,6 +160,11 @@ func (s *server) clientCommunication(port string, c *comm.Comm) (room string, er
 	if err != nil {
 		return
 	}
+	if bytes.Equal(Abytes, []byte("ping")) {
+		room = pingRoom
+		c.Send([]byte("pong"))
+		return
+	}
 	err = B.Update(Abytes)
 	if err != nil {
 		return
@@ -373,6 +387,25 @@ func pipe(conn1 net.Conn, conn2 net.Conn) {
 	}
 }
 
+func PingServer(address string) (err error) {
+	c, err := comm.NewConnection(address, 200*time.Millisecond)
+	if err != nil {
+		return
+	}
+	err = c.Send([]byte("ping"))
+	if err != nil {
+		return
+	}
+	b, err := c.Receive()
+	if err != nil {
+		return
+	}
+	if bytes.Equal(b, []byte("pong")) {
+		return nil
+	}
+	return fmt.Errorf("no pong")
+}
+
 // ConnectToTCPServer will initiate a new connection
 // to the specified address, room with optional time limit
 func ConnectToTCPServer(address, password, room string, timelimit ...time.Duration) (c *comm.Comm, banner string, ipaddr string, err error) {

src/tcp/tcp_test.go

@@ -25,6 +25,12 @@ func TestTCP(t *testing.T) {
 	log.SetLevel("error")
 	timeToRoomDeletion = 100 * time.Millisecond
 	go Run("debug", "8281", "pass123", "8282")
+	time.Sleep(100 * time.Millisecond)
+	err := PingServer("localhost:8281")
+	assert.Nil(t, err)
+	err = PingServer("localhost:8333")
+	assert.NotNil(t, err)
+
 	time.Sleep(100 * time.Millisecond)
 	c1, banner, _, err := ConnectToTCPServer("localhost:8281", "pass123", "testRoom", 1*time.Minute)
 	assert.Equal(t, banner, "8282")