From 73d84da01ca55b6b7bf945b4d75e05afc2072e66 Mon Sep 17 00:00:00 2001
From: Singee <git@singee.me>
Date: Thu, 10 Dec 2020 16:20:08 +0800
Subject: [PATCH] support write directory info

---
 .gitignore       |  3 +-
 protocol.go      | 68 +++++++++++++++++++++++++++++++++++++++++++-
 protocol_test.go | 73 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 142 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 723ef36..61acb19 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
-.idea
\ No newline at end of file
+.idea
+*.log
\ No newline at end of file
diff --git a/protocol.go b/protocol.go
index e9a4019..5f7ee91 100644
--- a/protocol.go
+++ b/protocol.go
@@ -7,6 +7,7 @@ import (
 	"github.com/alessio/shellescape"
 	"golang.org/x/crypto/ssh"
 	"io"
+	"os"
 	"strconv"
 	"strings"
 )
@@ -16,6 +17,8 @@ type RemoteClient struct {
 
 	stdout io.Reader
 	stdin  io.WriteCloser
+
+	verbose bool
 }
 
 func NewClient(session *ssh.Session) (*RemoteClient, error) {
@@ -40,7 +43,37 @@ func (c *RemoteClient) Start(filename string, isDirectory bool) error {
 	commandBuilder.WriteString("scp -t")
 
 	if isDirectory {
-		commandBuilder.WriteString(" -d")
+		commandBuilder.WriteString(" -r")
+	}
+	if c.verbose {
+		commandBuilder.WriteString(" -v")
+
+		// just for debug
+		stderr, err := c.session.StderrPipe()
+		if err != nil {
+			fmt.Println("Cannot connect to stderr")
+		} else {
+			f, err := os.Create("protocol.log")
+
+			if err != nil {
+				fmt.Println("Cannot create test file")
+			} else {
+				go func() {
+					buffer := make([]byte, 1024)
+					for {
+						read, _ := stderr.Read(buffer)
+						if read != 0 {
+							_, err := f.Write(buffer[:read])
+
+							if err != nil {
+								fmt.Println("Cannot write verbose info")
+							}
+						}
+					}
+				}()
+			}
+
+		}
 	}
 
 	commandBuilder.WriteByte(' ')
@@ -124,3 +157,36 @@ func (c *RemoteClient) WriteFile(perm string, size int64, filename string, data
 
 	return nil
 }
+
+func (c *RemoteClient) WriteDirectoryStart(perm string, filename string) error {
+	if len(filename) > 1024 {
+		// Unknown reason, just as same as https://github.com/openssh/openssh-portable/blob/master/scp.c#L1204
+		filename = filename[:1024]
+	}
+
+	_, err := fmt.Fprintln(c.stdin, "D"+perm, 0, filename)
+	if err != nil {
+		return err
+	}
+
+	err = c.checkResponse()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (c *RemoteClient) WriteDirectoryEnd() error {
+	_, err := fmt.Fprintln(c.stdin, "E")
+	if err != nil {
+		return err
+	}
+
+	err = c.checkResponse()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/protocol_test.go b/protocol_test.go
index ca14f34..7fed226 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -148,3 +148,76 @@ func TestWriteMultipleFilesToFolder(t *testing.T) {
 	tt.AssertEqual(t, "hahaha", string(readFile(t, "/test/test01")))
 	tt.AssertEqual(t, "xixixixi", string(readFile(t, "/test/test02")))
 }
+
+func TestWriteDirectoryToFolder(t *testing.T) {
+	reset(t)
+
+	session := getSshSession(t)
+	defer session.Close()
+
+	client, err := NewClient(session)
+	tt.AssertIsNil(t, err)
+
+	/*
+		Virtual File Tree
+			-- (Folder) a
+			  -- (File) b (Content: hahaha)
+			  -- (File) c (Content: xixixixi)
+			-- (Folder) b
+			-- (Folder) c
+			  -- (File) d (Empty)
+			-- (File) e (Content: root)
+	*/
+
+	err = client.Start("/test", true)
+	tt.AssertIsNil(t, err)
+
+	// first folder a
+	err = client.WriteDirectoryStart("0755", "a")
+	tt.AssertIsNil(t, err)
+
+	// in folder a, file b  (Content: hahaha)
+	err = client.WriteFile("0644", 6, "b", strings.NewReader("hahaha"))
+	tt.AssertIsNil(t, err)
+	// in folder a, file c (Content: xixixixi)
+	err = client.WriteFile("0644", 8, "c", strings.NewReader("xixixixi"))
+	tt.AssertIsNil(t, err)
+
+	// end first folder a
+	err = client.WriteDirectoryEnd()
+	tt.AssertIsNil(t, err)
+
+	// second folder b
+	err = client.WriteDirectoryStart("0755", "b")
+	tt.AssertIsNil(t, err)
+	err = client.WriteDirectoryEnd()
+	tt.AssertIsNil(t, err)
+
+	// third folder c
+	err = client.WriteDirectoryStart("0755", "c")
+	tt.AssertIsNil(t, err)
+	// in folder c, empty file d
+	err = client.WriteFile("0644", 0, "d", strings.NewReader(""))
+	tt.AssertIsNil(t, err)
+	err = client.WriteDirectoryEnd()
+	tt.AssertIsNil(t, err)
+
+	// last file e (Content: root)
+	err = client.WriteFile("0644", 4, "e", strings.NewReader("root"))
+	tt.AssertIsNil(t, err)
+
+	// Write End, Now Check
+
+	tt.AssertTrue(t, checkFileExist(t, "/test/a"))
+	tt.AssertTrue(t, checkFileExist(t, "/test/a/b"))
+	tt.AssertTrue(t, checkFileExist(t, "/test/a/c"))
+	tt.AssertTrue(t, checkFileExist(t, "/test/b"))
+	tt.AssertTrue(t, checkFileExist(t, "/test/c"))
+	tt.AssertTrue(t, checkFileExist(t, "/test/c/d"))
+	tt.AssertTrue(t, checkFileExist(t, "/test/e"))
+
+	tt.AssertEqual(t, "hahaha", string(readFile(t, "/test/a/b")))
+	tt.AssertEqual(t, "xixixixi", string(readFile(t, "/test/a/c")))
+	tt.AssertEqual(t, "", string(readFile(t, "/test/c/d")))
+	tt.AssertEqual(t, "root", string(readFile(t, "/test/e")))
+}