Merge pull request #26 from lummie/master

Issue #25

.gitignore

@@ -0,0 +1,3 @@
+# IDEs
+.idea/
+

connect.go

@@ -79,21 +79,16 @@ func (c *Connection) Run() error {
 		// check code
 		goodCode := true
 		m := strings.Split(c.Code, "-")
-		log.Debug(m)
 		numThreads, errParse := strconv.Atoi(m[0])
 		if len(m) < 2 {
 			goodCode = false
-			log.Debug("code too short")
 		} else if numThreads > MAX_NUMBER_THREADS || numThreads < 1 || (forceSingleThreaded && numThreads != 1) {
 			c.NumberOfConnections = MAX_NUMBER_THREADS
 			goodCode = false
-			log.Debug("incorrect number of threads")
 		} else if errParse != nil {
 			goodCode = false
-			log.Debug("problem parsing threads")
 		}
 		log.Debug(m)
-		log.Debug(goodCode)
 		if !goodCode {
 			if c.IsSender {
 				if forceSingleThreaded {
@@ -135,8 +130,6 @@ func (c *Connection) Run() error {
 		if err != nil {
 			return err
 		}
-		fmt.Printf("Sending %d byte file named '%s'\n", c.File.Size, c.File.Name)
-		fmt.Printf("Code is: %s\n", c.Code)
 	}
 
 	return c.runClient()
@@ -160,6 +153,7 @@ func (c *Connection) runClient() error {
 	}
 	gotOK := false
 	gotResponse := false
+	gotConnectionInUse := false
 	for id := 0; id < c.NumberOfConnections; id++ {
 		go func(id int) {
 			defer wg.Done()
@@ -187,6 +181,10 @@ func (c *Connection) runClient() error {
 			if c.IsSender { // this is a sender
 				logger.Debug("waiting for ok from relay")
 				message = receiveMessage(connection)
+				if message == "no" {
+					fmt.Println("The specifed code is already in use by a sender.")
+					gotConnectionInUse = true
+				} else {
 					logger.Debug("got ok from relay")
 					if id == 0 {
 						fmt.Printf("\nSending (->%s)..\n", message)
@@ -196,9 +194,15 @@ func (c *Connection) runClient() error {
 					// Write data from file
 					logger.Debug("send file")
 					c.sendFile(id, connection)
+					fmt.Println("File sent.")
+				}
 			} else { // this is a receiver
 				logger.Debug("waiting for meta data from sender")
 				message = receiveMessage(connection)
+				if message == "no" {
+					fmt.Println("The specifed code is already in use by a receiver.")
+					gotConnectionInUse = true
+				} else {
 					m := strings.Split(message, "-")
 					encryptedData, salt, iv, sendersAddress := m[0], m[1], m[2], m[3]
 					encryptedBytes, err := hex.DecodeString(encryptedData)
@@ -241,17 +245,22 @@ func (c *Connection) runClient() error {
 					} else {
 						sendMessage("ok", connection)
 						logger.Debug("receive file")
-					if id == 0 {
 						fmt.Printf("\n\nReceiving (<-%s)..\n", sendersAddress)
-					}
 						c.receiveFile(id, connection)
 					}
 				}
+			}
 		}(id)
 	}
 	wg.Wait()
 
-	if !c.IsSender {
+	if gotConnectionInUse {
+		return nil // connection was in use, just quit cleanly
+	}
+
+	if c.IsSender {
+		// TODO: Add confirmation
+	} else { // Is a Receiver
 		if !gotOK {
 			return errors.New("Transfer interrupted")
 		}
@@ -282,9 +291,6 @@ func (c *Connection) runClient() error {
 		} else {
 			fmt.Printf("\nReceived file written to %s", c.File.Name)
 		}
-	} else {
-		fmt.Println("File sent.")
-		// TODO: Add confirmation
 	}
 	return nil
 }
@@ -348,7 +354,6 @@ func (c *Connection) receiveFile(id int, connection net.Conn) error {
 
 	logger.Debug("waiting for file")
 	var receivedBytes int64
-	receivedFirstBytes := false
 	for {
 		if !c.Debug {
 			c.bars[id].Incr()
@@ -365,10 +370,6 @@ func (c *Connection) receiveFile(id int, connection net.Conn) error {
 		}
 		io.CopyN(newFile, connection, BUFFERSIZE)
 		receivedBytes += BUFFERSIZE
-		if !receivedFirstBytes {
-			receivedFirstBytes = true
-			logger.Debug("Receieved first bytes!")
-		}
 	}
 	logger.Debug("received file")
 	return nil
@@ -402,7 +403,7 @@ func (c *Connection) sendFile(id int, connection net.Conn) {
 	sendBuffer := make([]byte, BUFFERSIZE)
 
 	// open encrypted file
-	file, err := os.OpenFile(c.File.Name+".enc", os.O_RDONLY, 0755)
+	file, err := os.Open(c.File.Name + ".enc")
 	if err != nil {
 		log.Error(err)
 		return

crypto_test.go

@@ -45,4 +45,5 @@ func TestEncryptFiles(t *testing.T) {
 	}
 	os.Remove("temp.dec")
 	os.Remove("temp.enc")
+	os.Remove("temp")
 }

relay.go

@@ -14,12 +14,27 @@ import (
 const MAX_NUMBER_THREADS = 8
 
 type connectionMap struct {
-	reciever map[string]net.Conn
+	receiver           map[string]net.Conn
 	sender             map[string]net.Conn
 	metadata           map[string]string
+	potentialReceivers map[string]struct{}
 	sync.RWMutex
 }
 
+func (c *connectionMap) IsSenderConnected(key string) (found bool) {
+	c.RLock()
+	defer c.RUnlock()
+	_, found = c.sender[key]
+	return
+}
+
+func (c *connectionMap) IsPotentialReceiverConnected(key string) (found bool) {
+	c.RLock()
+	defer c.RUnlock()
+	_, found = c.potentialReceivers[key]
+	return
+}
+
 type Relay struct {
 	connections         connectionMap
 	Debug               bool
@@ -42,9 +57,10 @@ func NewRelay(flags *Flags) *Relay {
 func (r *Relay) Run() {
 	r.connections = connectionMap{}
 	r.connections.Lock()
-	r.connections.reciever = make(map[string]net.Conn)
+	r.connections.receiver = make(map[string]net.Conn)
 	r.connections.sender = make(map[string]net.Conn)
 	r.connections.metadata = make(map[string]string)
+	r.connections.potentialReceivers = make(map[string]struct{})
 	r.connections.Unlock()
 	r.runServer()
 }
@@ -107,7 +123,12 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
 		"codePhrase": codePhrase,
 	})
 
-	if connectionType == "s" {
+	if connectionType == "s" { // sender connection
+		if r.connections.IsSenderConnected(key) {
+			sendMessage("no", connection)
+			return
+		}
+
 		logger.Debug("got sender")
 		r.connections.Lock()
 		r.connections.metadata[key] = metaData
@@ -117,9 +138,9 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
 		receiversAddress := ""
 		for {
 			r.connections.RLock()
-			if _, ok := r.connections.reciever[key]; ok {
-				receiversAddress = r.connections.reciever[key].RemoteAddr().String()
-				logger.Debug("got reciever")
+			if _, ok := r.connections.receiver[key]; ok {
+				receiversAddress = r.connections.receiver[key].RemoteAddr().String()
+				logger.Debug("got receiver")
 				r.connections.RUnlock()
 				break
 			}
@@ -131,18 +152,29 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
 		logger.Debug("preparing pipe")
 		r.connections.Lock()
 		con1 := r.connections.sender[key]
-		con2 := r.connections.reciever[key]
+		con2 := r.connections.receiver[key]
 		r.connections.Unlock()
 		logger.Debug("piping connections")
 		Pipe(con1, con2)
 		logger.Debug("done piping")
 		r.connections.Lock()
 		delete(r.connections.sender, key)
-		delete(r.connections.reciever, key)
+		delete(r.connections.receiver, key)
 		delete(r.connections.metadata, key)
+		delete(r.connections.potentialReceivers, key)
 		r.connections.Unlock()
 		logger.Debug("deleted sender and receiver")
-	} else {
+	} else { //receiver connection "r"
+		if r.connections.IsPotentialReceiverConnected(key) {
+			sendMessage("no", connection)
+			return
+		}
+
+		// add as a potential receiver
+		r.connections.Lock()
+		r.connections.potentialReceivers[key] = struct{}{}
+		r.connections.Unlock()
+
 		// wait for sender's metadata
 		sendersAddress := ""
 		for {
@@ -168,7 +200,7 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
 		if consent == "ok" {
 			logger.Debug("got consent")
 			r.connections.Lock()
-			r.connections.reciever[key] = connection
+			r.connections.receiver[key] = connection
 			r.connections.Unlock()
 		}
 	}