From eb3af56ccf96d57b430e7b60ef6d33f97b253320 Mon Sep 17 00:00:00 2001
From: Bradley T Lunsford <lunsford94@gmail.com>
Date: Fri, 20 Oct 2017 14:51:30 -0700
Subject: [PATCH] adding wait funtionality: if waiting, will send a different
 identification to the server

---
 connect.go | 17 ++++++++++++++++-
 main.go    |  2 ++
 relay.go   |  4 ++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/connect.go b/connect.go
index 47feb7f8..0d6aa700 100644
--- a/connect.go
+++ b/connect.go
@@ -27,6 +27,7 @@ type Connection struct {
 	IsSender            bool
 	Debug               bool
 	DontEncrypt         bool
+        Wait                bool
 	bars                []*uiprogress.Bar
 	rate                int
 }
@@ -41,6 +42,7 @@ func NewConnection(flags *Flags) *Connection {
 	c := new(Connection)
 	c.Debug = flags.Debug
 	c.DontEncrypt = flags.DontEncrypt
+        c.Wait = flags.Wait
 	c.Server = flags.Server
 	c.Code = flags.Code
 	c.NumberOfConnections = flags.NumberOfConnections
@@ -160,6 +162,7 @@ func (c *Connection) runClient() error {
 	}
 	gotOK := false
 	gotResponse := false
+        notPresent := false
 	for id := 0; id < c.NumberOfConnections; id++ {
 		go func(id int) {
 			defer wg.Done()
@@ -182,7 +185,11 @@ func (c *Connection) runClient() error {
 				sendMessage("s."+c.HashedCode+"."+hex.EncodeToString(encryptedMetaData)+"-"+salt+"-"+iv, connection)
 			} else {
 				logger.Debugf("telling relay: %s", "r."+c.Code)
-				sendMessage("r."+c.HashedCode+".0.0.0", connection)
+                                if c.Wait {
+				    sendMessage("r."+c.HashedCode+".0.0.0", connection)
+                                } else {
+				    sendMessage("c."+c.HashedCode+".0.0.0", connection)
+                                }
 			}
 			if c.IsSender { // this is a sender
 				logger.Debug("waiting for ok from relay")
@@ -201,6 +208,10 @@ func (c *Connection) runClient() error {
 				message = receiveMessage(connection)
 				m := strings.Split(message, "-")
 				encryptedData, salt, iv, sendersAddress := m[0], m[1], m[2], m[3]
+                                if sendersAddress == "0.0.0.0" {
+                                        notPresent = true
+                                        return
+                                }
 				encryptedBytes, err := hex.DecodeString(encryptedData)
 				if err != nil {
 					log.Error(err)
@@ -252,6 +263,10 @@ func (c *Connection) runClient() error {
 	wg.Wait()
 
 	if !c.IsSender {
+                if notPresent {
+                    fmt.Println("Sender/Code not present")
+                    return nil
+                }
 		if !gotOK {
 			return errors.New("Transfer interrupted")
 		}
diff --git a/main.go b/main.go
index fbc719ec..199aa509 100644
--- a/main.go
+++ b/main.go
@@ -15,6 +15,7 @@ var oneGigabytePerSecond = 1000000 // expressed as kbps
 type Flags struct {
 	Relay               bool
 	Debug               bool
+	Wait               bool
 	DontEncrypt         bool
 	Server              string
 	File                string
@@ -37,6 +38,7 @@ croc version ` + version + `
 	flags := new(Flags)
 	flag.BoolVar(&flags.Relay, "relay", false, "run as relay")
 	flag.BoolVar(&flags.Debug, "debug", false, "debug mode")
+	flag.BoolVar(&flags.Wait, "wait", false, "wait for code to be sent")
 	flag.StringVar(&flags.Server, "server", "cowyo.com", "address of relay server")
 	flag.StringVar(&flags.File, "send", "", "file to send")
 	flag.StringVar(&flags.Code, "code", "", "use your own code phrase")
diff --git a/relay.go b/relay.go
index ef35c207..0e2d5561 100644
--- a/relay.go
+++ b/relay.go
@@ -156,6 +156,10 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
 				}
 			}
 			r.connections.RUnlock()
+                        if connectionType == "c" {
+		                sendMessage("0-0-0-0.0.0.0", connection)
+                                return
+                        }
 			time.Sleep(100 * time.Millisecond)
 		}
 		// send  meta data