From d0c2dc1be82c47ae8018a5b0694311518b7860a8 Mon Sep 17 00:00:00 2001
From: Zack Scholl <zack.scholl@gmail.com>
Date: Mon, 18 Nov 2019 13:09:17 -0800
Subject: [PATCH] use one-byte pings

---
 src/croc/croc.go | 16 ++++++++++++++--
 src/tcp/tcp.go   | 44 ++++++++++++++++++++++++++++++++++++++------
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/src/croc/croc.go b/src/croc/croc.go
index c7aa569..6877018 100644
--- a/src/croc/croc.go
+++ b/src/croc/croc.go
@@ -277,11 +277,13 @@ func (c *Client) transferOverLocalRelay(options TransferOptions, errchan chan<-
 		// not really an error because it will try to connect over the actual relay
 		return
 	}
-	log.Debugf("connection established: %+v", conn)
+	log.Debugf("local connection established: %+v", conn)
 	for {
 		data, _ := conn.Receive()
 		if bytes.Equal(data, []byte("handshake")) {
 			break
+		} else {
+			log.Debugf("instead of handshake got: %s", data)
 		}
 	}
 	c.conn[0] = conn
@@ -339,6 +341,7 @@ func (c *Client) Send(options TransferOptions) (err error) {
 		}
 		log.Debugf("connection established: %+v", conn)
 		for {
+			log.Debug("waiting for bytes")
 			data, errConn := conn.Receive()
 			if errConn != nil {
 				log.Debugf("[%+v] had error: %s", conn, errConn.Error())
@@ -360,10 +363,14 @@ func (c *Client) Send(options TransferOptions) (err error) {
 				conn.Send(bips)
 			} else if bytes.Equal(data, []byte("handshake")) {
 				break
+			} else if bytes.Equal(data, []byte{1}) {
+				log.Debug("got ping")
+				continue
 			} else {
 				log.Debugf("[%+v] got weird bytes: %+v", conn, data)
 				// throttle the reading
 				time.Sleep(100 * time.Millisecond)
+				break
 			}
 		}
 
@@ -434,11 +441,12 @@ func (c *Client) Receive() (err error) {
 		err = errors.Wrap(err, fmt.Sprintf("could not connect to %s", c.Options.RelayAddress))
 		return
 	}
-	log.Debugf("connection established: %+v", c.conn[0])
+	log.Debugf("receiver connection established: %+v", c.conn[0])
 
 	if !usingLocal && !c.Options.DisableLocal {
 		// ask the sender for their local ips and port
 		// and try to connect to them
+		log.Debug("sending ips?")
 		var data []byte
 		c.conn[0].Send([]byte("ips?"))
 		data, err = c.conn[0].Receive()
@@ -1049,6 +1057,10 @@ func (c *Client) receiveData(i int) {
 		if err != nil {
 			break
 		}
+		if bytes.Equal(data, []byte{1}) {
+			log.Debug("got ping")
+			continue
+		}
 
 		data, err = crypt.Decrypt(data, c.Key)
 		if err != nil {
diff --git a/src/tcp/tcp.go b/src/tcp/tcp.go
index f60f0e3..138e38f 100644
--- a/src/tcp/tcp.go
+++ b/src/tcp/tcp.go
@@ -96,15 +96,47 @@ func (s *server) run() (err error) {
 		}
 		log.Debugf("client %s connected", connection.RemoteAddr().String())
 		go func(port string, connection net.Conn) {
-			errCommunication := s.clientCommuncation(port, comm.New(connection))
+			c := comm.New(connection)
+			room, errCommunication := s.clientCommuncation(port, c)
 			if errCommunication != nil {
 				log.Warnf("relay-%s: %s", connection.RemoteAddr().String(), errCommunication.Error())
 			}
+			for {
+				// check connection
+				log.Debugf("checking connection of room %s for %+v", room, c)
+				deleteIt := false
+				s.rooms.Lock()
+				if _, ok := s.rooms.rooms[room]; !ok {
+					log.Debug("room is gone")
+					s.rooms.Unlock()
+					return
+				}
+				log.Debugf("room: %+v", s.rooms.rooms[room])
+				if s.rooms.rooms[room].first != nil && s.rooms.rooms[room].second != nil {
+					log.Debug("rooms ready")
+					s.rooms.Unlock()
+					break
+				} else {
+					if s.rooms.rooms[room].first != nil {
+						errSend := s.rooms.rooms[room].first.Send([]byte{1})
+						if errSend != nil {
+							log.Debug(errSend)
+							deleteIt = true
+						}
+					}
+				}
+				s.rooms.Unlock()
+				if deleteIt {
+					s.deleteRoom(room)
+					break
+				}
+				time.Sleep(1 * time.Second)
+			}
 		}(s.port, connection)
 	}
 }
 
-func (s *server) clientCommuncation(port string, c *comm.Comm) (err error) {
+func (s *server) clientCommuncation(port string, c *comm.Comm) (room string, err error) {
 	log.Debugf("waiting for password")
 	passwordBytes, err := c.Receive()
 	if err != nil {
@@ -133,7 +165,7 @@ func (s *server) clientCommuncation(port string, c *comm.Comm) (err error) {
 	if err != nil {
 		return
 	}
-	room := string(roomBytes)
+	room = string(roomBytes)
 
 	s.rooms.Lock()
 	// create the room if it is new
@@ -151,7 +183,7 @@ func (s *server) clientCommuncation(port string, c *comm.Comm) (err error) {
 			return
 		}
 		log.Debugf("room %s has 1", room)
-		return nil
+		return
 	}
 	if s.rooms.rooms[room].full {
 		s.rooms.Unlock()
@@ -161,7 +193,7 @@ func (s *server) clientCommuncation(port string, c *comm.Comm) (err error) {
 			s.deleteRoom(room)
 			return
 		}
-		return nil
+		return
 	}
 	log.Debugf("room %s has 2", room)
 	s.rooms.rooms[room] = roomInfo{
@@ -195,7 +227,7 @@ func (s *server) clientCommuncation(port string, c *comm.Comm) (err error) {
 
 	// delete room
 	s.deleteRoom(room)
-	return nil
+	return
 }
 
 func (s *server) deleteRoom(room string) {