From 034e0122273a7bcfb70c7cf32bd5a883809597f5 Mon Sep 17 00:00:00 2001
From: Zack Scholl <zack.scholl@gmail.com>
Date: Thu, 22 Oct 2020 10:09:04 -0700
Subject: [PATCH] allow pinging

---
 src/tcp/tcp.go      | 33 +++++++++++++++++++++++++++++++++
 src/tcp/tcp_test.go |  6 ++++++
 2 files changed, 39 insertions(+)

diff --git a/src/tcp/tcp.go b/src/tcp/tcp.go
index 38f5dbe..39664be 100644
--- a/src/tcp/tcp.go
+++ b/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) {
diff --git a/src/tcp/tcp_test.go b/src/tcp/tcp_test.go
index fda31f2..ccceba7 100644
--- a/src/tcp/tcp_test.go
+++ b/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")