Implementing the Upload Server correctly and refactoring config package a bit

README.md

@@ -17,6 +17,11 @@ If the upload server is running, the user will be able to put files
 on the **absoluteServePath**. After the file is uploaded successfully,
 if the timeout is not reached, the user will get back the filename.
 
+To send data to the upload server, the following command can be used:
+
+```nc ip port < gopher.png```
+
+
 #### Sending commands via netcat
 
 To grab a file the following command can be send:
@@ -48,6 +53,7 @@ Sample Configuration File:
     },
     "upload": {
         "enabled": false,
+        "directory": "upload",
         "timeout": 5,
         "address": "localhost",
         "port": 8081

server/server/commands.go

@@ -10,9 +10,6 @@ import (
 	"strconv"
 	"strings"
 
-	"bufio"
-	"time"
-
 	"github.com/spf13/viper"
 	"github.com/zyxar/image2ascii/ascii"
 )
@@ -27,34 +24,14 @@ func randSeq(n int) string {
 }
 
 // UploadFile uploads a file to the server
-func UploadFile(c Client) (string, error) {
-
-	input := bufio.NewScanner(c.Connection())
-	start := time.Now()
-	timeout := time.Duration(viper.GetInt("upload.timeout")) * time.Second
-
-	var filename = randSeq(10)
-	var _, err = os.Stat(MakePathFromStringStack(c.Stack()) + filename)
-
-	// Make sure that the filename is random.
-	for !os.IsNotExist(err) {
-		filename = randSeq(10)
-		_, err = os.Stat(MakePathFromStringStack(c.Stack()) + filename)
-	}
-
+func UploadFile(c Client, filename string) (string, error) {
 	f, err := os.Create(MakePathFromStringStack(c.Stack()) + filename)
 	if err != nil {
 		return filename, err
 	}
 	defer f.Close()
 
-	for input.Scan() {
-		if time.Since(start) >= timeout {
-			log.Println(c.Connection().RemoteAddr().String() + " has reached timeout!")
-			break
-		}
-		f.Write(input.Bytes())
-	}
+	io.Copy(f, c.Connection())
 
 	return filename, nil
 }

server/server/config/config.go

@@ -4,18 +4,13 @@ package config
 import (
 	"log"
 
-	"flag"
-
 	"github.com/fsnotify/fsnotify"
 	"github.com/spf13/viper"
 )
 
-// ConfigPath will be used via cmd to set the configuration path for the config file.
-var ConfigPath string
-
 // loadConfigFromFile tries to load the configuration file from the disk.
-func loadConfigFromFile() error {
-	viper.SetConfigName("config")
+func loadConfigFromFile(configName string) error {
+	viper.SetConfigName(configName)
 	viper.AddConfigPath(viper.GetString("ConfigPath"))
 
 	err := viper.ReadInConfig() // Find and read the config file
@@ -26,31 +21,30 @@ func loadConfigFromFile() error {
 }
 
 // setDefaultConfiguration will set the default configuration settings.
-func setDefaultConfiguration() {
+func setDefaultConfiguration(configPath string) {
 	viper.SetDefault("address", "localhost")
 	viper.SetDefault("port", 8080)
-	viper.SetDefault("configPath", ConfigPath)
+	viper.SetDefault("configPath", configPath)
 	viper.SetDefault("maxDirDepth", 30)
 	viper.SetDefault("absoluteServePath", "./")
 	viper.SetDefault("pic.x", 0)
 	viper.SetDefault("pic.y", 0)
 	viper.SetDefault("pic.color", false)
 	viper.SetDefault("upload.enabled", false)
+	viper.SetDefault("upload.directory", "upload")
+	viper.SetDefault("upload.timeout", 3)
 	viper.SetDefault("upload.address", "localhost")
 	viper.SetDefault("upload.port", 8081)
 }
 
 // InitializeConfiguration initializes the configuration for the application.
-func InitializeConfiguration() {
-	flag.StringVar(&ConfigPath, "config", ".", "Set the location of the config file.")
-	flag.Parse()
-
-	setDefaultConfiguration()
-	loadConfigFromFile()
+func InitializeConfiguration(configName string, configPath string) {
+	setDefaultConfiguration(configPath)
+	loadConfigFromFile(configName)
 
 	viper.WatchConfig()
 }
 
-func ConfigChangeCallback(cb func(event fsnotify.Event)) {
+func ChangeCallback(cb func(event fsnotify.Event)) {
 	viper.OnConfigChange(cb)
 }

server/server/config/config_test.go

@@ -8,7 +8,7 @@ import (
 
 func TestLoadConfigFromFile(t *testing.T) {
 	// SetDefaultConfiguration must be called BEFORE LoadConfigFromFile.
-	InitializeConfiguration()
+	InitializeConfiguration("config", "./")
 
 	Address := viper.GetString("address")
 	if Address == "" {

server/server/connection.go

@@ -7,6 +7,10 @@ import (
 	"log"
 	"net"
 
+	"os"
+
+	"time"
+
 	"github.com/fsnotify/fsnotify"
 	"github.com/metonimie/simpleFTP/server/server/config"
 	"github.com/spf13/viper"
@@ -17,6 +21,15 @@ import (
 // is also send to the client.
 const DataBufferSize = 1024 * 1024
 
+// ConfigPath is used by the config package to find the config file.
+var ConfigPath string
+
+// uploadDirectory is the directory where the files will be uploaded
+var uploadDirectory string
+
+// uploadTimeout is the amount in seconds the server will wait for a file to be uploaded
+var uploadTimeout time.Duration
+
 // Client interface provides the blueprints for the Client that is used by the server.
 type Client interface {
 	Connection() net.Conn        // Connection returns the connection stream.
@@ -25,6 +38,11 @@ type Client interface {
 	Stack() *StringStack         // Returns the underlying String Stack.
 }
 
+type UploadResult struct {
+	Filename string // Filename represents the file which was randomly created by the server
+	Err      error  // Err is the error that occurred while uploading the file.
+}
+
 // FTPClient represents a FTPClient connection, it holds a root cage and the underlying connection.
 type FTPClient struct {
 	rootCage   *StringStack // rootCage is a StringStack that is used to represent the current directory the client is in.
@@ -89,8 +107,8 @@ func HandleConnection(client Client) {
 }
 
 func Init() {
-	config.InitializeConfiguration()
-	config.ConfigChangeCallback(func(event fsnotify.Event) {
+	config.InitializeConfiguration("config", ConfigPath)
+	config.ChangeCallback(func(event fsnotify.Event) {
 		log.Println("Configuration reloaded successfully!")
 	})
 }
@@ -126,44 +144,101 @@ func StartFtpServer() {
 	}
 }
 
-func StartUploadServer() {
+func HandleUpload(conn net.Conn) {
+	// Initialize Client
+	client := FTPClient{}
+	client.SetStack(MakeStringStack(2))
+
+	// Upload directory
+	client.Stack().Push(uploadDirectory)
+	client.SetConnection(conn)
+
+	// Create the file on the disk and make sure that the filename is random.
+	var filename = randSeq(10)
+	var _, err = os.Stat(MakePathFromStringStack(client.Stack()) + filename)
+
+	for !os.IsNotExist(err) {
+		filename = randSeq(10)
+		_, err = os.Stat(MakePathFromStringStack(client.Stack()) + filename)
+	}
+
+	// This channel will be used to store the uploadResult
+	c1 := make(chan UploadResult, 1)
+	log.Println(conn.RemoteAddr().String() + " is uploading something.")
+
+	// Create a new Go routine for uploading
+	go func() {
+		fname, err := UploadFile(&client, filename)
+		c1 <- UploadResult{fname, err}
+	}()
+
+	// Wait for either UploadResult or Timeout
+	select {
+	case result := <-c1:
+		{
+			filename, err := result.Filename, result.Err
+
+			if err == nil {
+				io.WriteString(conn, filename)
+				log.Println(conn.RemoteAddr().String() + "'s upload finished.")
+			} else {
+				log.Println(fmt.Sprintf("%s: %s %s", "HandleUpload", conn.RemoteAddr().String(), err.Error()))
+
+				client.Stack().Push(filename)
+				os.Remove(MakePathFromStringStack(client.Stack()))
+
+				io.WriteString(conn, err.Error())
+			}
+
+			conn.Close()
+		}
+	case <-time.After(time.Second * uploadTimeout):
+		{
+			io.WriteString(conn, "Timeout")
+			conn.Close()
+		}
+	}
+}
+
+// StartUploadServer starts the uploading server
+func StartUploadServer() error {
 	if viper.GetBool("upload.enabled") == false {
 		log.Println("Uploading not enabled. To enable modify the config file and restart the server")
-		return
+		return ErrUploadServerFailure
 	}
 
-	Addr := viper.GetString("upload.address")
-	Port := viper.GetInt("upload.port")
+	addr := viper.GetString("upload.address")
+	port := viper.GetInt("upload.port")
+	uploadDirectory = viper.GetString("upload.directory")
+	uploadTimeout = time.Duration(viper.GetInt("upload.timeout"))
 
-	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", Addr, Port))
+	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", addr, port))
 	if err != nil {
-		log.Fatal(err)
+		log.Println(err)
+		return ErrUploadServerFailure
 	}
 
-	log.Println("Upload server running on:", Addr, "port", Port)
+	err = os.Mkdir(uploadDirectory, 0740)
+	if err != nil {
+		if _, err := os.Stat(uploadDirectory); err != nil {
+			if os.IsNotExist(err) {
+				log.Println("Can't create upload directory!")
+				return ErrUploadServerFailure
+			}
+		}
+	}
+
+	log.Println("Upload server running on:", addr, "port", port)
 
 	for {
-
 		conn, err := listener.Accept()
 		if err != nil {
 			log.Print(err)
 			continue
 		}
 
-		client := FTPClient{}
-		client.SetStack(MakeStringStack(1))
-		client.SetConnection(conn)
-
-		log.Println(conn.RemoteAddr().String() + " is uploading something.")
-
-		filename, err := UploadFile(&client)
-		if err == nil {
-			io.WriteString(conn, filename)
-		} else {
-			log.Print(conn.RemoteAddr().String())
-			log.Println(err)
+		go HandleUpload(conn)
 	}
 
-		log.Println(conn.RemoteAddr().String() + "'s upload finished.")
-	}
+	return nil
 }

server/server/errors.go

@@ -52,3 +52,8 @@ var (
 	ErrAlreadyAtBaseDirectory = PathError{errors.New("can't go past beyond root directory")}
 	ErrSlashNotAllowed        = PathError{errors.New("slash is not allowed in file names")}
 )
+
+// General Errors
+var (
+	ErrUploadServerFailure = errors.New("upload server failed to start")
+)

server/simplFTP.go

@@ -1,10 +1,15 @@
 package main
 
 import (
+	"flag"
+
 	"github.com/metonimie/simpleFTP/server/server"
 )
 
 func main() {
+	flag.StringVar(&server.ConfigPath, "config", ".", "Set the location of the config file.")
+	flag.Parse()
+
 	server.Init()
 
 	go server.StartUploadServer()