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,71 +181,86 @@ func (c *Connection) runClient() error {
 			if c.IsSender { // this is a sender
 				logger.Debug("waiting for ok from relay")
 				message = receiveMessage(connection)
-				logger.Debug("got ok from relay")
-				if id == 0 {
-					fmt.Printf("\nSending (->%s)..\n", message)
+				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)
+					}
+					// wait for pipe to be made
+					time.Sleep(100 * time.Millisecond)
+					// Write data from file
+					logger.Debug("send file")
+					c.sendFile(id, connection)
+					fmt.Println("File sent.")
 				}
-				// wait for pipe to be made
-				time.Sleep(100 * time.Millisecond)
-				// Write data from file
-				logger.Debug("send file")
-				c.sendFile(id, connection)
 			} else { // this is a receiver
 				logger.Debug("waiting for meta data from sender")
 				message = receiveMessage(connection)
-				m := strings.Split(message, "-")
-				encryptedData, salt, iv, sendersAddress := m[0], m[1], m[2], m[3]
-				encryptedBytes, err := hex.DecodeString(encryptedData)
-				if err != nil {
-					log.Error(err)
-					return
-				}
-				decryptedBytes, _ := Decrypt(encryptedBytes, c.Code, salt, iv, c.DontEncrypt)
-				err = json.Unmarshal(decryptedBytes, &c.File)
-				if err != nil {
-					log.Error(err)
-					return
-				}
-				log.Debugf("meta data received: %v", c.File)
-				// have the main thread ask for the okay
-				if id == 0 {
-					fmt.Printf("Receiving file (%d bytes) into: %s\n", c.File.Size, c.File.Name)
-					var sentFileNames []string
-
-					if fileAlreadyExists(sentFileNames, c.File.Name) {
-						fmt.Printf("Will not overwrite file!")
-						os.Exit(1)
-					}
-					getOK := getInput("ok? (y/n): ")
-					if getOK == "y" {
-						gotOK = true
-						sentFileNames = append(sentFileNames, c.File.Name)
-					}
-					gotResponse = true
-				}
-				// wait for the main thread to get the okay
-				for limit := 0; limit < 1000; limit++ {
-					if gotResponse {
-						break
-					}
-					time.Sleep(10 * time.Millisecond)
-				}
-				if !gotOK {
-					sendMessage("not ok", connection)
+				if message == "no" {
+					fmt.Println("The specifed code is already in use by a receiver.")
+					gotConnectionInUse = true
 				} else {
-					sendMessage("ok", connection)
-					logger.Debug("receive file")
-					if id == 0 {
-						fmt.Printf("\n\nReceiving (<-%s)..\n", sendersAddress)
+					m := strings.Split(message, "-")
+					encryptedData, salt, iv, sendersAddress := m[0], m[1], m[2], m[3]
+					encryptedBytes, err := hex.DecodeString(encryptedData)
+					if err != nil {
+						log.Error(err)
+						return
+					}
+					decryptedBytes, _ := Decrypt(encryptedBytes, c.Code, salt, iv, c.DontEncrypt)
+					err = json.Unmarshal(decryptedBytes, &c.File)
+					if err != nil {
+						log.Error(err)
+						return
+					}
+					log.Debugf("meta data received: %v", c.File)
+					// have the main thread ask for the okay
+					if id == 0 {
+						fmt.Printf("Receiving file (%d bytes) into: %s\n", c.File.Size, c.File.Name)
+						var sentFileNames []string
+
+						if fileAlreadyExists(sentFileNames, c.File.Name) {
+							fmt.Printf("Will not overwrite file!")
+							os.Exit(1)
+						}
+						getOK := getInput("ok? (y/n): ")
+						if getOK == "y" {
+							gotOK = true
+							sentFileNames = append(sentFileNames, c.File.Name)
+						}
+						gotResponse = true
+					}
+					// wait for the main thread to get the okay
+					for limit := 0; limit < 1000; limit++ {
+						if gotResponse {
+							break
+						}
+						time.Sleep(10 * time.Millisecond)
+					}
+					if !gotOK {
+						sendMessage("not ok", connection)
+					} else {
+						sendMessage("ok", connection)
+						logger.Debug("receive file")
+						fmt.Printf("\n\nReceiving (<-%s)..\n", sendersAddress)
+						c.receiveFile(id, connection)
 					}
-					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
-	sender   map[string]net.Conn
-	metadata map[string]string
+	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()
 		}
 	}