Merge pull request #1 from zaz600/fix-multiline

rewrite multiline code (merge from zaz600)

ftp.go

@@ -15,6 +15,7 @@ import (
 )
 
 var REGEX_PWD_PATH *regexp.Regexp = regexp.MustCompile(`\"(.*)\"`)
+var errNoCode = errors.New("No code")
 
 type FTP struct {
 	conn net.Conn
@@ -355,35 +356,55 @@ func (ftp *FTP) receiveLine() (string, error) {
 	if ftp.debug {
 		log.Printf("< %s", line)
 	}
-
 	return line, err
 }
 
+func (ftp *FTP) getCode(line string) (int, bool, error) {
+	trimmedLine := strings.Trim(line, " \r\n")
+	fields := strings.Fields(trimmedLine)
+	if len(fields) == 0 {
+		return 0, false, errNoCode
+	}
+
+	if len(fields[0]) == 3 {
+		if code, err := strconv.Atoi(fields[0]); err == nil {
+			return code, false, nil
+		}
+	} else if len(fields[0]) == 4 && fields[0][3] == '-' {
+		if code, err := strconv.Atoi(fields[0][:3]); err == nil {
+			return code, true, nil
+		}
+	} else {
+		if len(trimmedLine) >= 4 && trimmedLine[3] == '-' {
+			if code, err := strconv.Atoi(trimmedLine[:3]); err == nil {
+				return code, true, nil
+			}
+		}
+	}
+	return 0, false, errNoCode
+}
+
 func (ftp *FTP) receive() (string, error) {
 	line, err := ftp.receiveLine()
 
 	if err != nil {
 		return line, err
 	}
-
-	if (len(line) >= 4) && (line[3] == '-') {
+	code, beginMultiline, err := ftp.getCode(line)
+	if err == nil && beginMultiline {
 		//Multiline response
-		closingCode := line[:3] + " "
+		closingCode := code
 		for {
 			str, err := ftp.receiveLine()
 			line = line + str
 			if err != nil {
 				return line, err
 			}
-			if len(str) < 4 {
-				if ftp.debug {
-					log.Println("Uncorrectly terminated response")
-				}
+			code, _, err := ftp.getCode(str)
+			if err == nil && code == closingCode {
 				break
 			} else {
-				if str[:4] == closingCode {
-					break
-				}
+				//check error codes 5xx 4xx
 			}
 		}
 	}

ftp_test.go

@@ -168,3 +168,34 @@ func TestGetFilesListOnGoodServer(t *testing.T) {
 
 	connection.Close()
 }
+
+type getCodeResult struct {
+	code           int
+	beginMultiline bool
+	err            error
+}
+
+func TestGetCode(t *testing.T) {
+	var tests = []struct {
+		input string
+		want  getCodeResult
+	}{
+		{"220 test", getCodeResult{220, false, nil}},
+		{"220     test", getCodeResult{220, false, nil}},
+		{"  220     test", getCodeResult{220, false, nil}},
+		{"220- test", getCodeResult{220, true, nil}},
+		{"220asdf test", getCodeResult{0, false, errNoCode}},
+		{"", getCodeResult{0, false, errNoCode}},
+		{"\r\n", getCodeResult{0, false, errNoCode}},
+		{"123-Firstline", getCodeResult{123, true, nil}},
+		{"123-First line", getCodeResult{123, true, nil}},
+	}
+	ftp := &FTP{}
+	for _, test := range tests {
+		code, beginMultiline, err := ftp.getCode(test.input)
+		res := getCodeResult{code, beginMultiline, err}
+		if res != test.want {
+			t.Errorf("want: %#v, expected: %#v", test.want, res)
+		}
+	}
+}